[Mlir-commits] [mlir] [mlir][spirv] Implement lowering `gpu.subgroup_reduce` with cluster size for SPIRV (PR #141402)

[Jakub Kuderski]

================
@@ -464,27 +464,39 @@ LogicalResult GPUShuffleConversion::matchAndRewrite(
 
 template <typename UniformOp, typename NonUniformOp>
 static Value createGroupReduceOpImpl(OpBuilder &builder, Location loc,
-                                     Value arg, bool isGroup, bool isUniform) {
+                                     Value arg, bool isGroup, bool isUniform,
+                                     std::optional<uint32_t> clusterSize) {
   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(),
-                                                spirv::GroupOperation::Reduce);
+  auto groupOp = spirv::GroupOperationAttr::get(
+      builder.getContext(), clusterSize.has_value()
+                                ? spirv::GroupOperation::ClusteredReduce
+                                : spirv::GroupOperation::Reduce);
   if (isUniform) {
     return builder.create<UniformOp>(loc, type, scope, groupOp, arg)
         .getResult();
   }
-  return builder.create<NonUniformOp>(loc, type, scope, groupOp, arg, Value{})
+
+  Value clusterSizeValue =
+      clusterSize.has_value()
+          ? builder.create<spirv::ConstantOp>(
+                loc, builder.getI32Type(),
+                builder.getIntegerAttr(builder.getI32Type(), *clusterSize))
+          : Value{};
----------------
kuhar wrote:

nit: prefer if statements over complicated ternaries like this one

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