[mlir][gpu] Add `uniform` flag to gpu reduction ops

[Hardcode84]

ping

[ThomasRaoux]

Sorry for the slow review. Why are we adding it without any lowering? The goal is to lower to SPIRV that has those semantics?

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

692

I wonder if the flag should be reverted and be uniform instead. non_uniform should be the default as it is the most conservative case. This way unless specified the op is created without the uniform assumption.
What do you think?

mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp

398

Should this be an error?

[Hardcode84]

In D138758#3989123, @ThomasRaoux wrote:

Sorry for the slow review. Why are we adding it without any lowering? The goal is to lower to SPIRV that has those semantics?

I was planning to add spirv lowering in separate review.

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

692

The idea itself is ok to me, but:

• We are changing default semantics and it can broke user code without any warning
• Historically, at least in spirv, uniform ops appeared much earlier than non-uniform (i.e. GroupFAdd is 1.0 but GroupNonUniformFAdd is 1.3)

mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp

398

There is no good way to return error from rewrite, I think. Even if we use emitError it will only change error message, but won't propagate it.

[ThomasRaoux]

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

692

using non-uniform version should be conservative and correct at the GPU level. The part that does break is the lowering and there should be an explicit lowering failure there.

Historically, at least in spirv, uniform ops appeared much earlier than non-uniform (i.e. GroupFAdd is 1.0 but GroupNonUniformFAdd is 1.3)

I wonder if this is due to the order in which those got added. If there were added at the same time and/or we didn't have to keep compatibility I would guess that's not how they would have been done.

mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp

398

Yes, in general we should just fail the pattern, here because this pattern tries to convert the whole function this is a bit awkward.
I think the best would be to fail the pattern and return notifyMatchFailure.

[Hardcode84]

rebase, invert flag, fix pattern

[Hardcode84]

Fix for integration tests https://reviews.llvm.org/D140014

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

def GPU_AllReduceOperation : I32EnumAttr<"AllReduceOperation",
   let cppNamespace = "::mlir::gpu";
 }
 def GPU_AllReduceOperationAttr : EnumAttr<GPU_Dialect, GPU_AllReduceOperation,
                                           "all_reduce_op">;
 
 def GPU_AllReduceOp : GPU_Op<"all_reduce",
     [SameOperandsAndResultType, IsolatedFromAbove]>,
     Arguments<(ins AnyType:$value,
-               OptionalAttr<GPU_AllReduceOperationAttr>:$op)>,
+               OptionalAttr<GPU_AllReduceOperationAttr>:$op,
+               UnitAttr:$uniform)>,
     Results<(outs AnyType)> {
   let summary = "Reduce values among workgroup.";
   let description = [{
     The `all_reduce` op reduces the value of every work item across a local
     workgroup. The result is equal for all work items of a workgroup.
 
     For example, both
 
     ```mlir
     %1 = gpu.all_reduce add %0 {} : (f32) -> (f32)
     %2 = gpu.all_reduce %0 {
     ^bb(%lhs : f32, %rhs : f32):
       %sum = arith.addf %lhs, %rhs : f32
       "gpu.yield"(%sum) : (f32) -> ()
     } : (f32) -> (f32)
     ```
 
     compute the sum of each work item's %0 value. The first version specifies
     the accumulation as operation, whereas the second version specifies the
     accumulation as code region. The accumulation operation must be one of:
     `add`, `and`, `max`, `min`, `mul`, `or`, `xor`.
 
-    Either none or all work items of a workgroup need to execute this op
-    in convergence.
+    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 $body attr-dict
+  let assemblyFormat = [{ custom<AllReduceOperation>($op) $value
+                          (`uniform` $uniform^)? $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)>,
+               GPU_AllReduceOperationAttr:$op,
+               UnitAttr:$uniform)>,
     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.
+    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 attr-dict
+  let assemblyFormat = [{ custom<AllReduceOperation>($op) $value
+                          (`uniform` $uniform^)? 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">;

mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp

 
 struct GpuAllReduceConversion : public RewritePattern {
   explicit GpuAllReduceConversion(MLIRContext *context)
       : RewritePattern(gpu::GPUFuncOp::getOperationName(), 1, context) {}
 
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     auto funcOp = cast<gpu::GPUFuncOp>(op);
-    auto callback = [&](gpu::AllReduceOp reduceOp) {
-      GpuAllReduceRewriter(funcOp, reduceOp, rewriter).rewrite();
-      // Performing a rewrite invalidates the walk iterator. Report interrupt
-      // so that we can start a new walk until all all_reduce ops are replaced.
+
+    SmallVector<gpu::AllReduceOp> reduceOps;
+    auto callback = [&](gpu::AllReduceOp reduceOp) -> WalkResult {
+      if (!reduceOp.getUniform())
         return WalkResult::interrupt();
+
+      reduceOps.emplace_back(reduceOp);
+      return WalkResult::advance();
     };
-    while (funcOp.walk(callback).wasInterrupted()) {
-    }
+
+    if (funcOp.walk(callback).wasInterrupted())
+      return rewriter.notifyMatchFailure(
+          op, "Non uniform reductions are not supported yet.");
+
+    for (gpu::AllReduceOp reduceOp : reduceOps)
+      GpuAllReduceRewriter(funcOp, reduceOp, rewriter).rewrite();
+
     return success();
   }
 };
 } // namespace
 
 void mlir::populateGpuAllReducePatterns(RewritePatternSet &patterns) {
   patterns.add<GpuAllReduceConversion>(patterns.getContext());
 }

mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir

 gpu.module @test_module {
   // CHECK-LABEL: func @gpu_all_reduce_op()
   gpu.func @gpu_all_reduce_op() {
     %arg0 = arith.constant 1.0 : f32
     // TODO: Check full IR expansion once lowering has settled.
     // CHECK: nvvm.shfl.sync bfly {{.*}}
     // CHECK: nvvm.barrier0
     // CHECK: llvm.fadd
-    %result = gpu.all_reduce add %arg0 {} : (f32) -> (f32)
+    %result = gpu.all_reduce add %arg0 uniform {} : (f32) -> (f32)
 
     gpu.return
   }
 }
 
 // -----
 
 gpu.module @test_module {
   // CHECK-LABEL: func @gpu_all_reduce_region()
   gpu.func @gpu_all_reduce_region() {
     %arg0 = arith.constant 1 : i32
     // TODO: Check full IR expansion once lowering has settled.
     // CHECK: nvvm.shfl.sync bfly {{.*}}
     // CHECK: nvvm.barrier0
-    %result = gpu.all_reduce %arg0 {
+    %result = gpu.all_reduce %arg0 uniform {
     ^bb(%lhs : i32, %rhs : i32):
       %xor = arith.xori %lhs, %rhs : i32
       "gpu.yield"(%xor) : (i32) -> ()
     } : (i32) -> (i32)
     gpu.return
   }
 }
 

mlir/test/Dialect/GPU/all-reduce-max.mlir

gpu.func @kernel(%arg0 : f32) kernel {
     // CHECK:   cf.br ^bb40([[VAL_132]] : f32)
     // CHECK: ^bb40([[VAL_133:%.*]]: f32):
     // CHECK:   store [[VAL_133]], [[VAL_1]]{{\[}}[[VAL_4]]] : memref<32xf32, 3>
     // CHECK:   cf.br ^bb42
     // CHECK: ^bb41:
     // CHECK:   cf.br ^bb42
     // CHECK: ^bb42:
     // CHECK:   gpu.barrier
-    %sum = gpu.all_reduce max %arg0 {} : (f32) -> (f32)
+    %sum = gpu.all_reduce max %arg0 uniform {} : (f32) -> (f32)
     gpu.return
   }
 
 }

mlir/test/Dialect/GPU/all-reduce.mlir

gpu.func @kernel(%arg0 : f32) kernel {
     // CHECK:   cf.br ^bb40([[VAL_112]] : f32)
     // CHECK: ^bb40([[VAL_113:%.*]]: f32):
     // CHECK:   store [[VAL_113]], [[VAL_1]]{{\[}}[[VAL_4]]] : memref<32xf32, 3>
     // CHECK:   cf.br ^bb42
     // CHECK: ^bb41:
     // CHECK:   cf.br ^bb42
     // CHECK: ^bb42:
     // CHECK:   gpu.barrier
-    %sum = gpu.all_reduce add %arg0 {} : (f32) -> (f32)
+    %sum = gpu.all_reduce add %arg0 uniform {} : (f32) -> (f32)
     gpu.return
   }
 
 }

mlir/test/Dialect/GPU/multiple-all-reduce.mlir

 // RUN: mlir-opt --gpu-kernel-outlining --convert-gpu-to-nvvm %s | FileCheck %s
 
 func.func @main() {
   %data = memref.alloc() : memref<2x6xf32>
   %sum = memref.alloc() : memref<2xf32>
   %mul = memref.alloc() : memref<2xf32>
   %c1 = arith.constant 1 : index
 
   // ADD + MUL
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1, %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %c1, %block_y = %c1, %block_z = %c1) {
     %val = memref.load %data[%bx, %tx] : memref<2x6xf32>
-    %reduced0 = gpu.all_reduce add %val {} : (f32) -> (f32)
+    %reduced0 = gpu.all_reduce add %val uniform {} : (f32) -> (f32)
     memref.store %reduced0, %sum[%bx] : memref<2xf32>
-    %reduced1 = gpu.all_reduce mul %val {} : (f32) -> (f32)
+    %reduced1 = gpu.all_reduce mul %val uniform {} : (f32) -> (f32)
     memref.store %reduced1, %mul[%bx] : memref<2xf32>
     gpu.terminator
   }
 
 // CHECK:      gpu.module @main_kernel {
 // CHECK-NEXT:   llvm.mlir.global internal @{{.*}}() {addr_space = 3 : i32} : !llvm.array<32 x f32>
 // CHECK-NEXT:   llvm.mlir.global internal @{{.*}}() {addr_space = 3 : i32} : !llvm.array<32 x f32>
 
   return
 }

mlir/test/Dialect/GPU/ops.mlir

gpu.func @kernel_1(%arg0 : f32, %arg1 : memref<?xf32, 1>) kernel {
       %gIdY = gpu.global_id y
       %gIdZ = gpu.global_id z
 
       %sgId = gpu.subgroup_id : index
       %numSg = gpu.num_subgroups : index
       %SgSi = gpu.subgroup_size : index
 
       %one = arith.constant 1.0 : f32
+
+      // CHECK: %{{.*}} = gpu.all_reduce add %{{.*}} {
+      // CHECK-NEXT: } : (f32) -> f32
       %sum = gpu.all_reduce add %one {} : (f32) -> (f32)
 
+      // CHECK: %{{.*}} = gpu.all_reduce add %{{.*}} uniform {
+      // CHECK-NEXT: } : (f32) -> f32
+      %sum1 = gpu.all_reduce add %one uniform {} : (f32) -> f32
+
       // CHECK: %{{.*}} = gpu.subgroup_reduce add %{{.*}} : (f32) -> f32
       %sum_subgroup = gpu.subgroup_reduce add %one : (f32) -> f32
 
+      // CHECK: %{{.*}} = gpu.subgroup_reduce add %{{.*}} uniform : (f32) -> f32
+      %sum_subgroup1 = gpu.subgroup_reduce add %one uniform : (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