[InstCombine] Fold masked add/lshr/add runtime-unroll remainder using modular arithmetic
- 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
Assessment
This issue has not been assessed yet.