llvm-project

b9d15ca7 - [mlir][acc] Decide reduction atomics by destination instead of by partitioning (#235099)

[mlir][acc] Decide reduction atomics by destination instead of by partitioning (#235099) A reduction combine became atomic only if the reduction's parallel dims named a block dimension, so the choice rested entirely on how the reduction was partitioned. Partitioning answers whether a merge is needed at all. Whether that merge has to be atomic should depend instead on how many blocks can reach the destination, which is a property of the launch. A thread-only reduction shows the difference: it names no block dimension, yet every block still combines its own partial into the same shared slot, so the combine stayed non-atomic, leaving those updates open to a data race. An array accumulate partitioned only across threads avoided emitting that unsynchronized merge by reporting not implemented whenever the launch had a block dimension. The test was on the dimension's presence and not its width, so it also rejected launches whose block dimension is one wide and therefore partitioned nothing. Repro: ```mlir // repro.mlir func.func @thread_only_array_reduction() { %c8_pw = arith.constant 8 : index %c128 = arith.constant 128 : index %bx = acc.par_width %c8_pw par_dim(#acc.par_dim<block_x>) %tx = acc.par_width %c128 par_dim(#acc.par_dim<thread_x>) acc.compute_region launch(%kbx = %bx, %ktx = %tx) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c8 = arith.constant 8 : index %c0_i32 = arith.constant 0 : i32 %local = memref.alloca() : memref<8xi32> scf.for %i = %c0 to %c8 step %c1 { memref.store %c0_i32, %local[%i] : memref<8xi32> } %bounds = acc.bounds extent(%c8 : index) acc.reduction_accumulate_array %local bounds(%bounds) <add> par_dims(#acc<par_dims[thread_x]>) : memref<8xi32> acc.yield } <{origin = "acc.parallel"}> return } ``` ``` mlir-opt repro.mlir --pass-pipeline="builtin.module(func.func(acc-cg-to-gpu))" ``` This reported not implemented prior to this change: ``` error: not yet implemented: reduction: thread-only array reduction accumulate ``` ### Fix Decide the atomic from how many blocks can reach the destination. A block dimension in the reduction implies more than one, and so does a thread-only reduction in a multi-block launch; a thread-private destination is reached by a single thread either way. A launch width that is not a compile-time constant is counted conservatively as more than one block. The accumulate in turn skips dimensions launched one wide, since those partition nothing. With the merge across blocks now handled correctly by the combine, the accumulate has no case left that it cannot lower, so the diagnostic is removed. The same command now lowers to a per-element reduction over the array: ```mlir scf.for %arg12 = %c0_2 to %1 step %c1_3 { %2 = memref.load %alloca[%arg12] : memref<8xi32> %3 = gpu.all_reduce add %2 uniform { } {acc.par_dims = #acc<par_dims[thread_x]>} : (i32) -> i32 memref.store %3, %alloca[%arg12] : memref<8xi32> } ``` The combine change is not array-specific. It closes the same exposure for scalar thread-only reductions in multi-block launches, which were already reachable and lowered to a plain store into a shared slot. One conservative choice is left in place in the combine: a reduction whose own parallel dims name a block dimension takes the atomic path even when that dimension is one wide. Narrowing that to the launch width should be a separate follow-up.