llvm / llvm/llvm-project

[InstCombine] Fold masked add/lshr/add runtime-unroll remainder using modular arithmetic

Open
#217,230 2 comments 0 reactions 1 assignee Claimed by @ParkHanbum View on GitHub
llvm:instcombine missed-optimization
Dominant language
LLVM
Stars
40.5k
Forks
18.7k
PR merge metrics
PR metrics pending

Description

## Summary

InstCombine misses the following simplification:

```llvm
%a = add i64 %x, 12
%b = lshr i64 %a, 2
%c = add nuw nsw i64 %b, 1
%r = and i64 %c, 3
```

This is equivalent to:

```llvm
%b = lshr i64 %x, 2
%r = and i64 %b, 3
```

The pattern was found more than once in optimized ONNX Runtime IR for
`onnxruntime/core/mlas/lib/activate.cpp`, in
`_Z14MlasActivationPK15MLAS_ACTIVATIONPfPKfmmm`.

It is not merely source-level arithmetic. It is produced by LLVM while forming
the remainder count for a runtime-unrolled loop and then partially
canonicalizing that count.

## Minimal reproducer

```llvm
; RUN: opt -S -passes=instcombine %s | FileCheck %s

define i64 @masked_runtime_unroll_remainder(i64 %x) {
; CHECK-LABEL: @masked_runtime_unroll_remainder(
; CHECK-NEXT: [[Q:%.*]] = lshr i64 [[X:%.*]], 2
; CHECK-NEXT: [[R:%.*]] = and i64 [[Q]], 3
; CHECK-NEXT: ret i64 [[R]]
;
%a = add i64 %x, 12
%b = lshr i64 %a, 2
%c = add nuw nsw i64 %b, 1
%r = and i64 %c, 3
ret i64 %r
}
```

The real-world form also contains a guard, but the guard is not needed for the
arithmetic simplification:

```llvm
define range(i64 0, 4) i64 @src(i64 noundef %N) {
entry:
%guard = icmp ugt i64 %N, 3
%n4 = add i64 %N, 12
%q = lshr i64 %n4, 2
%trip = add nuw nsw i64 %q, 1
%unroll_count = and i64 %trip, 3
%ret = select i1 %guard, i64 %unroll_count, i64 0
ret i64 %ret
}
```

Expected:

```llvm
define range(i64 0, 4) i64 @src(i64 noundef %N) {
entry:
%guard = icmp ugt i64 %N, 3
%q = lshr i64 %N, 2
%unroll_count = and i64 %q, 3
%ret = select i1 %guard, i64 %unroll_count, i64 0
ret i64 %ret
}
```

## Why the transform is correct

Write the unsigned `i64` input as:

```text
N = 4Q + R, where 0 <= R < 4
```

Then `Q = N >> 2`. Integer addition is modulo `2^64`, so:

```text
((N + 12) mod 2^64) >> 2
= (Q + 3) mod 2^62
```

Therefore:

```text
((((N + 12) mod 2^64) >> 2) + 1) & 3
= (((Q + 3) mod 2^62) + 1) mod 4
= (Q + 4) mod 4
= Q mod 4
= (N >> 2) & 3
```

The possible wraparound of `N + 12` does not invalidate the transform because
`2^62` is a multiple of 4, and the final `and 3` observes only the value modulo
4.

The flags on the second addition are also valid: after `lshr i64 ..., 2`, its
operand is in `[0, 2^62 - 1]`, so adding one cannot overflow either signed or
unsigned `i64`.

## How this IR is formed

The relevant ONNX Runtime template loop processes four floats per iteration:

```cpp
size_t n = N;

if (n >= 4) {
do {
// Process four floats.
buffer += 4;
n -= 4;
} while (n >= 4);
}
```

Source:
https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/mlas/lib/activate.cpp#L220-L290

For `N >= 4`, the loop trip count is:

```text
((N - 4) >> 2) + 1
```

When LLVM runtime-unrolls this loop by four, it needs:

```text
trip_count % 4
```

`CreateTripRemainder()` implements a power-of-two remainder as
`and TripCount, Count - 1`, which produces the final `and ..., 3`:

https://github.com/llvm/llvm-project/blob/main/llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp#L562-L580

A pass dump from a reduced C++ reproducer shows the following transition.
Before InstCombine:

```llvm
%sub = sub nuw i64 %N, 4
%q = lshr i64 %sub, 2
%trip = add nuw nsw i64 %q, 1
%rem = and i64 %trip, 3
```

After InstCombine:

```llvm
%add = add i64 %N, 12
%q = lshr i64 %add, 2
%trip = add nuw nsw i64 %q, 1
%rem = and i64 %trip, 3
```

This intermediate canonicalization is valid because `-4` and `+12` differ by
16, which is unobservable after shifting by two and masking with three.
SimplifyCFG can then fold the guarded branches into the `select` shown above,
but the final modular simplification is missed.

The same source loop is instantiated for multiple activation kinds and for the
bias/no-bias variants, explaining why the pattern can occur multiple times in
the same optimized `MlasActivation` function:

https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/mlas/lib/activate.cpp#L378-L482

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.