[mlir][gpu] Add NoMemoryEffect to gpu.subgroup_reduce op (#180227)

fabrizio-indirli · web-flow · commit 368d7ea12c0b · 2026-02-06T17:41:26.000Z
Add missing MemoryEffectOpInterface to gpu.subgroup_reduce

Signed-off-by: Fabrizio Indirli &lt;fabrizio.indirli@arm.com&gt;
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
@@ -1226,7 +1226,7 @@ def GPU_AllReduceOp : GPU_Op<"all_reduce",
 def AnyIntegerOrFloatOr1DVector :
   AnyTypeOf<[AnyIntegerOrFloat, FixedVectorOfRankAndType<[1], [AnyIntegerOrFloat]>]>;
 
-def GPU_SubgroupReduceOp : GPU_Op<"subgroup_reduce", [SameOperandsAndResultType]> {
+def GPU_SubgroupReduceOp : GPU_Op<"subgroup_reduce", [SameOperandsAndResultType, NoMemoryEffect]> {
   let summary = "Reduce values among subgroup.";
   let description = [{
     The `subgroup_reduce` op reduces the values of lanes (work items) across a
diff --git a/mlir/test/Conversion/ConvertToSPIRV/gpu.mlir b/mlir/test/Conversion/ConvertToSPIRV/gpu.mlir
@@ -27,11 +27,12 @@ module attributes {
 
 gpu.module @kernels {
   // CHECK-LABEL: spirv.func @subgroup_reduce
-  // CHECK-SAME: (%[[ARG0:.*]]: f32)
+  // CHECK-SAME: (%[[ARG0:.*]]: f32, [[BUF:.*]]: !spirv.ptr{{[^)]*}})
   // CHECK: %{{.*}} = spirv.GroupNonUniformFAdd <Subgroup> <Reduce> %[[ARG0]] : f32 -> f32
-  gpu.func @subgroup_reduce(%arg0 : f32) kernel
+  gpu.func @subgroup_reduce(%arg0 : f32, %buf : memref<f32>) kernel
     attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
     %reduced = gpu.subgroup_reduce add %arg0 {} : (f32) -> (f32)
+    memref.store %reduced, %buf[] : memref<f32>
     gpu.return
   }
 }
diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
@@ -644,39 +644,45 @@ gpu.module @test_module_29 {
 
 gpu.module @test_module_30 {
   // CHECK-LABEL: func @subgroup_reduce_add
-  gpu.func @subgroup_reduce_add(%arg0 : i32) {
+  gpu.func @subgroup_reduce_add(%arg0 : i32, %buf : memref<i32>) {
     // CHECK: nvvm.redux.sync add {{.*}}
     %result = gpu.subgroup_reduce add %arg0 uniform {} : (i32) -> (i32)
+    memref.store %result, %buf[] : memref<i32>
     gpu.return
   }
   // CHECK-LABEL: @subgroup_reduce_minsi
-  gpu.func @subgroup_reduce_minsi(%arg0 : i32) {
+  gpu.func @subgroup_reduce_minsi(%arg0 : i32, %buf : memref<i32>) {
     // CHECK: nvvm.redux.sync min {{.*}}
     %result = gpu.subgroup_reduce minsi %arg0 uniform {} : (i32) -> (i32)
+    memref.store %result, %buf[] : memref<i32>
     gpu.return
   }
   // CHECK-LABEL:  @subgroup_reduce_maxsi
-  gpu.func @subgroup_reduce_maxsi(%arg0 : i32) {
+  gpu.func @subgroup_reduce_maxsi(%arg0 : i32, %buf : memref<i32>) {
     // CHECK: nvvm.redux.sync max {{.*}}
     %result = gpu.subgroup_reduce maxsi %arg0 uniform {} : (i32) -> (i32)
+    memref.store %result, %buf[] : memref<i32>
     gpu.return
   }
   // CHECK-LABEL: func @subgroup_reduce_and
-  gpu.func @subgroup_reduce_and(%arg0 : i32) {
+  gpu.func @subgroup_reduce_and(%arg0 : i32, %buf : memref<i32>) {
     // CHECK: nvvm.redux.sync and {{.*}}
     %result = gpu.subgroup_reduce and %arg0 uniform {} : (i32) -> (i32)
+    memref.store %result, %buf[] : memref<i32>
     gpu.return
   }
   // CHECK-LABEL:  @subgroup_reduce_or
-  gpu.func @subgroup_reduce_or(%arg0 : i32) {
+  gpu.func @subgroup_reduce_or(%arg0 : i32, %buf : memref<i32>) {
     // CHECK: nvvm.redux.sync or {{.*}}
     %result = gpu.subgroup_reduce or %arg0 uniform {} : (i32) -> (i32)
+    memref.store %result, %buf[] : memref<i32>
     gpu.return
   }
   // CHECK-LABEL: @subgroup_reduce_xor
-  gpu.func @subgroup_reduce_xor(%arg0 : i32) {
+  gpu.func @subgroup_reduce_xor(%arg0 : i32, %buf : memref<i32>) {
     // CHECK: nvvm.redux.sync xor {{.*}}
     %result = gpu.subgroup_reduce xor %arg0 uniform {} : (i32) -> (i32)
+    memref.store %result, %buf[] : memref<i32>
     gpu.return
   }
 }

Original file line number	Diff line number	Diff line change
`@@ -27,11 +27,12 @@ module attributes {`
`27`	`27`
`28`	`28`	`gpu.module @kernels {`
`29`	`29`	`// CHECK-LABEL: spirv.func @subgroup_reduce`
`30`		`- // CHECK-SAME: (%[[ARG0:.*]]: f32)`
	`30`	`+ // CHECK-SAME: (%[[ARG0:.]]: f32, [[BUF:.]]: !spirv.ptr{{[^)]*}})`
`31`	`31`	`// CHECK: %{{.*}} = spirv.GroupNonUniformFAdd <Subgroup> <Reduce> %[[ARG0]] : f32 -> f32`
`32`		`- gpu.func @subgroup_reduce(%arg0 : f32) kernel`
	`32`	`+ gpu.func @subgroup_reduce(%arg0 : f32, %buf : memref<f32>) kernel`
`33`	`33`	`attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {`
`34`	`34`	`%reduced = gpu.subgroup_reduce add %arg0 {} : (f32) -> (f32)`
	`35`	`+ memref.store %reduced, %buf[] : memref<f32>`
`35`	`36`	`gpu.return`
`36`	`37`	`}`
`37`	`38`	`}`