[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

690

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

399

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

690

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

399

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

690

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

399

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:$non_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 `non_uniform` flag is not 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
+                          (`non_uniform` $non_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:$non_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 `non_uniform` flag is not 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
+                          (`non_uniform` $non_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) {
+      if (reduceOp.getNonUniform()) {
+        reduceOp.emitWarning("Non uniform reductions are not supported yet.");
+        return WalkResult::advance();
+      }
+
       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.
       return WalkResult::interrupt();
     };
     while (funcOp.walk(callback).wasInterrupted()) {
     }
     return success();
   }
 };
 } // namespace
 
 void mlir::populateGpuAllReducePatterns(RewritePatternSet &patterns) {
   patterns.add<GpuAllReduceConversion>(patterns.getContext());
 }

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 %{{.*}} non_uniform {
+      // CHECK-NEXT: } : (f32) -> f32
+      %sum1 = gpu.all_reduce add %one non_uniform {} : (f32) -> f32
+
       // CHECK: %{{.*}} = gpu.subgroup_reduce add %{{.*}} : (f32) -> f32
       %sum_subgroup = gpu.subgroup_reduce add %one : (f32) -> f32
 
+      // CHECK: %{{.*}} = gpu.subgroup_reduce add %{{.*}} non_uniform : (f32) -> f32
+      %sum_subgroup1 = gpu.subgroup_reduce add %one non_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