[Mlir-commits] [mlir] ff46298 - [mlir][SPIR-V] Support bitwise and/or/xor in gpu reduction lowering (#202351)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jul 28 04:25:05 PDT 2026
Author: Arseniy Obolenskiy
Date: 2026-07-28T13:25:00+02:00
New Revision: ff46298372dc64d8fb56f1397c2914772b815605
URL: https://github.com/llvm/llvm-project/commit/ff46298372dc64d8fb56f1397c2914772b815605
DIFF: https://github.com/llvm/llvm-project/commit/ff46298372dc64d8fb56f1397c2914772b815605.diff
LOG: [mlir][SPIR-V] Support bitwise and/or/xor in gpu reduction lowering (#202351)
Added:
Modified:
mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
mlir/test/Conversion/GPUToSPIRV/reductions.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
index 87a032bfd766b..48623e2fac47c 100644
--- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
+++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
@@ -718,6 +718,33 @@ static Value createGroupReduceOpImpl(OpBuilder &builder, Location loc,
.getResult();
}
+template <typename NonUniformOp>
+static Value createGroupNonUniformBitwiseReduceOpImpl(
+ OpBuilder &builder, Location loc, Value arg, bool isGroup, bool isUniform,
+ std::optional<uint32_t> clusterSize) {
+ spirv::Scope scope =
+ isGroup ? spirv::Scope::Workgroup : spirv::Scope::Subgroup;
+ if (isUniform || scope != spirv::Scope::Subgroup)
+ return Value();
+
+ Type type = arg.getType();
+ auto scopeAttr = mlir::spirv::ScopeAttr::get(builder.getContext(), scope);
+ auto groupOp = spirv::GroupOperationAttr::get(
+ builder.getContext(), clusterSize.has_value()
+ ? spirv::GroupOperation::ClusteredReduce
+ : spirv::GroupOperation::Reduce);
+
+ Value clusterSizeValue;
+ if (clusterSize.has_value())
+ clusterSizeValue = spirv::ConstantOp::create(
+ builder, loc, builder.getI32Type(),
+ builder.getIntegerAttr(builder.getI32Type(), *clusterSize));
+
+ return NonUniformOp::create(builder, loc, type, scopeAttr, groupOp, arg,
+ clusterSizeValue)
+ .getResult();
+}
+
static std::optional<Value>
createGroupReduceOp(OpBuilder &builder, Location loc, Value arg,
gpu::AllReduceOperation opType, bool isGroup,
@@ -784,11 +811,22 @@ createGroupReduceOp(OpBuilder &builder, Location loc, Value arg,
spirv::GroupNonUniformFMinOp>},
{ReduceType::MAXIMUMF, ElemType::Float,
&createGroupReduceOpImpl<spirv::GroupFMaxOp,
- spirv::GroupNonUniformFMaxOp>}};
+ spirv::GroupNonUniformFMaxOp>},
+ {ReduceType::AND, ElemType::Integer,
+ &createGroupNonUniformBitwiseReduceOpImpl<
+ spirv::GroupNonUniformBitwiseAndOp>},
+ {ReduceType::OR, ElemType::Integer,
+ &createGroupNonUniformBitwiseReduceOpImpl<
+ spirv::GroupNonUniformBitwiseOrOp>},
+ {ReduceType::XOR, ElemType::Integer,
+ &createGroupNonUniformBitwiseReduceOpImpl<
+ spirv::GroupNonUniformBitwiseXorOp>}};
for (const OpHandler &handler : handlers)
if (handler.kind == opType && elementType == handler.elemType)
- return handler.func(builder, loc, arg, isGroup, isUniform, clusterSize);
+ if (Value result =
+ handler.func(builder, loc, arg, isGroup, isUniform, clusterSize))
+ return result;
return std::nullopt;
}
diff --git a/mlir/test/Conversion/GPUToSPIRV/reductions.mlir b/mlir/test/Conversion/GPUToSPIRV/reductions.mlir
index e7e0fa296c98a..99c9bcab78187 100644
--- a/mlir/test/Conversion/GPUToSPIRV/reductions.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/reductions.mlir
@@ -675,6 +675,148 @@ gpu.module @kernels {
// -----
+// Bitwise reductions have no SPIR-V group op with Workgroup scope, so
+// gpu.all_reduce fails to legalize.
+
+module attributes {
+ gpu.container_module,
+ spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Kernel, Addresses, Groups, GroupNonUniformArithmetic, GroupUniformArithmeticKHR], []>, #spirv.resource_limits<>>
+} {
+
+gpu.module @kernels {
+ gpu.func @and(%arg : i32) kernel
+ attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
+ // expected-error @+1 {{failed to legalize operation 'gpu.all_reduce'}}
+ %r0 = gpu.all_reduce and %arg {} : (i32) -> (i32)
+ gpu.return
+ }
+}
+
+}
+
+// -----
+
+// Bitwise reductions have no SPIR-V group op with Workgroup scope, so
+// gpu.all_reduce fails to legalize.
+
+module attributes {
+ gpu.container_module,
+ spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Kernel, Addresses, Groups, GroupNonUniformArithmetic, GroupUniformArithmeticKHR], []>, #spirv.resource_limits<>>
+} {
+
+gpu.module @kernels {
+ gpu.func @or(%arg : i32) kernel
+ attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
+ // expected-error @+1 {{failed to legalize operation 'gpu.all_reduce'}}
+ %r0 = gpu.all_reduce or %arg {} : (i32) -> (i32)
+ gpu.return
+ }
+}
+
+}
+
+// -----
+
+// Bitwise reductions have no SPIR-V group op with Workgroup scope, so
+// gpu.all_reduce fails to legalize.
+
+module attributes {
+ gpu.container_module,
+ spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Kernel, Addresses, Groups, GroupNonUniformArithmetic, GroupUniformArithmeticKHR], []>, #spirv.resource_limits<>>
+} {
+
+gpu.module @kernels {
+ gpu.func @xor(%arg : i32) kernel
+ attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
+ // expected-error @+1 {{failed to legalize operation 'gpu.all_reduce'}}
+ %r0 = gpu.all_reduce xor %arg {} : (i32) -> (i32)
+ gpu.return
+ }
+}
+
+}
+
+// -----
+
+module attributes {
+ gpu.container_module,
+ spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Kernel, Addresses, Groups, GroupNonUniformArithmetic, GroupUniformArithmeticKHR], []>, #spirv.resource_limits<>>
+} {
+
+gpu.module @kernels {
+ // CHECK-LABEL: spirv.func @and
+ // CHECK-SAME: (%[[ARG:.*]]: i32)
+ gpu.func @and(%arg : i32) kernel
+ attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
+ // CHECK: %{{.*}} = spirv.GroupNonUniformBitwiseAnd <Subgroup> <Reduce> %[[ARG]] : i32 -> i32
+ %r0 = gpu.subgroup_reduce and %arg : (i32) -> (i32)
+ gpu.return
+ }
+}
+
+}
+
+// -----
+
+module attributes {
+ gpu.container_module,
+ spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Kernel, Addresses, Groups, GroupNonUniformArithmetic, GroupUniformArithmeticKHR], []>, #spirv.resource_limits<>>
+} {
+
+gpu.module @kernels {
+ // CHECK-LABEL: spirv.func @or
+ // CHECK-SAME: (%[[ARG:.*]]: i32)
+ gpu.func @or(%arg : i32) kernel
+ attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
+ // CHECK: %{{.*}} = spirv.GroupNonUniformBitwiseOr <Subgroup> <Reduce> %[[ARG]] : i32 -> i32
+ %r0 = gpu.subgroup_reduce or %arg : (i32) -> (i32)
+ gpu.return
+ }
+}
+
+}
+
+// -----
+
+module attributes {
+ gpu.container_module,
+ spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Kernel, Addresses, Groups, GroupNonUniformArithmetic, GroupUniformArithmeticKHR], []>, #spirv.resource_limits<>>
+} {
+
+gpu.module @kernels {
+ // CHECK-LABEL: spirv.func @xor
+ // CHECK-SAME: (%[[ARG:.*]]: i32)
+ gpu.func @xor(%arg : i32) kernel
+ attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
+ // CHECK: %{{.*}} = spirv.GroupNonUniformBitwiseXor <Subgroup> <Reduce> %[[ARG]] : i32 -> i32
+ %r0 = gpu.subgroup_reduce xor %arg : (i32) -> (i32)
+ gpu.return
+ }
+}
+
+}
+
+// -----
+
+// Bitwise reductions have no uniform SPIR-V group op, so a uniform request
+// fails to legalize.
+
+module attributes {
+ gpu.container_module,
+ spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Kernel, Addresses, Groups, GroupNonUniformArithmetic, GroupUniformArithmeticKHR], []>, #spirv.resource_limits<>>
+} {
+gpu.module @kernels {
+ gpu.func @and(%arg : i32) kernel
+ attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
+ // expected-error @+1 {{failed to legalize operation 'gpu.subgroup_reduce'}}
+ %r0 = gpu.subgroup_reduce and %arg uniform : (i32) -> (i32)
+ gpu.return
+ }
+}
+}
+
+// -----
+
// TODO: Handle boolean reductions.
module attributes {
More information about the Mlir-commits
mailing list