[mlir][gpu] Reduction ops canonicalizatios

[Hardcode84]

Summary

Make group ops uniform if gpu.launch is their direct parent.

[Hardcode84]

rebase

[Hardcode84]

ping

[ThomasRaoux]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

435–438

this is only correct if we assume all control flow is structured control flow right? I believe LaunchOp region is allowed to have multiple blocks. I wonder if we should check that the op belongs to the first block of LaunchOp region instead?

[Hardcode84]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

435–438

I believe LaunchOp region is allowed to have multiple blocks

Does it really? I'm only using structured cf so I've never checked it.

I will add a check so it will only do fold for entry blocks then.

[Hardcode84]

rebase, only convert ops in entry block.

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

let description = [{
 
     If `uniform` flag is set either none or all work items of a workgroup
     need to execute this op in convergence.
   }];
   let regions = (region AnyRegion:$body);
   let assemblyFormat = [{ custom<AllReduceOperation>($op) $value
                           (`uniform` $uniform^)? $body attr-dict
                           `:` functional-type(operands, results) }];
+
+  let hasFolder = 1;
   let hasRegionVerifier = 1;
 }
 
 def GPU_SubgroupReduceOp : GPU_Op<"subgroup_reduce",
     [SameOperandsAndResultType]>,
     Arguments<(ins AnyType:$value,
                GPU_AllReduceOperationAttr:$op,
                UnitAttr:$uniform)>,
let description = [{
     ```
 
     If `uniform` flag is set either none or all work items of a subgroup
     need to execute this op in convergence.
   }];
   let assemblyFormat = [{ custom<AllReduceOperation>($op) $value
                           (`uniform` $uniform^)? attr-dict
                           `:` functional-type(operands, results) }];
+
+  let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 def GPU_ShuffleOpXor  : I32EnumAttrCase<"XOR",  0, "xor">;
 def GPU_ShuffleOpDown : I32EnumAttrCase<"DOWN", 1, "down">;
 def GPU_ShuffleOpUp   : I32EnumAttrCase<"UP",   2, "up">;
 def GPU_ShuffleOpIdx  : I32EnumAttrCase<"IDX",  3, "idx">;
 

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

if (!verifyReduceOpAndType(opName, getType())) {
       return emitError()
              << '`' << gpu::stringifyAllReduceOperation(opName)
              << "` accumulator is only compatible with Integer type";
     }
   }
   return success();
 }
 
+static bool canMakeGroupOpUniform(Operation *op) {
+  auto launchOp = dyn_cast<gpu::LaunchOp>(op->getParentOp());
+  if (!launchOp)
+    return false;
+
+  Region &body = launchOp.getBody();
+  assert(!body.empty() && "Invalid region");
+
+  // Only convert ops in gpu::launch entry block for now.
+  return op->getBlock() == &body.front();
+}
+
+OpFoldResult gpu::AllReduceOp::fold(FoldAdaptor /*adaptor*/) {
+  if (!getUniform() && canMakeGroupOpUniform(*this)) {
+    setUniform(true);
+    return getResult();
+  }
+
+  return nullptr;
+}
+
 // TODO: Support optional custom attributes (without dialect prefix).
 static ParseResult parseAllReduceOperation(AsmParser &parser,
                                            AllReduceOperationAttr &attr) {
   StringRef enumStr;
   if (!parser.parseOptionalKeyword(&enumStr)) {
     std::optional<AllReduceOperation> op =
         gpu::symbolizeAllReduceOperation(enumStr);
     if (!op)
LogicalResult gpu::SubgroupReduceOp::verify() {
   gpu::AllReduceOperation opName = getOp();
   if (!verifyReduceOpAndType(opName, getType())) {
     return emitError() << '`' << gpu::stringifyAllReduceOperation(opName)
                        << "` accumulator is only compatible with Integer type";
   }
   return success();
 }
 
+OpFoldResult gpu::SubgroupReduceOp::fold(FoldAdaptor /*adaptor*/) {
+  if (!getUniform() && canMakeGroupOpUniform(*this)) {
+    setUniform(true);
+    return getResult();
+  }
+
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // AsyncOpInterface
 //===----------------------------------------------------------------------===//
 
 void gpu::addAsyncDependency(Operation *op, Value token) {
   op->insertOperands(0, {token});
   if (!op->template hasTrait<OpTrait::AttrSizedOperandSegments>())
     return;

mlir/test/Dialect/GPU/canonicalize.mlir

gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1)
     gpu.terminator
   }
   return
 }
 
 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
 // CHECK: gpu.launch blocks(%{{.*}}, %{{.*}}, %{{.*}}) in (%{{.*}} = %[[C1]], %{{.*}} = %[[C1]], %{{.*}} = %[[C1]]) threads(%[[TIDX:.*]], %{{.*}}, %{{.*}}) in (%{{.*}} = %c32, %{{.*}} = %[[C1]], %{{.*}} = %[[C1]]) {
-// CHECK-NEXT:    arith.divui %[[TIDX]], %c32 : index
-// CHECK-NEXT:    arith.muli %{{.*}}, %c2 : index
+// CHECK-NEXT:    arith.divui %[[TIDX]], %c32 : index
+// CHECK-NEXT:    arith.muli %{{.*}}, %c2 : index
 // CHECK-NEXT:    memref.load %memref[%{{.*}}, %[[C0]], %[[C0]]] : memref<2x16x16xf32>
 // CHECK-NEXT:    arith.addi %{{.*}}, %[[C1]] : index
 // CHECK-NEXT:    memref.load %memref[%{{.*}}, %[[C0]], %[[C0]]] : memref<2x16x16xf32>
 // CHECK-NEXT:    memref.store %{{.*}}, %memref[%{{.*}}, %[[C0]], %[[C0]]] : memref<2x16x16xf32>
 // CHECK-NEXT:    memref.store %{{.*}}, %memref[%{{.*}}, %[[C0]], %[[C0]]] : memref<2x16x16xf32>
 // CHECK-NEXT:    gpu.terminator
 // CHECK-NEXT:  }
+
+// -----
+
+// CHECK-LABEL: func @make_reduce_uniform
+//       CHECK: gpu.launch blocks
+//       CHECK: %[[V1:.*]] = "test.test2"() : () -> i32
+//       CHECK: %[[V2:.*]] = gpu.all_reduce add %[[V1]] uniform {
+//       CHECK: "test.test3"(%[[V2]]) : (i32) -> ()
+func.func @make_reduce_uniform() {
+  %0:6 = "test.test1"() : () -> (index, index, index, index, index, index)
+  gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %0#0, %arg7 = %0#1, %arg8 = %0#2)
+    threads(%arg3, %arg4, %arg5) in (%arg9 = %0#3, %arg10 = %0#4, %arg11 = %0#5) {
+    %1 = "test.test2"() : () -> i32
+    %2 = gpu.all_reduce add %1 {} : (i32) -> (i32)
+    "test.test3"(%2) : (i32) -> ()
+    gpu.terminator
+  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @make_subgroup_reduce_uniform
+//       CHECK: gpu.launch blocks
+//       CHECK: %[[V1:.*]] = "test.test2"() : () -> i32
+//       CHECK: %[[V2:.*]] = gpu.subgroup_reduce add %[[V1]] uniform
+//       CHECK: "test.test3"(%[[V2]]) : (i32) -> ()
+func.func @make_subgroup_reduce_uniform() {
+  %0:6 = "test.test1"() : () -> (index, index, index, index, index, index)
+  gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %0#0, %arg7 = %0#1, %arg8 = %0#2)
+    threads(%arg3, %arg4, %arg5) in (%arg9 = %0#3, %arg10 = %0#4, %arg11 = %0#5) {
+    %1 = "test.test2"() : () -> i32
+    %2 = gpu.subgroup_reduce add %1 : (i32) -> (i32)
+    "test.test3"(%2) : (i32) -> ()
+    gpu.terminator
+  }
+  return
+}