[Mlir-commits] [mlir] [mlir][acc] Reduce a complex accumulator componentwise (PR #234276)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Oct 21 09:49:41 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Vijay Kandiah (VijayKandiah)
<details>
<summary>Changes</summary>
`gpu.all_reduce` accepts only integer and float operands, including in its
custom-region form. A thread-partitioned `acc.reduction_accumulate` or
`acc.reduction_accumulate_array` with a complex accumulator therefore produced
IR that failed to verify:
```
error: 'gpu.all_reduce' op operand #<!-- -->0 must be Integer or Float, but got 'complex<f32>'
```
MLIR repro:
```mlir
// repro.mlir
func.func @<!-- -->complex_reduction(%slot: memref<complex<f32>>) {
%c128 = arith.constant 128 : index
%tx = acc.par_width %c128 par_dim(#acc.par_dim<thread_x>)
acc.compute_region launch(%ktx = %tx) ins(%s = %slot) : (memref<complex<f32>>) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%one = arith.constant 1.0 : f32
%v = complex.create %one, %one : complex<f32>
scf.parallel (%i) = (%c0) to (%ktx) step (%c1) {
acc.reduction_accumulate %v to %s <add> par_dims(#acc<par_dims[thread_x]>) : complex<f32> -> memref<complex<f32>>
scf.reduce
} {acc.par_dims = #acc<par_dims[thread_x]>}
acc.yield
} <{origin = "acc.parallel"}>
return
}
```
```
mlir-opt repro.mlir --pass-pipeline="builtin.module(func.func(acc-cg-to-gpu))"
```
This change reduces the real and imaginary parts separately and rebuilds the value from the
two results:
```mlir
%1 = complex.re %0 : complex<f32>
%2 = gpu.all_reduce add %1 uniform {
} {acc.par_dims = #acc<par_dims[thread_x]>} : (f32) -> f32
%3 = complex.im %0 : complex<f32>
%4 = gpu.all_reduce add %3 uniform {
} {acc.par_dims = #acc<par_dims[thread_x]>} : (f32) -> f32
%5 = complex.create %2, %4 : complex<f32>
```
Addition is the only operator for which this rewrite is valid, since each
component of a sum depends only on that component of the inputs.
Multiplication mixes the two components, and min and max have no ordering on
complex at all. The split is therefore allowed for floating-point addition
alone; every other operator on a complex accumulator reports not implemented
rather than being split in a way that would silently give a wrong answer.
Reporting not-implemented also marks the conversion as failed, so the pass
leaves the `acc.reduction_accumulate` in place and fails rather than erasing
the accumulate and emitting incomplete IR.
---
Full diff: https://github.com/llvm/llvm-project/pull/234276.diff
3 Files Affected:
- (modified) mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp (+38-4)
- (added) mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-complex-unsupported.mlir (+35)
- (added) mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-complex.mlir (+111)
``````````diff
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
index fd99bd71795d6c..db017f6a1852a9 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCCGToGPU.cpp
@@ -3568,8 +3568,38 @@ void ACCCGToGPULowering::createGPUAllReduceOp(
bool isPerThreadPrivateTarget) {
gpu::AllReduceOperationAttr attr = gpu::AllReduceOperationAttr::get(
computeRegion->getContext(), getAllReduceOperation(kind));
- auto allReduceOp = gpu::AllReduceOp::create(rewriter, loc, input, attr, true);
- mlir::acc::setParDimsAttr(allReduceOp, parDimsAttr);
+ auto emitAllReduce = [&](Value operand) -> Value {
+ auto allReduceOp =
+ gpu::AllReduceOp::create(rewriter, loc, operand, attr, true);
+ mlir::acc::setParDimsAttr(allReduceOp, parDimsAttr);
+ return allReduceOp;
+ };
+
+ Value reducedValue;
+ if (auto complexType = dyn_cast<ComplexType>(input.getType())) {
+ // gpu.all_reduce takes only integers and floats, so a complex value has to
+ // go through as components. Addition is the one operator reaching here
+ // where that is equivalent: multiplication mixes the components, and
+ // min/max have no ordering on complex at all.
+ if (kind != arith::AtomicRMWKind::addf) {
+ std::string msg;
+ llvm::raw_string_ostream os(msg);
+ os << "reduction: " << arith::stringifyAtomicRMWKind(kind) << " of "
+ << complexType << " across threads";
+ (void)accSupport.emitNYI(loc, os.str());
+ hasFailed = true;
+ return;
+ }
+ Type elementType = complexType.getElementType();
+ Value realResult =
+ emitAllReduce(complex::ReOp::create(rewriter, loc, elementType, input));
+ Value imagResult =
+ emitAllReduce(complex::ImOp::create(rewriter, loc, elementType, input));
+ reducedValue = complex::CreateOp::create(rewriter, loc, complexType,
+ realResult, imagResult);
+ } else {
+ reducedValue = emitAllReduce(input);
+ }
// Predicate the store on the thread-level dimensions being reduced so that
// only one thread per reduced group writes the result. Only dimensions in
// parDimsAttr are included; sweeping over all dimensions between the highest
@@ -3613,13 +3643,13 @@ void ACCCGToGPULowering::createGPUAllReduceOp(
Block &thenBlock = thenRegion.back();
rewriter.setInsertionPoint(thenBlock.getTerminator());
}
- memref::StoreOp::create(rewriter, loc, allReduceOp, memref, indices);
+ memref::StoreOp::create(rewriter, loc, reducedValue, memref, indices);
if (predicate && !isPerThreadPrivate)
rewriter.setInsertionPointAfter(ifOp);
// A later block combine reuses this instead of reloading the slot. This only
// applies to scalar accumulators; array elements are indexed individually.
if (indices.empty())
- reductionAccumValue[memref] = allReduceOp;
+ reductionAccumValue[memref] = reducedValue;
}
void ACCCGToGPULowering::postprocessAccumulateOp(
@@ -3666,6 +3696,10 @@ void ACCCGToGPULowering::postprocessAccumulateOp(
} else if (hasThreadDim) {
createGPUAllReduceOp(loc, op.getValue(), op.getMemref(), *kind,
op.getParDims());
+ // Nothing replaced the accumulate, so leave it in place rather than
+ // erasing it below.
+ if (hasFailed)
+ return;
} else {
// Block-only: no gpu.all_reduce needed (all threads have the same
// value after broadcast). Just store the value to the memref.
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-complex-unsupported.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-complex-unsupported.mlir
new file mode 100644
index 00000000000000..62a4e0ecc4c745
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-complex-unsupported.mlir
@@ -0,0 +1,35 @@
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(acc-cg-to-gpu))" -verify-diagnostics
+
+// A complex accumulator is reduced one component at a time, which only gives
+// the same answer when each component depends on itself alone. Multiplication
+// mixes the two, so it is rejected rather than reduced componentwise.
+
+func.func @complex_mul_thread_reduction(%arg_slot: memref<complex<f32>>) {
+ %c1_pw = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %bx = acc.par_width %c1_pw par_dim(#acc.par_dim<block_x>)
+ %tx = acc.par_width %c128 par_dim(#acc.par_dim<thread_x>)
+ // expected-error @below {{failed to legalize operation 'acc.compute_region'}}
+ acc.compute_region launch(%kbx = %bx, %ktx = %tx) ins(%a_slot = %arg_slot) : (memref<complex<f32>>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+ %one = arith.constant 1.0 : f32
+ %init = complex.create %one, %one : complex<f32>
+ %contrib = complex.create %one, %one : complex<f32>
+ scf.parallel (%tx_iv) = (%c0) to (%ktx) step (%c1) {
+ %inner_red = scf.parallel (%i) = (%tx_iv) to (%c8) step (%ktx) init (%init) -> complex<f32> {
+ scf.reduce(%contrib : complex<f32>) {
+ ^bb0(%lhs: complex<f32>, %rhs: complex<f32>):
+ %prod = complex.mul %lhs, %rhs : complex<f32>
+ scf.reduce.return %prod : complex<f32>
+ }
+ } {acc.par_dims = #acc<par_dims[sequential]>}
+ // expected-error @below {{not yet implemented: reduction: mulf of complex<f32> across threads}}
+ acc.reduction_accumulate %inner_red to %a_slot <mul> par_dims(#acc<par_dims[thread_x]>) : complex<f32> -> memref<complex<f32>>
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[thread_x]>}
+ acc.yield
+ } <{origin = "acc.parallel"}>
+ return
+}
diff --git a/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-complex.mlir b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-complex.mlir
new file mode 100644
index 00000000000000..b2ae75e69ea7b5
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/acc-cg-to-gpu-reduction-complex.mlir
@@ -0,0 +1,111 @@
+// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(acc-cg-to-gpu))" | FileCheck %s
+
+// A cross-thread reduction takes integers and floats, so a complex accumulator
+// is reduced one component at a time and rebuilt from the two results. Addition
+// is componentwise, so reducing the parts separately gives the same answer.
+
+// CHECK-LABEL: func.func @complex_thread_reduction
+// CHECK: gpu.launch
+// CHECK: %[[RE:.*]] = complex.re %[[VAL:.*]] : complex<f32>
+// CHECK: %[[RERED:.*]] = gpu.all_reduce add %[[RE]]
+// CHECK: %[[IM:.*]] = complex.im %[[VAL]] : complex<f32>
+// CHECK: %[[IMRED:.*]] = gpu.all_reduce add %[[IM]]
+// CHECK: %[[RES:.*]] = complex.create %[[RERED]], %[[IMRED]] : complex<f32>
+// CHECK: memref.store %[[RES]]
+
+func.func @complex_thread_reduction(%arg_slot: memref<complex<f32>>) {
+ %c1_pw = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %bx = acc.par_width %c1_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) ins(%a_slot = %arg_slot) : (memref<complex<f32>>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c8 = arith.constant 8 : index
+ %zero = arith.constant 0.0 : f32
+ %one = arith.constant 1.0 : f32
+ %init = complex.create %zero, %zero : complex<f32>
+ %contrib = complex.create %one, %one : complex<f32>
+ scf.parallel (%tx_iv) = (%c0) to (%ktx) step (%c1) {
+ %inner_red = scf.parallel (%i) = (%tx_iv) to (%c8) step (%ktx) init (%init) -> complex<f32> {
+ scf.reduce(%contrib : complex<f32>) {
+ ^bb0(%lhs: complex<f32>, %rhs: complex<f32>):
+ %sum = complex.add %lhs, %rhs : complex<f32>
+ scf.reduce.return %sum : complex<f32>
+ }
+ } {acc.par_dims = #acc<par_dims[sequential]>}
+ acc.reduction_accumulate %inner_red to %a_slot <add> par_dims(#acc<par_dims[thread_x]>) : complex<f32> -> memref<complex<f32>>
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[thread_x]>}
+ acc.yield
+ } <{origin = "acc.parallel"}>
+ return
+}
+
+// -----
+
+// An array accumulator reduces element by element, and each element goes
+// through the same componentwise path, so a complex array needs no separate
+// handling.
+
+// CHECK-LABEL: func.func @complex_array_reduction
+// CHECK: gpu.launch
+// CHECK-NOT: acc.reduction_accumulate_array
+// CHECK: scf.for %[[IV:.*]] = %{{.*}} to %{{.*}} step %{{.*}} {
+// CHECK: %[[ELT:.*]] = memref.load %[[ALLOCA:.*]][%[[IV]]] : memref<2xcomplex<f32>>
+// CHECK: %[[RE:.*]] = complex.re %[[ELT]] : complex<f32>
+// CHECK: %[[RERED:.*]] = gpu.all_reduce add %[[RE]]
+// CHECK: %[[IM:.*]] = complex.im %[[ELT]] : complex<f32>
+// CHECK: %[[IMRED:.*]] = gpu.all_reduce add %[[IM]]
+// CHECK: %[[RES:.*]] = complex.create %[[RERED]], %[[IMRED]] : complex<f32>
+// CHECK: memref.store %[[RES]], %[[ALLOCA]][%[[IV]]] : memref<2xcomplex<f32>>
+// CHECK: }
+
+func.func @complex_array_reduction(%arg0: memref<2xcomplex<f32>>) {
+ %0 = acc.copyin varPtr(%arg0 : memref<2xcomplex<f32>>) dataClause(acc_reduction) implicit(true) name("r") -> memref<2xcomplex<f32>>
+ acc.kernel_environment dataOperands(%0 : memref<2xcomplex<f32>>) {
+ %c1_pw = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %bx = acc.par_width %c1_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) ins(%arg2 = %0) : (memref<2xcomplex<f32>>) {
+ %c2 = arith.constant 2 : index
+ %c1 = arith.constant 1 : index
+ %c0 = arith.constant 0 : index
+ %zero = arith.constant 0.0 : f32
+ %one = arith.constant 1.0 : f32
+ %zc = complex.create %zero, %zero : complex<f32>
+ %oc = complex.create %one, %one : complex<f32>
+ %2 = acc.reduction_init %arg2 <add> : memref<2xcomplex<f32>> {
+ %alloca = memref.alloca() : memref<2xcomplex<f32>>
+ scf.parallel (%i) = (%c0) to (%c2) step (%c1) {
+ memref.store %zc, %alloca[%i] : memref<2xcomplex<f32>>
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[thread_x]>}
+ acc.yield %alloca : memref<2xcomplex<f32>>
+ }
+ scf.parallel (%bx_iv) = (%c0) to (%kbx) step (%c1) {
+ scf.parallel (%tx_iv) = (%c0) to (%ktx) step (%c1) {
+ %3 = memref.load %2[%c0] : memref<2xcomplex<f32>>
+ %4 = complex.add %3, %oc : complex<f32>
+ memref.store %4, %2[%c0] : memref<2xcomplex<f32>>
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[thread_x]>}
+ scf.reduce
+ } {acc.par_dims = #acc<par_dims[block_x]>}
+ %b = acc.bounds extent(%c2 : index)
+ acc.reduction_accumulate_array %2 bounds(%b) <add> par_dims(#acc<par_dims[block_x, thread_x]>) : memref<2xcomplex<f32>>
+ acc.reduction_combine_region %2 into %arg2 : memref<2xcomplex<f32>> {
+ scf.for %i = %c0 to %c2 step %c1 {
+ %3 = memref.load %2[%i] : memref<2xcomplex<f32>>
+ %4 = memref.load %arg2[%i] : memref<2xcomplex<f32>>
+ %5 = complex.add %3, %4 : complex<f32>
+ memref.store %5, %arg2[%i] : memref<2xcomplex<f32>>
+ }
+ }
+ acc.yield
+ } <{origin = "acc.parallel"}>
+ }
+ acc.copyout accPtr(%0 : memref<2xcomplex<f32>>) to varPtr(%arg0 : memref<2xcomplex<f32>>) dataClause(acc_reduction) implicit(true) name("r")
+ return
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/234276
More information about the Mlir-commits
mailing list