[Mlir-commits] [mlir] [mlir][SPIR-V] Accept spec constants for GroupNonUniform cluster size (PR #214408)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-spirv

Author: Arseniy Obolenskiy (aobolensk)

<details>
<summary>Changes</summary>

The SPIR-V spec requires ClusterSize to come from a constant instruction, which includes OpSpecConstant

The verifiers only accepted spirv.Constant, rejecting valid modules that size subgroup clusters with a tunable spec constant

---
Full diff: https://github.com/llvm/llvm-project/pull/214408.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/SPIRV/IR/GroupOps.cpp (+12-7) 
- (modified) mlir/test/Dialect/SPIRV/IR/non-uniform-ops.mlir (+26) 


``````````diff
diff --git a/mlir/lib/Dialect/SPIRV/IR/GroupOps.cpp b/mlir/lib/Dialect/SPIRV/IR/GroupOps.cpp
index fe6f00e9e5bca..6a206f03a62b7 100644
--- a/mlir/lib/Dialect/SPIRV/IR/GroupOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/GroupOps.cpp
@@ -33,14 +33,16 @@ static LogicalResult verifyGroupNonUniformArithmeticOp(Operation *groupOp) {
                                 "'ClusteredReduce' group operation");
   if (groupOp->getNumOperands() > 1) {
     Operation *sizeOp = groupOp->getOperand(1).getDefiningOp();
-    int32_t clusterSize = 0;
 
-    // TODO: support specialization constant here.
-    if (failed(extractValueFromConstOp(sizeOp, clusterSize)))
+    // SPIR-V spec: "ClusterSize must come from a constant instruction",
+    // which includes both a normal constant and a specialization constant.
+    if (!isa_and_present<spirv::ConstantOp, spirv::ReferenceOfOp>(sizeOp))
       return groupOp->emitOpError(
           "cluster size operand must come from a constant op");
 
-    if (!llvm::isPowerOf2_32(clusterSize))
+    int32_t clusterSize = 0;
+    if (succeeded(extractValueFromConstOp(sizeOp, clusterSize)) &&
+        !llvm::isPowerOf2_32(clusterSize))
       return groupOp->emitOpError(
           "cluster size operand must be a power of two");
   }
@@ -242,12 +244,15 @@ LogicalResult GroupNonUniformLogicalXorOp::verify() {
 LogicalResult GroupNonUniformRotateKHROp::verify() {
   if (Value clusterSizeVal = getClusterSize()) {
     mlir::Operation *defOp = clusterSizeVal.getDefiningOp();
-    int32_t clusterSize = 0;
 
-    if (failed(extractValueFromConstOp(defOp, clusterSize)))
+    // SPIR-V spec: "ClusterSize must come from a constant instruction",
+    // which includes both a normal constant and a specialization constant.
+    if (!isa_and_present<spirv::ConstantOp, spirv::ReferenceOfOp>(defOp))
       return emitOpError("cluster size operand must come from a constant op");
 
-    if (!llvm::isPowerOf2_32(clusterSize))
+    int32_t clusterSize = 0;
+    if (succeeded(extractValueFromConstOp(defOp, clusterSize)) &&
+        !llvm::isPowerOf2_32(clusterSize))
       return emitOpError("cluster size operand must be a power of two");
   }
 
diff --git a/mlir/test/Dialect/SPIRV/IR/non-uniform-ops.mlir b/mlir/test/Dialect/SPIRV/IR/non-uniform-ops.mlir
index 9c927331a08b4..0ce076a27ed4b 100644
--- a/mlir/test/Dialect/SPIRV/IR/non-uniform-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/non-uniform-ops.mlir
@@ -312,6 +312,19 @@ func.func @group_non_uniform_iadd_clustered_reduce(%val: vector<2xi32>) -> vecto
 
 // -----
 
+spirv.module Logical GLSL450 {
+  spirv.SpecConstant @cluster_size = 4 : i32
+  // CHECK-LABEL: @group_non_uniform_iadd_clustered_reduce_spec_const
+  spirv.func @group_non_uniform_iadd_clustered_reduce_spec_const(%val: vector<2xi32>) -> vector<2xi32> "None" {
+    %size = spirv.mlir.referenceof @cluster_size : i32
+    // CHECK: %{{.+}} = spirv.GroupNonUniformIAdd <Subgroup> <ClusteredReduce> %{{.+}} cluster_size(%{{.+}}) : vector<2xi32>, i32 -> vector<2xi32>
+    %0 = spirv.GroupNonUniformIAdd <Subgroup> <ClusteredReduce> %val cluster_size(%size) : vector<2xi32>, i32 -> vector<2xi32>
+    spirv.ReturnValue %0 : vector<2xi32>
+  }
+}
+
+// -----
+
 func.func @group_non_uniform_iadd_reduce(%val: i32) -> i32 {
   // expected-error @+1 {{execution_scope must be 'Subgroup'}}
   %0 = spirv.GroupNonUniformIAdd <Device> <Reduce> %val : i32 -> i32
@@ -775,6 +788,19 @@ func.func @group_non_uniform_rotate_khr(%val: f32, %delta: i32) -> f32 {
 
 // -----
 
+spirv.module Logical GLSL450 {
+  spirv.SpecConstant @cluster_size = 4 : i32
+  // CHECK-LABEL: @group_non_uniform_rotate_khr_spec_const
+  spirv.func @group_non_uniform_rotate_khr_spec_const(%val: f32, %delta: i32) -> f32 "None" {
+    %size = spirv.mlir.referenceof @cluster_size : i32
+    // CHECK: %{{.+}} = spirv.GroupNonUniformRotateKHR <Workgroup> %{{.+}} : f32, i32, i32 -> f32
+    %0 = spirv.GroupNonUniformRotateKHR <Workgroup> %val, %delta, cluster_size(%size) : f32, i32, i32 -> f32
+    spirv.ReturnValue %0 : f32
+  }
+}
+
+// -----
+
 func.func @group_non_uniform_rotate_khr(%val: f32, %delta: i32) -> f32 {
   %four = spirv.Constant 4 : i32
   // expected-error @+1 {{execution_scope must be 'Workgroup' or 'Subgroup'}}

``````````

</details>


https://github.com/llvm/llvm-project/pull/214408