[mlir][gpu] Add 'cluster_stride' attribute to gpu.subgroup_reduce (#107142)

andfau-amd · web-flow · commit 3d01f0a33b9a · 2024-09-05T09:03:22.000-04:00
Follow-up to 7aa22f0, adding an additional attribute needed in some applications.
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
@@ -1200,10 +1200,12 @@ def GPU_SubgroupReduceOp : GPU_Op<"subgroup_reduce", [SameOperandsAndResultType]
     The `subgroup_reduce` op reduces the values of lanes (work items) across a
     subgroup.
 
-    The subgroup is divided into clusters of `cluster_size` contiguous lanes
-    each, and a reduction is done for every lane of each cluster (in parallel).
-    The result is equal for all lanes in a cluster. When `cluster_size` is
-    omitted, there is a single cluster covering the entire subgroup.
+    The subgroup is divided into clusters starting at lane index 0. Within each
+    cluster, there are `size` lanes, and the lane index advances by `stride`.
+    A reduction is done for each cluster in parallel: every lane in the cluster
+    is reduced, and the result is equal for all lanes in the cluster. If `size`
+    is omitted, there is a single cluster covering the entire subgroup. If
+    `stride` is omitted, the stride is 1 (the cluster's lanes are contiguous).
 
     When the reduced value is of a vector type, each vector element is reduced
     independently. Only 1-d vector types are allowed.
@@ -1213,7 +1215,8 @@ def GPU_SubgroupReduceOp : GPU_Op<"subgroup_reduce", [SameOperandsAndResultType]
     ```mlir
     %1 = gpu.subgroup_reduce add %a : (f32) -> f32
     %2 = gpu.subgroup_reduce add %b : (vector<4xf16>) -> vector<4xf16>
-    %3 = gpu.subgroup_reduce add %c cluster_size(4) : (f32) -> f32
+    %3 = gpu.subgroup_reduce add %c cluster(size = 4) : (f32) -> f32
+    %3 = gpu.subgroup_reduce add %c cluster(size = 4, stride = 2) : (f32) -> f32
     ```
 
     If `uniform` flag is set either none or all lanes of a subgroup need to execute
@@ -1230,27 +1233,38 @@ def GPU_SubgroupReduceOp : GPU_Op<"subgroup_reduce", [SameOperandsAndResultType]
     AnyIntegerOrFloatOr1DVector:$value,
     GPU_AllReduceOperationAttr:$op,
     UnitAttr:$uniform,
-    OptionalAttr<I32Attr>:$cluster_size
+    OptionalAttr<I32Attr>:$cluster_size,
+    DefaultValuedAttr<I32Attr,"1">:$cluster_stride
   );
   let results = (outs AnyIntegerOrFloatOr1DVector:$result);
 
   let builders = [
     OpBuilder<(ins "Value":$value,
                "::mlir::gpu::AllReduceOperation":$op,
                "bool":$uniform), [{
-      build($_builder, $_state, value, op, uniform, /*cluster_size=*/ nullptr);
+      build($_builder, $_state, value, op, uniform, std::nullopt);
     }]>,
     OpBuilder<(ins "Value":$value,
                "::mlir::gpu::AllReduceOperation":$op,
                "bool":$uniform,
                "std::optional<uint32_t>":$cluster_size), [{
-      build($_builder, $_state, value, op, uniform, cluster_size ? $_builder.getI32IntegerAttr(*cluster_size) : nullptr);
+      build($_builder, $_state, value, op, uniform,
+            cluster_size ? $_builder.getI32IntegerAttr(*cluster_size) : nullptr);
+    }]>,
+    OpBuilder<(ins "Value":$value,
+               "::mlir::gpu::AllReduceOperation":$op,
+               "bool":$uniform,
+               "std::optional<uint32_t>":$cluster_size,
+               "uint32_t":$cluster_stride), [{
+      build($_builder, $_state, value, op, uniform,
+            cluster_size ? $_builder.getI32IntegerAttr(*cluster_size) : nullptr,
+            cluster_stride);
     }]>
   ];
 
   let assemblyFormat = [{ custom<AllReduceOperation>($op) $value
                           (`uniform` $uniform^)?
-                          (`cluster_size` `(` $cluster_size^ `)`)?
+                          (`cluster` `(` `size` `=` $cluster_size^ (`,` `stride` `=` $cluster_stride^)? `)`)?
                           attr-dict
                           `:` functional-type(operands, results) }];
 
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -621,14 +621,25 @@ LogicalResult gpu::SubgroupReduceOp::verify() {
                        << getType();
   }
 
-  if (auto clusterSize = getClusterSize()) {
+  auto clusterSize = getClusterSize();
+  if (clusterSize) {
     uint32_t size = *clusterSize;
     if (!llvm::isPowerOf2_32(size)) {
       return emitOpError() << "cluster size " << size
                            << " is not a power of two";
     }
   }
 
+  uint32_t stride = getClusterStride();
+  if (stride != 1 && !clusterSize) {
+    return emitOpError() << "cluster stride can only be specified if cluster "
+                            "size is specified";
+  }
+  if (!llvm::isPowerOf2_32(stride)) {
+    return emitOpError() << "cluster stride " << stride
+                         << " is not a power of two";
+  }
+
   return success();
 }
 
diff --git a/mlir/lib/Dialect/GPU/Transforms/SubgroupReduceLowering.cpp b/mlir/lib/Dialect/GPU/Transforms/SubgroupReduceLowering.cpp
@@ -50,8 +50,6 @@ struct BreakDownSubgroupReduce final : OpRewritePattern<gpu::SubgroupReduceOp> {
 
   LogicalResult matchAndRewrite(gpu::SubgroupReduceOp op,
                                 PatternRewriter &rewriter) const override {
-    std::optional<uint32_t> clusterSize = op.getClusterSize();
-
     auto vecTy = dyn_cast<VectorType>(op.getType());
     if (!vecTy || vecTy.getNumElements() < 2)
       return rewriter.notifyMatchFailure(op, "not a multi-element reduction");
@@ -97,7 +95,8 @@ struct BreakDownSubgroupReduce final : OpRewritePattern<gpu::SubgroupReduceOp> {
       }
 
       Value reduce = rewriter.create<gpu::SubgroupReduceOp>(
-          loc, extracted, op.getOp(), op.getUniform(), clusterSize);
+          loc, extracted, op.getOp(), op.getUniform(), op.getClusterSize(),
+          op.getClusterStride());
       if (numElems == 1) {
         res = rewriter.create<vector::InsertOp>(loc, reduce, res, startIdx);
         continue;
@@ -129,8 +128,6 @@ struct ScalarizeSingleElementReduce final
 
   LogicalResult matchAndRewrite(gpu::SubgroupReduceOp op,
                                 PatternRewriter &rewriter) const override {
-    std::optional<uint32_t> clusterSize = op.getClusterSize();
-
     auto vecTy = dyn_cast<VectorType>(op.getType());
     if (!vecTy || vecTy.getNumElements() != 1)
       return rewriter.notifyMatchFailure(op, "not a single-element reduction");
@@ -140,34 +137,64 @@ struct ScalarizeSingleElementReduce final
     Location loc = op.getLoc();
     Value extracted = rewriter.create<vector::ExtractOp>(loc, op.getValue(), 0);
     Value reduce = rewriter.create<gpu::SubgroupReduceOp>(
-        loc, extracted, op.getOp(), op.getUniform(), clusterSize);
+        loc, extracted, op.getOp(), op.getUniform(), op.getClusterSize(),
+        op.getClusterStride());
     rewriter.replaceOpWithNewOp<vector::BroadcastOp>(op, vecTy, reduce);
     return success();
   }
 };
 
+struct ClusterInfo {
+  unsigned clusterStride;
+  unsigned clusterSize;
+  unsigned subgroupSize;
+};
+
+static FailureOr<ClusterInfo>
+getAndValidateClusterInfo(gpu::SubgroupReduceOp op, unsigned subgroupSize) {
+  assert(llvm::isPowerOf2_32(subgroupSize));
+
+  std::optional<uint32_t> clusterSize = op.getClusterSize();
+  assert(!clusterSize ||
+         llvm::isPowerOf2_32(*clusterSize)); // Verifier should've caught this.
+  if (clusterSize && *clusterSize > subgroupSize)
+    return op.emitOpError()
+           << "cluster size " << *clusterSize
+           << " is greater than subgroup size " << subgroupSize;
+  unsigned effectiveClusterSize = clusterSize.value_or(subgroupSize);
+
+  auto clusterStride = op.getClusterStride();
+  assert(llvm::isPowerOf2_32(clusterStride)); // Verifier should've caught this.
+  if (clusterStride >= subgroupSize)
+    return op.emitOpError()
+           << "cluster stride " << clusterStride
+           << " is not less than subgroup size " << subgroupSize;
+
+  return ClusterInfo{clusterStride, effectiveClusterSize, subgroupSize};
+}
+
 /// Emits a subgroup reduction using a sequence of shuffles. Uses the `packFn`
 /// and `unpackFn` to convert to the native shuffle type and to the reduction
 /// type, respectively. For example, with `input` of type `f16`, `packFn` could
 /// build ops to cast the value to `i32` to perform shuffles, while `unpackFn`
 /// would cast it back to `f16` to perform arithmetic reduction on. Assumes that
 /// the subgroup is `subgroupSize` lanes wide and divides it into clusters of
-/// `clusterSize` lanes, reducing all lanes in each cluster in parallel.
-static Value createSubgroupShuffleReduction(
-    OpBuilder &builder, Location loc, Value input, gpu::AllReduceOperation mode,
-    unsigned clusterSize, unsigned subgroupSize,
-    function_ref<Value(Value)> packFn, function_ref<Value(Value)> unpackFn) {
-  assert(llvm::isPowerOf2_32(clusterSize));
-  assert(llvm::isPowerOf2_32(subgroupSize));
-  assert(clusterSize <= subgroupSize);
+/// `clusterSize` lanes starting at lane 0 with a stride of `clusterStride` for
+/// lanes within a cluster, reducing all lanes in each cluster in parallel.
+Value createSubgroupShuffleReduction(OpBuilder &builder, Location loc,
+                                     Value input, gpu::AllReduceOperation mode,
+                                     const ClusterInfo &ci,
+                                     function_ref<Value(Value)> packFn,
+                                     function_ref<Value(Value)> unpackFn) {
   // Lane value always stays in the original type. We use it to perform arith
   // reductions.
   Value laneVal = input;
   // Parallel reduction using butterfly shuffles.
-  for (unsigned i = 1; i < clusterSize; i <<= 1) {
+  for (unsigned i = ci.clusterStride; i < ci.clusterStride * ci.clusterSize;
+       i <<= 1) {
     Value shuffled = builder
                          .create<gpu::ShuffleOp>(loc, packFn(laneVal), i,
-                                                 /*width=*/subgroupSize,
+                                                 /*width=*/ci.subgroupSize,
                                                  /*mode=*/gpu::ShuffleMode::XOR)
                          .getShuffleResult();
     laneVal = vector::makeArithReduction(builder, loc,
@@ -190,12 +217,9 @@ struct ScalarSubgroupReduceToShuffles final
 
   LogicalResult matchAndRewrite(gpu::SubgroupReduceOp op,
                                 PatternRewriter &rewriter) const override {
-    std::optional<uint32_t> clusterSize = op.getClusterSize();
-    if (clusterSize && *clusterSize > subgroupSize)
-      return op.emitOpError()
-             << "cluster size " << *clusterSize
-             << " is greater than subgroup size " << subgroupSize;
-    unsigned effectiveClusterSize = clusterSize.value_or(subgroupSize);
+    auto ci = getAndValidateClusterInfo(op, subgroupSize);
+    if (failed(ci))
+      return failure();
 
     Type valueTy = op.getType();
     unsigned elemBitwidth =
@@ -209,9 +233,8 @@ struct ScalarSubgroupReduceToShuffles final
     if (elemBitwidth == shuffleBitwidth) {
       auto identityFn = [](Value v) { return v; };
       rewriter.replaceOp(op, createSubgroupShuffleReduction(
-                                 rewriter, loc, op.getValue(), op.getOp(),
-                                 effectiveClusterSize, subgroupSize, identityFn,
-                                 identityFn));
+                                 rewriter, loc, op.getValue(), op.getOp(), *ci,
+                                 identityFn, identityFn));
       return success();
     }
 
@@ -232,8 +255,7 @@ struct ScalarSubgroupReduceToShuffles final
 
     rewriter.replaceOp(
         op, createSubgroupShuffleReduction(rewriter, loc, op.getValue(),
-                                           op.getOp(), effectiveClusterSize,
-                                           subgroupSize, packFn, unpackFn));
+                                           op.getOp(), *ci, packFn, unpackFn));
     return success();
   }
 
@@ -253,12 +275,9 @@ struct VectorSubgroupReduceToShuffles final
 
   LogicalResult matchAndRewrite(gpu::SubgroupReduceOp op,
                                 PatternRewriter &rewriter) const override {
-    std::optional<uint32_t> clusterSize = op.getClusterSize();
-    if (clusterSize && *clusterSize > subgroupSize)
-      return op.emitOpError()
-             << "cluster size " << *clusterSize
-             << " is greater than subgroup size " << subgroupSize;
-    unsigned effectiveClusterSize = clusterSize.value_or(subgroupSize);
+    auto ci = getAndValidateClusterInfo(op, subgroupSize);
+    if (failed(ci))
+      return failure();
 
     auto vecTy = dyn_cast<VectorType>(op.getType());
     if (!vecTy)
@@ -308,9 +327,8 @@ struct VectorSubgroupReduceToShuffles final
       return rewriter.create<vector::BitCastOp>(loc, extendedVecTy, asIntVec);
     };
 
-    Value res = createSubgroupShuffleReduction(rewriter, loc, extendedInput,
-                                               op.getOp(), effectiveClusterSize,
-                                               subgroupSize, packFn, unpackFn);
+    Value res = createSubgroupShuffleReduction(
+        rewriter, loc, extendedInput, op.getOp(), *ci, packFn, unpackFn);
 
     if (vecBitwidth < shuffleBitwidth) {
       res = rewriter.create<vector::ExtractStridedSliceOp>(
diff --git a/mlir/test/Dialect/GPU/canonicalize.mlir b/mlir/test/Dialect/GPU/canonicalize.mlir
@@ -255,7 +255,7 @@ func.func @subgroup_reduce_cluster_size_1() {
   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 cluster_size(1) : (i32) -> (i32)
+    %2 = gpu.subgroup_reduce add %1 cluster(size=1) : (i32) -> (i32)
     "test.test3"(%2) : (i32) -> ()
     gpu.terminator
   }
diff --git a/mlir/test/Dialect/GPU/invalid.mlir b/mlir/test/Dialect/GPU/invalid.mlir
@@ -335,18 +335,35 @@ func.func @reduce_invalid_op_type_maximumf(%arg0 : i32) {
 
 func.func @subgroup_reduce_zero_cluster_size(%arg0 : vector<4xf32>) {
   // expected-error@+1 {{cluster size 0 is not a power of two}}
-  %res = gpu.subgroup_reduce add %arg0 cluster_size(0) : (vector<4xf32>) -> vector<4xf32>
+  %res = gpu.subgroup_reduce add %arg0 cluster(size = 0) : (vector<4xf32>) -> vector<4xf32>
   return
 }
 
 // -----
 
 func.func @subgroup_reduce_npot_cluster_size(%arg0 : vector<4xf32>) {
   // expected-error@+1 {{cluster size 3 is not a power of two}}
-  %res = gpu.subgroup_reduce add %arg0 cluster_size(3) : (vector<4xf32>) -> vector<4xf32>
+  %res = gpu.subgroup_reduce add %arg0 cluster(size = 3) : (vector<4xf32>) -> vector<4xf32>
   return
 }
 
+// -----
+
+func.func @subgroup_reduce_zero_cluster_stride(%arg0 : vector<4xf32>) {
+  // expected-error@+1 {{cluster stride 0 is not a power of two}}
+  %res = gpu.subgroup_reduce add %arg0 cluster(size = 4, stride = 0) : (vector<4xf32>) -> vector<4xf32>
+  return
+}
+
+// -----
+
+func.func @subgroup_reduce_cluster_stride_without_size(%arg0 : vector<4xf32>) {
+  // expected-error@+1 {{cluster stride can only be specified if cluster size is specified}}
+  %res = gpu.subgroup_reduce add %arg0 { cluster_stride = 2 : i32 } : (vector<4xf32>) -> vector<4xf32>
+  return
+}
+
+
 // -----
 
 func.func @subgroup_reduce_bad_type(%arg0 : vector<2x2xf32>) {
diff --git a/mlir/test/Dialect/GPU/subgroup-reduce-lowering.mlir b/mlir/test/Dialect/GPU/subgroup-reduce-lowering.mlir

Original file line number	Diff line number	Diff line change
`@@ -621,14 +621,25 @@ LogicalResult gpu::SubgroupReduceOp::verify() {`
`621`	`621`	`<< getType();`
`622`	`622`	`}`
`623`	`623`
`624`		`- if (auto clusterSize = getClusterSize()) {`
	`624`	`+ auto clusterSize = getClusterSize();`
	`625`	`+ if (clusterSize) {`
`625`	`626`	`uint32_t size = *clusterSize;`
`626`	`627`	`if (!llvm::isPowerOf2_32(size)) {`
`627`	`628`	`return emitOpError() << "cluster size " << size`
`628`	`629`	`<< " is not a power of two";`
`629`	`630`	`}`
`630`	`631`	`}`
`631`	`632`
	`633`	`+ uint32_t stride = getClusterStride();`
	`634`	`+ if (stride != 1 && !clusterSize) {`
	`635`	`+ return emitOpError() << "cluster stride can only be specified if cluster "`
	`636`	`+ "size is specified";`
	`637`	`+ }`
	`638`	`+ if (!llvm::isPowerOf2_32(stride)) {`
	`639`	`+ return emitOpError() << "cluster stride " << stride`
	`640`	`+ << " is not a power of two";`
	`641`	`+ }`
	`642`	`+`
`632`	`643`	`return success();`
`633`	`644`	`}`
`634`	`645`
Original file line number	Diff line number	Diff line change
`@@ -255,7 +255,7 @@ func.func @subgroup_reduce_cluster_size_1() {`
`255`	`255`	`gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %0#0, %arg7 = %0#1, %arg8 = %0#2)`
`256`	`256`	`threads(%arg3, %arg4, %arg5) in (%arg9 = %0#3, %arg10 = %0#4, %arg11 = %0#5) {`
`257`	`257`	`%1 = "test.test2"() : () -> i32`
`258`		`- %2 = gpu.subgroup_reduce add %1 cluster_size(1) : (i32) -> (i32)`
	`258`	`+ %2 = gpu.subgroup_reduce add %1 cluster(size=1) : (i32) -> (i32)`
`259`	`259`	`"test.test3"(%2) : (i32) -> ()`
`260`	`260`	`gpu.terminator`
`261`	`261`	`}`