llvm / llvm/llvm-project

[mlir][linalg] partial reduction tiling propagates unjustified integer overflow flags

Open
#218,538 1 comment 1 reaction 0 assignees View on GitHub
mlir
Dominant language
LLVM
Stars
40.5k
Forks
18.7k
PR merge metrics
PR metrics pending

Description

PR https://github.com/llvm/llvm-project/pull/214033 added support for `arith.subi` reducers, which pointed out a potential issue. I explored and wrote up this issue with the help of Claude.

Sorry for the length of this issue, but this came out of a very long discussion in the PR, and I wanted to summarize why it looks to me like the overflow flags are a potential issue. I have never seen these flags used in the wild; it is simply an observation that came from writing the PR since it is slightly different from the other reducers.

Tiling a `linalg` reduction with `PartialReductionOpInterface` copies the combiner's `arith` integer overflow flags (`nsw` / `nuw`) onto both the partial accumulation and the merge. Those flags are an assertion about the *original* computation. Tiling replaces it with a different one, so the assertion can be false which makes the tiled reduction produce poison where the untiled one is well defined.

## Why overflow flags no longer apply

`tileToPartialReduction` clones the combiner into the partial reduction body, and `mergeReductions` builds the merge from the same combiner. Both keep the overflow flags. But a partial reduction is not the accumulation that the flag was attached to: it starts from the neutral element instead of from `init`.

`init` is what keeps the untiled accumulator in range. `linalg.generic` does not fix the order in which it visits its inputs, so a flag is only honest on the source op if it holds for *every* order, that is if `init` accumulated with any subset of the inputs stays in range. A partial accumulates a subset of the inputs too, and every intermediate step of it is again a subset, but it is seeded with the neutral element, which carries none of the range `init` provided.

Which inputs a partial receives is decided by the tiling, so the grouping determines *whether* a given input triggers the problem. It is not a second cause: regrouping on its own cannot break a flag that held for every order, but using the neutral element instead of `init` makes this a computation that could not have happened in the untiled form, which is what the overflow flag refers to.

The merge is likewise safe on its own since it starts from `init`, so when it overflows as well, it does so because a partial has already wrapped.

## Reproducing with `arith.subi`

The PR explicitly maintains flag propagation when adding support for `subi` so as to be consistent with existing reducers, but it points out that maintaining these flags is not always safe. These are theoretical edge cases I ran into:

### Counterexample: `arith.subi` with `nsw`

Let's work with `i8` and have a subtracting accumulation with
```
i8 range: [-128, 127]
init = 101
inputs = [100, 0, 100, 0] // can be applied in any order
```

Regardless of how we permute the inputs, the value of `init` keeps the accumulation in the range `-99` to `101`, and never overflows: `101 - 100 - 100 - 0 - 0 = -99`. If this is tiled by `2`, then we have 2 partials (each taking inputs strided by 2) starting from neutral element `0`, and a combiner:

```
p_0 = 0 - 100 - 100 = -200 // out of [-128, 127], `nsw` violated; wraps to 56
p_1 = 0 - 0 - 0 = 0
sum = 101 + 56 = 157 // out of [-128, 127], `nsw` violated; wraps to -99
```

Here is the above example as a repro case that can be run through `mlir-opt %s -transform-interpreter -test-transform-dialect-erase-schedule -canonicalize -cse`
```mlir
#in_map = affine_map<(d0) -> (d0)>
#out_map = affine_map<(d0) -> ()>

module {
func.func @sub_reduce_nsw() -> tensor {
%input = arith.constant dense<[100, 0, 100, 0]> : tensor<4xi8>
%init = arith.constant dense<101> : tensor

%result = linalg.generic {
indexing_maps = [#in_map, #out_map],
iterator_types = ["reduction"]
}
ins(%input : tensor<4xi8>)
outs(%init : tensor) {
^bb0(%in: i8, %acc: i8):
%next = arith.subi %acc, %in overflow : i8
linalg.yield %next : i8
} -> tensor

return %result : tensor
}

module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
%0 = transform.structured.match ops{["linalg.generic"]} in %arg0
: (!transform.any_op) -> !transform.any_op
%1, %2, %3, %loop = transform.structured.tile_reduction_using_for %0 by tile_sizes = [2]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
transform.yield
}
}
}
```
This is assuming that the above PR is merged so that it transforms properly. Here is an abbreviated output showing the pattern from above where the two `overflow` flags are violated.
```mlir
%1 = linalg.fill ins(%c0_i8 : i8) outs(%0 : tensor<2xi8>) -> tensor<2xi8>
%2 = scf.for %arg0 = %c0 to %c4 step %c2 iter_args(%arg1 = %1) -> (tensor<2xi8>) {
%extracted_slice = tensor.extract_slice %cst[%arg0] [2] [1] : tensor<4xi8> to tensor<2xi8>
%3 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%extracted_slice : tensor<2xi8>) outs(%arg1 : tensor<2xi8>) {
^bb0(%in: i8, %out: i8):
%4 = arith.subi %out, %in overflow : i8
linalg.yield %4 : i8
} -> tensor<2xi8>
scf.yield %3 : tensor<2xi8>
}
%reduced = linalg.reduce ins(%2 : tensor<2xi8>) outs(%cst_0 : tensor) dimensions = [0]
(%in: i8, %init: i8) {
%3 = arith.addi %in, %init overflow : i8
linalg.yield %3 : i8
}
return %reduced : tensor
```
The flags survive lowering, so the assertion reaches the backend unchanged:

```mlir
%83 = llvm.sub %82, %79 overflow : i8
%125 = llvm.add %122, %124 overflow : i8
```

### Counterexample: `arith.subi` with `nuw`

`nuw` is easier to show since the neutral element is `0`, the `0 - x` fails `nuw` for any non-zero `x`.

```
i8 range: [-128, 127]
init = 100
inputs = [50, 30, 10, 5] // sum to 95 so non-tiled will not overflow when subtracted from `100`
p_0 = 0 - 50 // fail on the very first partial
```

Again, the untiled case is fine since all permutations of the inputs stay within range `[5, 100]`. But the very first subtraction in the partial immediately underflows.

Here's the full repro case:
```mlir
#in_map = affine_map<(d0) -> (d0)>
#out_map = affine_map<(d0) -> ()>

module {
func.func @sub_reduce_nuw() -> tensor {
%input = arith.constant dense<[50, 30, 10, 5]> : tensor<4xi8>
%init = arith.constant dense<100> : tensor

%result = linalg.generic {
indexing_maps = [#in_map, #out_map],
iterator_types = ["reduction"]
}
ins(%input : tensor<4xi8>)
outs(%init : tensor) {
^bb0(%in: i8, %acc: i8):
%next = arith.subi %acc, %in overflow : i8
linalg.yield %next : i8
} -> tensor

return %result : tensor
}

module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
%0 = transform.structured.match ops{["linalg.generic"]} in %arg0
: (!transform.any_op) -> !transform.any_op
%1, %2, %3, %loop = transform.structured.tile_reduction_using_for %0 by tile_sizes = [2]
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
transform.yield
}
}
}
```

Tiled by `2`, the partials are seeded with `linalg.fill` of `0` and keep the overflow flag:
```mlir
%1 = linalg.fill ins(%c0_i8 : i8) outs(%0 : tensor<2xi8>) -> tensor<2xi8>
%2 = scf.for %arg0 = %c0 to %c4 step %c2 iter_args(%arg1 = %1) -> (tensor<2xi8>) {
...
%4 = arith.subi %out, %in overflow : i8
```
Every `nuw` flagged subtracting accumulation with a non-zero input has this problem, regardless of `init` and of the tile size.

## `arith.addi` and `arith.muli`

The same issue applies to the combiners that partial reduction tiling has always supported. This isn't new to subtracting reduction.

**`addi overflow`** on `i8`:
```
i8 range: [-128, 127]
init = -100
inputs = [100, -14, 100, -14] // untiled stays in [-128, 100] for all permutations
p_0 = 0 + 100 + 100 = 200 // out of [-128, 127], `nsw` violated; wraps to -56
p_1 = 0 - 14 - 14 = -28
sum = -100 + -56 // = -156; out of [-128, 127], `nsw` violated; wraps to 100
+ -28 = 72
```

**`muli overflow`** on `i32`:
```
i32 range: [-2^31, 2^31-1]
init = 0
inputs = [2^20, 1, 2^20, 1]
```
Multiplying by `init` first pins the untiled accumulator to `0` whatever the order, so `nsw` is trivially satisfied. The partials are seeded with the neutral element `1` instead,
```mlir
%1 = linalg.fill ins(%c1_i32 : i32) outs(%0 : tensor<2xi32>) -> tensor<2xi32>
...
%4 = arith.muli %out, %in overflow : i32
```
so `p_0 = 1 * 2^20 * 2^20 = 2^40`, which overflows `i32`. The merge multiplies by `init` and returns `0` as before, but the partial is poison.

## Why testing does not catch it

An execution test cannot tell the overflow flag was violated because the wrapping still results in the correct bit pattern. The damage is the unjustified `nsw` / `nuw` on the emitted IR. It only becomes observable once a later pass exploits the flag, at which point the reduction can be folded to something the source program never permitted.

The simplest fix (which is the one suggested in the original PR) is to simply drop the overflow flags if they can't be guaranteed. But this would affect several reducers which currently just use `clone`, so is a larger discussion.

Contributor guide

Open the contributing guide

Research direction

Start with the tileToPartialReduction and mergeReductions entry points described in the issue, then run the supplied mlir-opt reproducer for the nsw and nuw cases. Trace how the combiner's overflow flags reach the partial and merge operations. Done means the tiled reduction no longer carries overflow assertions that the transformed computation cannot justify, with coverage for the shown cases.

Written by the indexing model from the issue text.

Assessment

Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.