[Mlir-commits] [mlir] [mlir][linalg] Support subtracting accumulation in partial reduction … (PR #214033)

[Federico Bruzzone]

FedericoBruzzone wrote:

> @FedericoBruzzone You asked about overflow flags, and that was really fair. There was definitely an issue here:
> 
> * `overflow<nuw>`: `0 - x` wraps for every non-zero `x`, the result being negative.
> * `overflow<nsw>`: `0 - x` wraps for `INT_MIN`.
> 
> I added new lit tests with overflow flags. The problem is handled by declining to treat a flagged subtraction as an accumulation (so it stays untiled instead of doing it wrong).
> 
> I got Claude to analyze this a bit and it concluded that there was an existing issue with other combiners. It's outside the scope of what I'm trying to fix, but I'll include the analysis here (better to have Claude's verbose description than have me summarize and get it subtly wrong) and let you decide if there should be a new issue to address this. It all sounds reasonable to me, but I saw this code for the first time two days ago, so am missing a lot of context 😄.
> 
> Checking the overflow flags on `arith.subi` turned up the same hazard on the existing `arith.addi` / `arith.muli` path.
> 
> ## The hazard
> `tileToPartialReduction` clones the combiner op into the partial reduction body, overflow flags included. `mergeReductions` likewise clones it to combine the partial results.
> 
> `nsw`/`nuw` are a promise about _one specific computation_: this addition, on these values, does not wrap. Tiling replaces that computation with a different one — each partial result accumulates a **strided subset** of the inputs starting from the neutral element instead of from the real init value. The promise does not carry over, but the flag does.
> 
> If the flag is wrong, the result is poison, so later passes are entitled to optimize on an assumption that no longer holds.
> 
> ## Counterexample for `addi overflow<nsw>`
> `i32`, init `0`, four inputs, tile size 2:
> 
> ```
> x = [2^30, -2^30, 2^30, -2^30]
> ```
> 
> Untiled, the running accumulator is `2^30, 0, 2^30, 0` — every intermediate is within `i32`, so `nsw` holds and the op is well formed.
> 
> Tiled, partial column `j` accumulates input columns `j`, `j + 2`, ... , so:
> 
> ```
> p_0 = x_0 + x_2 = 2^30 + 2^30 = 2^31   // overflows i32 signed, nsw violated
> p_1 = x_1 + x_3 = -2^30 + -2^30        // overflows i32 signed, nsw violated
> ```
> 
> The strided grouping is confirmed by the emitted IR — the loop steps by the tile size over the reduction dimension while the partial slice is always taken at offset 0:
> 
> ```mlir
> %2 = scf.for %arg2 = %c0 to %c4 step %c2 iter_args(%arg3 = %1) -> (tensor<1x2xi32>) {
>   %extracted_slice = tensor.extract_slice %arg0[0, %arg2] [1, 2] [1, 1] : tensor<1x4xi32> to tensor<1x2xi32>
>   ...
>     %4 = arith.addi %out, %in overflow<nsw> : i32
> ```
> 
> and the flag is propagated to the merge as well:
> 
> ```mlir
> %reduced = linalg.reduce ins(%2 : tensor<1x2xi32>) outs(%arg1 : tensor<1xi32>) dimensions = [1]
>     %3 = arith.addi %in, %init overflow<nsw> : i32
> ```
> 
> ## Why there is no runnable repro
> Wrapping two's complement arithmetic is associative and commutative modulo `2^n`, so the tiled and untiled reductions compute the **same bit pattern**. An execution test cannot distinguish them.
> 
> The damage is not a wrong number, it is an unjustified `nsw`/`nuw` on the emitted IR. It only becomes visible once a later pass exploits the flag — after lowering to LLVM IR, for example. That makes this hard to catch by testing and worth fixing at the source.
> 
> ## Possible fix
> Drop the overflow flags when cloning the combiner in `tileToPartialReduction` and `mergeReductions`.
> 
> Dropping a flag is always safe: it removes an assumption rather than adding one. The cost is that partial and merged accumulations lose the optimization opportunity the flag would have provided.
> 
> This would also make the subtraction case tileable again — the PR currently bails on flagged `arith.subi` precisely because the cloned flag would be wrong, and stripping it removes the reason to bail.

Thanks, bailing out is the right call here.

Please correct me if I'm missing something, but `p_1` doesn't overflow (`-2^30 + -2^30 = -2^31 = INT_MIN` is _representable_), so the counterexample rests on `p_0` alone. Right? 
And `addi nuw` is safe, so scope the issue to `nsw`.

It also misses `muli`, broken for both flags. For instance, consider: i32, `init = 1`, `x = [2^20, 0, 2^20, 1]`, tile size 2.



https://github.com/llvm/llvm-project/pull/214033