[mlir][gpu] Add `subgroup_reduce` operation

[Hardcode84]

Summary

Introduce subgroup_reduce operation, similar to all_reduce, but operating on subgroup scope instead of workgroup.
It is intended as low-level building block for more high level abstractions (e.g for workgroup-wide all_reduce ops).
Only introduce version taking reduce operation enum for simplicity sake.

[Hardcode84]

fix typos

[ThomasRaoux]

are you planning to add some lowering for it? This can be useful to at least "enforce the semantic". For instance it should be easy to have a lowering of this shuffle ops.

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

727–728

can you clarify the behavior when some lanes in the subgroups are inactive?

[Hardcode84]

are you planning to add some lowering for it? This can be useful to at least "enforce the semantic". For instance it should be easy to have a lowering of this shuffle ops.

We are planning to lower them directly to GroupNonUniformXYZ spirv ops (those ops allows both Workgroup and Subgroup scope, but Intel Level Zero/Intel Graphics Compiler we are targeting only supports Subgroup version at the moment).

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

727–728

spirv ops like GroupNonUniformFAddOp we are targeting allow non-uniform execution, but we don't need it at the moment. I will add uniform requirement similar to all_reduce. We can always add something like non-uniform flag to these ops later if needed.

[Hardcode84]

add uniform requirement

[ThomasRaoux]

LGTM

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

let description = [{
     in convergence.
   }];
   let regions = (region AnyRegion:$body);
   let assemblyFormat = [{ custom<AllReduceOperation>($op) $value $body attr-dict
                           `:` functional-type(operands, results) }];
   let hasRegionVerifier = 1;
 }
 
+def GPU_SubgroupReduceOp : GPU_Op<"subgroup_reduce",
+    [SameOperandsAndResultType]>,
+    Arguments<(ins AnyType:$value,
+               GPU_AllReduceOperationAttr:$op)>,
+    Results<(outs AnyType)> {
+  let summary = "Reduce values among subgroup.";
+  let description = [{
+    The `subgroup_reduce` op reduces the value of every work item across a
+    subgroup. The result is equal for all work items of a subgroup.
+
+    Example:
+
+    ```mlir
+    %1 = gpu.subgroup_reduce add %0 : (f32) -> (f32)
+    ```
+
+    Either none or all work items of a subgroup need to execute this op
+    in convergence.
+  }];
+  let assemblyFormat = [{ custom<AllReduceOperation>($op) $value attr-dict
+                          `:` functional-type(operands, results) }];
+  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">;
 
 def GPU_ShuffleMode : I32EnumAttr<"ShuffleMode",
     "Indexing modes supported by gpu.shuffle.",
     [

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

static void printAsyncDependencies(OpAsmPrinter &printer, Operation *op,
   llvm::interleaveComma(asyncDependencies, printer);
   printer << ']';
 }
 
 //===----------------------------------------------------------------------===//
 // AllReduceOp
 //===----------------------------------------------------------------------===//
 
+static bool verifyReduceOpAndType(gpu::AllReduceOperation opName,
+                                  Type resType) {
+  if ((opName == gpu::AllReduceOperation::AND ||
+       opName == gpu::AllReduceOperation::OR ||
+       opName == gpu::AllReduceOperation::XOR) &&
+      !resType.isa<IntegerType>())
+    return false;
+
+  return true;
+}
+
 LogicalResult gpu::AllReduceOp::verifyRegions() {
   if (getBody().empty() != getOp().has_value())
     return emitError("expected either an op attribute or a non-empty body");
   if (!getBody().empty()) {
     if (getBody().getNumArguments() != 2)
       return emitError("expected two region arguments");
     for (auto argument : getBody().getArguments()) {
       if (argument.getType() != getType())
         return emitError("incorrect region argument type");
     }
     unsigned yieldCount = 0;
     for (Block &block : getBody()) {
       if (auto yield = dyn_cast<gpu::YieldOp>(block.getTerminator())) {
         if (yield.getNumOperands() != 1)
           return emitError("expected one gpu.yield operand");
         if (yield.getOperand(0).getType() != getType())
           return emitError("incorrect gpu.yield type");
         ++yieldCount;
       }
     }
     if (yieldCount == 0)
       return emitError("expected gpu.yield op in region");
   } else {
     gpu::AllReduceOperation opName = *getOp();
-    if ((opName == gpu::AllReduceOperation::AND ||
-         opName == gpu::AllReduceOperation::OR ||
-         opName == gpu::AllReduceOperation::XOR) &&
-        !getType().isa<IntegerType>()) {
+    if (!verifyReduceOpAndType(opName, getType())) {
       return emitError()
              << '`' << gpu::stringifyAllReduceOperation(opName)
              << "` accumulator is only compatible with Integer type";
     }
   }
   return success();
 }
 

 
 static void printAllReduceOperation(AsmPrinter &printer, Operation *op,
                                     AllReduceOperationAttr attr) {
   if (attr)
     attr.print(printer);
 }
 
 //===----------------------------------------------------------------------===//
+// SubgroupReduceOp
+//===----------------------------------------------------------------------===//
+
+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();
+}
+
+//===----------------------------------------------------------------------===//
 // AsyncOpInterface
 //===----------------------------------------------------------------------===//
 
 void gpu::addAsyncDependency(Operation *op, Value token) {
   op->insertOperands(0, {token});
   if (!op->template hasTrait<OpTrait::AttrSizedOperandSegments>())
     return;
   auto attrName =

mlir/test/Dialect/GPU/invalid.mlir

 func.func @reduce_invalid_op_type(%arg0 : f32) {
   // expected-error@+1 {{`and` accumulator is only compatible with Integer type}}
   %res = gpu.all_reduce and %arg0 {} : (f32) -> (f32)
   return
 }
 
 // -----
 
+func.func @subgroup_reduce_invalid_op_type(%arg0 : f32) {
+  // expected-error@+1 {{`and` accumulator is only compatible with Integer type}}
+  %res = gpu.subgroup_reduce and %arg0 : (f32) -> (f32)
+  return
+}
+
+// -----
+
 func.func @reduce_incorrect_region_arguments(%arg0 : f32) {
   // expected-error@+1 {{expected two region arguments}}
   %res = gpu.all_reduce %arg0 {
   ^bb(%lhs : f32):
     "gpu.yield"(%lhs) : (f32) -> ()
   } : (f32) -> (f32)
   return
 }

mlir/test/Dialect/GPU/ops.mlir

gpu.func @kernel_1(%arg0 : f32, %arg1 : memref<?xf32, 1>) kernel {
 
       %sgId = gpu.subgroup_id : index
       %numSg = gpu.num_subgroups : index
       %SgSi = gpu.subgroup_size : index
 
       %one = arith.constant 1.0 : f32
       %sum = gpu.all_reduce add %one {} : (f32) -> (f32)
 
+      // CHECK: %{{.*}} = gpu.subgroup_reduce add %{{.*}} : (f32) -> f32
+      %sum_subgroup = gpu.subgroup_reduce add %one : (f32) -> f32
+
       %width = arith.constant 7 : i32
       %offset = arith.constant 3 : i32
       // CHECK: gpu.shuffle xor %{{.*}}, %{{.*}}, %{{.*}} : f32
       %shfl, %pred = gpu.shuffle xor %arg0, %offset, %width : f32
       // CHECK: gpu.shuffle up %{{.*}}, %{{.*}}, %{{.*}} : f32
       %shfl1, %pred1 = gpu.shuffle up %arg0, %offset, %width : f32
       // CHECK: gpu.shuffle down %{{.*}}, %{{.*}}, %{{.*}} : f32
       %shfl2, %pred2 = gpu.shuffle down %arg0, %offset, %width : f32