[Mlir-commits] [mlir] [mlir][SPIR-V] Support bitwise and/or/xor in gpu reduction lowering (PR #202351)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: Arseniy Obolenskiy (aobolensk)

<details>
<summary>Changes</summary>



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


2 Files Affected:

- (modified) mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp (+40-2) 
- (modified) mlir/test/Conversion/GPUToSPIRV/reductions.mlir (+67) 


``````````diff
diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
index 87a032bfd766b..144c08dec2cea 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) {
+  if (isUniform)
+    return Value();
+
+  Type type = arg.getType();
+  auto scope = mlir::spirv::ScopeAttr::get(builder.getContext(),
+                                           isGroup ? spirv::Scope::Workgroup
+                                                   : spirv::Scope::Subgroup);
+  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, scope, 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..fae4c8104cf8e 100644
--- a/mlir/test/Conversion/GPUToSPIRV/reductions.mlir
+++ b/mlir/test/Conversion/GPUToSPIRV/reductions.mlir
@@ -675,6 +675,73 @@ gpu.module @kernels {
 
 // -----
 
+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 @test
+  //  CHECK-SAME: (%[[ARG:.*]]: i32)
+  gpu.func @test(%arg : i32) kernel
+    attributes {spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [16, 1, 1]>} {
+    // CHECK: %{{.*}} = spirv.GroupNonUniformBitwiseAnd <Workgroup> <Reduce> %[[ARG]] : i32 -> i32
+    // CHECK: %{{.*}} = spirv.GroupNonUniformBitwiseOr <Workgroup> <Reduce> %[[ARG]] : i32 -> i32
+    // CHECK: %{{.*}} = spirv.GroupNonUniformBitwiseXor <Workgroup> <Reduce> %[[ARG]] : i32 -> i32
+    %r0 = gpu.all_reduce and %arg {} : (i32) -> (i32)
+    %r1 = gpu.all_reduce or %arg {} : (i32) -> (i32)
+    %r2 = 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 @test
+  //  CHECK-SAME: (%[[ARG:.*]]: i32)
+  gpu.func @test(%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
+    // CHECK: %{{.*}} = spirv.GroupNonUniformBitwiseOr <Subgroup> <Reduce> %[[ARG]] : i32 -> i32
+    // CHECK: %{{.*}} = spirv.GroupNonUniformBitwiseXor <Subgroup> <Reduce> %[[ARG]] : i32 -> i32
+    %r0 = gpu.subgroup_reduce and %arg : (i32) -> (i32)
+    %r1 = gpu.subgroup_reduce or %arg : (i32) -> (i32)
+    %r2 = 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 {

``````````

</details>


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