Recognize software IEEE FP16 conversion idioms as fptrunc/fpext
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
## Description
While inspecting ONNX Runtime compiled for `x86-64-v4`, I found that a software implementation of FP32-to-FP16 conversion is vectorized as a long sequence of integer/FP operations instead of being recognized as an FP conversion and lowered to `vcvtps2ph`.
The pattern originates from ONNX Runtime's `MLFloat16` conversion implementation.
A real-world occurrence is in:
```cpp
template
Status MoE::ProcessExpertBatch(...) const {
...
if (fc1_bias) {
for (int64_t batch = 0; batch < batch_size; ++batch) {
T* batch_output = fc1_output + batch * fc1_output_size;
// Explicit loop for better vectorization
for (int64_t i = 0; i < fc1_output_size; ++i) {
batch_output[i] =
static_cast(
static_cast(batch_output[i]) +
static_cast(fc1_bias[i]));
}
}
}
...
}
```
For the `MoE` instantiation, this effectively becomes:
```cpp
float a = static_cast(batch_output[i]);
float b = static_cast(fc1_bias[i]);
float sum = a + b;
batch_output[i] = MLFloat16(sum);
```
`MLFloat16::ToFloatImpl()` and `ToUint16Impl()` implement the conversions using portable bit manipulation rather than native FP16 operations.
LLVM successfully vectorizes the loop, but it also vectorizes the software conversion implementation itself.
For example, the FP32-to-FP16 part becomes roughly:
```llvm
%abs = call <32 x float> @llvm.fabs.v32f32(<32 x float> %x)
%bits = bitcast <32 x float> %abs to <32 x i32>
%small = icmp ult <32 x i32> %bits, splat (i32 1199570944)
%offset = add <32 x i32> %bits, splat (i32 -947912704)
%normal = icmp ult <32 x i32> %offset, splat (i32 251658240)
%t0 = lshr <32 x i32> %bits, splat (i32 13)
%odd = and <32 x i32> %t0, splat (i32 1)
%t1 = add <32 x i32> %bits, splat (i32 134221823)
%t2 = add <32 x i32> %t1, %odd
%rounded = lshr <32 x i32> %t2, splat (i32 13)
...
%isnan = icmp ugt <32 x i32> %bits, splat (i32 2139095040)
%special = select <32 x i1> %isnan,
<32 x i32> splat (i32 32256),
<32 x i32> splat (i32 31744)
...
%result = trunc <32 x i32> %result32 to <32 x i16>
```
The important constants correspond to the software FP16 conversion algorithm:
```text
0x47800000 -> 1199570944
0x38800000 -> 947912704
0x08000fff -> 134221823
0x7f800000 -> 2139095040
0x7e00 -> 32256
0x7c00 -> 31744
```
On an `x86-64-v4` target, native FP16 conversion is already available, so ideally the non-NaN part of this idiom could be represented as:
```llvm
%h = fptrunc <16 x float> %x to <16 x half>
%r = bitcast <16 x half> %h to <16 x i16>
```
which X86 can lower to `vcvtps2ph`.
## Simple case when NaNs are known not to occur
If the input is known never to be NaN, e.g. from `nnan` or equivalent analysis, the whole software conversion should be replaceable by `fptrunc`.
### Before
```llvm
; Simplified form of the software FP32 -> FP16 conversion.
; The real pattern contains range checks, rounding and special-case handling.
define <16 x i16> @before(<16 x float> %x) {
entry:
%abs = call <16 x float> @llvm.fabs.v16f32(<16 x float> %x)
%bits = bitcast <16 x float> %abs to <16 x i32>
%t0 = lshr <16 x i32> %bits, splat (i32 13)
%odd = and <16 x i32> %t0, splat (i32 1)
%t1 = add <16 x i32> %bits, splat (i32 134221823)
%t2 = add <16 x i32> %t1, %odd
%rounded = lshr <16 x i32> %t2, splat (i32 13)
; ... range/subnormal/Inf handling ...
%r = trunc <16 x i32> %rounded to <16 x i16>
ret <16 x i16> %r
}
```
### Expected canonical form
```llvm
define <16 x i16> @after(<16 x float> %x) {
entry:
%h = fptrunc nnan <16 x float> %x to <16 x half>
%r = bitcast <16 x half> %h to <16 x i16>
ret <16 x i16> %r
}
```
For X86 this can then lower to `vcvtps2ph`.
## Case without `nnan`
The ONNX Runtime implementation canonicalizes NaNs explicitly:
```text
NaN -> 0x7e00 | sign
```
Therefore replacing the entire software implementation with plain `fptrunc` is not necessarily bit-equivalent for NaNs.
However, most of the conversion could still potentially use `fptrunc`, with only NaNs handled separately:
```llvm
define <16 x i16> @after_with_nan_fixup(<16 x float> %x) {
entry:
%h = fptrunc <16 x float> %x to <16 x half>
%converted = bitcast <16 x half> %h to <16 x i16>
%xbits = bitcast <16 x float> %x to <16 x i32>
%abs = and <16 x i32> %xbits, splat (i32 2147483647)
%isnan = icmp ugt <16 x i32> %abs,
splat (i32 2139095040)
%sign32 = lshr <16 x i32> %xbits, splat (i32 16)
%sign16.tmp = trunc <16 x i32> %sign32 to <16 x i16>
%sign16 = and <16 x i16> %sign16.tmp, splat (i16 -32768)
%canonical.nan =
or <16 x i16> %sign16, splat (i16 32256)
%r = select <16 x i1> %isnan,
<16 x i16> %canonical.nan,
<16 x i16> %converted
ret <16 x i16> %r
}
```
This should retain the application's NaN canonicalization while allowing the common non-NaN path to use native FP16 conversion.
## FP16-to-FP32 has a similar pattern
The same ONNX Runtime loop also contains two software FP16-to-FP32 conversions.
They are currently vectorized into patterns such as:
```llvm
%mag16 = and <32 x i16> %x, splat (i16 32767)
%mag32 = zext <32 x i16> %mag16 to <32 x i32>
%shift = shl <32 x i32> %mag32, splat (i32 13)
%exp = and <32 x i32> %shift, splat (i32 260046848)
...
%normal = add <32 x i32> %shift, splat (i32 939524096)
...
%result = bitcast <32 x i32> %bits to <32 x float>
```
Conceptually this is:
```llvm
%h = bitcast <32 x i16> %x to <32 x half>
%r = fpext <32 x half> %h to <32 x float>
```
and could similarly enable `vcvtph2ps`.
It may make sense to handle FP16-to-FP32 separately, since the legality requirements around NaN representation may differ.
## Why this seems worth optimizing
This is not a synthetic pattern.
It occurs in ONNX Runtime's MoE implementation in an explicitly vectorized bias-add loop:
```cpp
batch_output[i] =
static_cast(
static_cast(batch_output[i]) +
static_cast(fc1_bias[i]));
```
LLVM successfully vectorizes the loop, but because the `MLFloat16` conversions are implemented as portable bit manipulation, the resulting vector loop contains a large number of integer operations, compares and selects instead of native FP16 conversion instructions.
The backend already knows how to efficiently lower `fpext`/`fptrunc` between `half` and `float`. The missing piece appears to be recognizing the software IEEE FP16 conversion idiom and canonicalizing it back to FP conversion operations when the semantics permit.
A conservative first step could be recognizing the FP32-to-FP16 idiom only when the input is known never to be NaN. A later extension could preserve explicit NaN canonicalization while using `fptrunc` for the remaining values.
Contributor guide
Assessment
This issue has not been assessed yet.