[Mlir-commits] [mlir] [mlir][SPIR-V] Accept spec constants for GroupNonUniform cluster size (PR #214408)
Arseniy Obolenskiy
llvmlistbot at llvm.org
Mon Sep 21 05:11:26 PDT 2026
https://github.com/aobolensk updated https://github.com/llvm/llvm-project/pull/214408
>From 5c9d4bcc65ee3316a1df7b8356f4493d08f119ec Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Thu, 6 Aug 2026 08:33:19 +0200
Subject: [PATCH 1/2] [mlir][SPIR-V] Accept spec constants for GroupNonUniform
cluster size
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
---
mlir/lib/Dialect/SPIRV/IR/GroupOps.cpp | 19 +++++++++-----
.../Dialect/SPIRV/IR/non-uniform-ops.mlir | 26 +++++++++++++++++++
2 files changed, 38 insertions(+), 7 deletions(-)
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'}}
>From b45367d1573ffe1004956e2879a2af0918630dcc Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Mon, 21 Sep 2026 14:11:11 +0200
Subject: [PATCH 2/2] remove checks
---
mlir/lib/Dialect/SPIRV/IR/GroupOps.cpp | 33 +-----------------
.../Dialect/SPIRV/IR/non-uniform-ops.mlir | 34 -------------------
2 files changed, 1 insertion(+), 66 deletions(-)
diff --git a/mlir/lib/Dialect/SPIRV/IR/GroupOps.cpp b/mlir/lib/Dialect/SPIRV/IR/GroupOps.cpp
index 78ab374b9713f..22d02622b0a19 100644
--- a/mlir/lib/Dialect/SPIRV/IR/GroupOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/GroupOps.cpp
@@ -27,21 +27,6 @@ static LogicalResult verifyGroupNonUniformArithmeticOp(Operation *groupOp) {
groupOp->getNumOperands() == 1)
return groupOp->emitOpError("cluster size operand must be provided for "
"'ClusteredReduce' group operation");
- if (groupOp->getNumOperands() > 1) {
- Operation *sizeOp = groupOp->getOperand(1).getDefiningOp();
-
- // 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");
-
- 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");
- }
return success();
}
@@ -237,22 +222,6 @@ LogicalResult GroupNonUniformLogicalXorOp::verify() {
// spirv.GroupNonUniformRotateKHR
//===----------------------------------------------------------------------===//
-LogicalResult GroupNonUniformRotateKHROp::verify() {
- if (Value clusterSizeVal = getClusterSize()) {
- mlir::Operation *defOp = clusterSizeVal.getDefiningOp();
-
- // 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");
-
- int32_t clusterSize = 0;
- if (succeeded(extractValueFromConstOp(defOp, clusterSize)) &&
- !llvm::isPowerOf2_32(clusterSize))
- return emitOpError("cluster size operand must be a power of two");
- }
-
- return success();
-}
+LogicalResult GroupNonUniformRotateKHROp::verify() { return success(); }
} // namespace mlir::spirv
diff --git a/mlir/test/Dialect/SPIRV/IR/non-uniform-ops.mlir b/mlir/test/Dialect/SPIRV/IR/non-uniform-ops.mlir
index 0ce076a27ed4b..ceabc0e40c021 100644
--- a/mlir/test/Dialect/SPIRV/IR/non-uniform-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/non-uniform-ops.mlir
@@ -341,23 +341,6 @@ func.func @group_non_uniform_iadd_clustered_reduce(%val: vector<2xi32>) -> vecto
// -----
-func.func @group_non_uniform_iadd_clustered_reduce(%val: vector<2xi32>, %size: i32) -> vector<2xi32> {
- // expected-error @+1 {{cluster size operand must come from a constant op}}
- %0 = spirv.GroupNonUniformIAdd <Subgroup> <ClusteredReduce> %val cluster_size(%size) : vector<2xi32>, i32 -> vector<2xi32>
- return %0: vector<2xi32>
-}
-
-// -----
-
-func.func @group_non_uniform_iadd_clustered_reduce(%val: vector<2xi32>) -> vector<2xi32> {
- %five = spirv.Constant 5 : i32
- // expected-error @+1 {{cluster size operand must be a power of two}}
- %0 = spirv.GroupNonUniformIAdd <Subgroup> <ClusteredReduce> %val cluster_size(%five) : vector<2xi32>, i32 -> vector<2xi32>
- return %0: vector<2xi32>
-}
-
-// -----
-
//===----------------------------------------------------------------------===//
// spirv.GroupNonUniformIMul
//===----------------------------------------------------------------------===//
@@ -828,23 +811,6 @@ func.func @group_non_uniform_rotate_khr(%val: f32, %delta: i32) -> f32 {
// -----
-func.func @group_non_uniform_rotate_khr(%val: f32, %delta: i32, %four: i32) -> f32 {
- // expected-error @+1 {{cluster size operand must come from a constant op}}
- %0 = spirv.GroupNonUniformRotateKHR <Subgroup> %val, %delta, cluster_size(%four) : f32, i32, i32 -> f32
- return %0: f32
-}
-
-// -----
-
-func.func @group_non_uniform_rotate_khr(%val: f32, %delta: i32) -> f32 {
- %five = spirv.Constant 5 : i32
- // expected-error @+1 {{cluster size operand must be a power of two}}
- %0 = spirv.GroupNonUniformRotateKHR <Subgroup> %val, %delta, cluster_size(%five) : f32, i32, i32 -> f32
- return %0: f32
-}
-
-// -----
-
//===----------------------------------------------------------------------===//
// spirv.GroupNonUniformAll
//===----------------------------------------------------------------------===//
More information about the Mlir-commits
mailing list