llvm / llvm/llvm-project

Recognize software IEEE FP16 conversion idioms as fptrunc/fpext

Open
#214,802 1 comment 0 reactions 1 assignee Claimed by @ParkHanbum View on GitHub
backend:X86 llvm:SelectionDAG missed-optimization
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.