[mlir][spirv] Add GroupNonUniform arithmetic operations.

[denis13]

Summary

Add GroupNonUniform arithmetic operations: FAdd, FMul, IMul.
Unify parser, printer, verifier for GroupNonUniform arithmetic
operations.

[merge_guards_bot]

Unit tests: pass. 62214 tests passed, 0 failed and 815 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[mravishankar]

The (de)serialization needs to be implemented. Will that come later?

mlir/include/mlir/Dialect/SPIRV/SPIRVNonUniformOps.td

334

The SPIR-V spec says that the cluster_size must come from a constant. Do we want to just have an integer-literal here? We need to create the constant op while serialization too...

[denis13]

mlir/include/mlir/Dialect/SPIRV/SPIRVNonUniformOps.td

334

@mravishankar thanks for review, can you please clarify a little bit, the cluster_size is a SSA value and not a integral-literal, also check that cluster_size comes from constant at L662, and test for serialization. Do I miss something? Thanks.

[mravishankar]

mlir/include/mlir/Dialect/SPIRV/SPIRVNonUniformOps.td

334

Sorry about the serialization comment. The test indeed works with serialization.

In SPIR-V spec, the cluster_size is an SSA value, but it also says it has to come as a result of a constant instruction. One way to implement this is to accept an integer-literal in the op itself (so that you dont have an extra constant instruction generated to support this). But now that I think about this, the value could also come from a specialization constant. So its probably better to leave this as is, with the added advantage the the serialization will be auto-generated. So ignore this comment as a whole.

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

631

One suggestion is to make this a templated method (templated on OpTy). Then you could use ".execution_scope()" and ".group_operation()" to get the attribute values

646

Same here. Maybe this could be templated on OpTy.

Please add comments (one-liner should be fine) as well to describe what this method is for.

[antiagainst]

mlir/include/mlir/Dialect/SPIRV/SPIRVNonUniformOps.td

334

Right. In SPIR-V constant instructions include both normal OpConstant* and specialization OpSpecConstant*, so normally we don't want to represent a constant operand onto the op itself as an attribute. We do it for cases like for scope and memory semantic <id>s because the validation rule explicitly said that "All <id> used for Scope and Memory Semantics must be of an OpConstant". Making scope and memory semantics as attributes indeed improves the IR readability and handling. However, the specific validation rule is for Shader capability though; we need to relax this if, in the future, we support Kernel capability cases, or we need to design our own IntegerAttr that can encode the spec constant ID. At that time, maybe we can rethink here too whether we want to use that special IntegerAttr. It's a broader scope change though: we need modelling EnumAttr for it and all the stuff associated. Would prefer to defer until necessary. :)

[denis13]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

631

@mravishankar @antiagainst I agree that templated methods look more readable, but on the other hand for each GroupNonUniform arithmetic operation will be created 3 methods with only one difference - function signature. They will implement the same logic and code. Also I was looking on the other parser/printer/verifier functions for operations from the same category, for example cast op, and they use Operation* as function parameter, what do you think? Thanks.

[denis13]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

631

I was wrong about 3 methods, since parser is already generic, so it will be 2 methods for each GroupNonUniform arithmetic operation.

[antiagainst]

Thanks Denis for adding these ops!

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

631

That's reasonable. This is a minor issue anyway; I'm fine with the current implementation.

mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td

 def SPV_OC_OpBranchConditional         : I32EnumAttrCase<"OpBranchConditional", 250>;
 def SPV_OC_OpReturn                    : I32EnumAttrCase<"OpReturn", 253>;
 def SPV_OC_OpReturnValue               : I32EnumAttrCase<"OpReturnValue", 254>;
 def SPV_OC_OpUnreachable               : I32EnumAttrCase<"OpUnreachable", 255>;
 def SPV_OC_OpModuleProcessed           : I32EnumAttrCase<"OpModuleProcessed", 330>;
 def SPV_OC_OpGroupNonUniformElect      : I32EnumAttrCase<"OpGroupNonUniformElect", 333>;
 def SPV_OC_OpGroupNonUniformBallot     : I32EnumAttrCase<"OpGroupNonUniformBallot", 339>;
 def SPV_OC_OpGroupNonUniformIAdd       : I32EnumAttrCase<"OpGroupNonUniformIAdd", 349>;
+def SPV_OC_OpGroupNonUniformFAdd       : I32EnumAttrCase<"OpGroupNonUniformFAdd", 350>;
+def SPV_OC_OpGroupNonUniformIMul       : I32EnumAttrCase<"OpGroupNonUniformIMul", 351>;
+def SPV_OC_OpGroupNonUniformFMul       : I32EnumAttrCase<"OpGroupNonUniformFMul", 352>;
 def SPV_OC_OpSubgroupBallotKHR         : I32EnumAttrCase<"OpSubgroupBallotKHR", 4421>;
 
 def SPV_OpcodeAttr :
     SPV_I32EnumAttr<"Opcode", "valid SPIR-V instructions", [
       SPV_OC_OpNop, SPV_OC_OpUndef, SPV_OC_OpSourceContinued, SPV_OC_OpSource,
       SPV_OC_OpSourceExtension, SPV_OC_OpName, SPV_OC_OpMemberName, SPV_OC_OpString,
       SPV_OC_OpExtension, SPV_OC_OpExtInstImport, SPV_OC_OpExtInst,
       SPV_OC_OpMemoryModel, SPV_OC_OpEntryPoint, SPV_OC_OpExecutionMode,
SPV_I32EnumAttr<"Opcode", "valid SPIR-V instructions", [
       SPV_OC_OpAtomicCompareExchangeWeak, SPV_OC_OpAtomicIIncrement,
       SPV_OC_OpAtomicIDecrement, SPV_OC_OpAtomicIAdd, SPV_OC_OpAtomicISub,
       SPV_OC_OpAtomicSMin, SPV_OC_OpAtomicUMin, SPV_OC_OpAtomicSMax,
       SPV_OC_OpAtomicUMax, SPV_OC_OpAtomicAnd, SPV_OC_OpAtomicOr, SPV_OC_OpAtomicXor,
       SPV_OC_OpPhi, SPV_OC_OpLoopMerge, SPV_OC_OpSelectionMerge, SPV_OC_OpLabel,
       SPV_OC_OpBranch, SPV_OC_OpBranchConditional, SPV_OC_OpReturn,
       SPV_OC_OpReturnValue, SPV_OC_OpUnreachable, SPV_OC_OpModuleProcessed,
       SPV_OC_OpGroupNonUniformElect, SPV_OC_OpGroupNonUniformBallot,
-      SPV_OC_OpGroupNonUniformIAdd, SPV_OC_OpSubgroupBallotKHR
+      SPV_OC_OpGroupNonUniformIAdd, SPV_OC_OpGroupNonUniformFAdd,
+      SPV_OC_OpGroupNonUniformIMul, SPV_OC_OpGroupNonUniformFMul,
+      SPV_OC_OpSubgroupBallotKHR
     ]>;
 
 // End opcode section. Generated from SPIR-V spec; DO NOT MODIFY!
 
 //===----------------------------------------------------------------------===//
 // SPIR-V op definitions
 //===----------------------------------------------------------------------===//
 

mlir/include/mlir/Dialect/SPIRV/SPIRVNonUniformOps.td

 //===-- SPIRVNonUniformOps.td - MLIR SPIR-V NonUniform Ops -*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains non-uniform ops for the SPIR-V dialect. It corresponds to
 // "3.32.24. Non-Uniform Instructions" of the SPIR-V specification.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SPIRV_NON_UNIFORM_OPS
 #define SPIRV_NON_UNIFORM_OPS
 
+class SPV_GroupNonUniformArithmeticOp<string mnemonic, Type type,
+      list<OpTrait> traits = []> : SPV_Op<mnemonic, traits> {
+
+  let availability = [
+    MinVersion<SPV_V_1_3>,
+    MaxVersion<SPV_V_1_5>,
+    Extension<[]>,
+    Capability<[SPV_C_GroupNonUniformArithmetic,
+                SPV_C_GroupNonUniformClustered,
+                SPV_C_GroupNonUniformPartitionedNV]>
+  ];
+
+  let arguments = (ins
+    SPV_ScopeAttr:$execution_scope,
+    SPV_GroupOperationAttr:$group_operation,
+    SPV_ScalarOrVectorOf<type>:$value,
+    SPV_Optional<SPV_Integer>:$cluster_size
+  );
+
+  let results = (outs
+    SPV_ScalarOrVectorOf<type>:$result
+  );
+
+  let parser = [{ return parseGroupNonUniformArithmeticOp(parser, result); }];
+  let printer = [{ printGroupNonUniformArithmeticOp(getOperation(), p); }];
+  let verifier = [{ return ::verifyGroupNonUniformArithmeticOp(getOperation()); }];
+
+}
+
 // -----
 
 def SPV_GroupNonUniformBallotOp : SPV_Op<"GroupNonUniformBallot", []> {
   let summary = [{
     Returns a bitfield value combining the Predicate value from all
     invocations in the group that execute the same dynamic instance of this
     instruction. The bit is set to one if the corresponding invocation is
     active and the Predicate for that invocation evaluated to true;
def SPV_GroupNonUniformElectOp : SPV_Op<"GroupNonUniformElect", []> {
 
   let builders = [
     OpBuilder<[{Builder *builder, OperationState &state, spirv::Scope}]>
   ];
 }
 
 // -----
 
-def SPV_GroupNonUniformIAddOp : SPV_Op<"GroupNonUniformIAdd", []> {
+def SPV_GroupNonUniformFAddOp :
+    SPV_GroupNonUniformArithmeticOp<"GroupNonUniformFAdd", SPV_Float, []> {
+  let summary = [{
+    A floating point add group operation of all Value operands contributed
+    by active invocations in the group.
+  }];
+
+  let description = [{
+    Result Type  must be a scalar or vector of floating-point type.
+
+    Execution must be Workgroup or Subgroup Scope.
+
+    The identity I for Operation is 0. If Operation is ClusteredReduce,
+    ClusterSize must be specified.
+
+     The type of Value must be the same as Result Type.  The method used to
+    perform the group operation on the contributed Value(s) from active
+    invocations is implementation defined.
+
+    ClusterSize is the size of cluster to use. ClusterSize must be a scalar
+    of integer type, whose Signedness operand is 0. ClusterSize must come
+    from a constant instruction. ClusterSize must be at least 1, and must be
+    a power of 2. If ClusterSize is greater than the declared SubGroupSize,
+    executing this instruction results in undefined behavior.
+
+    ### Custom assembly form
+
+    ```
+    scope ::= `"Workgroup"` | `"Subgroup"`
+    operation ::= `"Reduce"` | `"InclusiveScan"` | `"ExclusiveScan"` | ...
+    float-scalar-vector-type ::= float-type |
+                                 `vector<` integer-literal `x` float-type `>`
+    non-uniform-fadd-op ::= ssa-id `=` `spv.GroupNonUniformFAdd` scope operation
+                            ssa-use ( `cluster_size` `(` ssa_use `)` )?
+                            `:` float-scalar-vector-type
+    ```
+
+    For example:
+
+    ```
+    %four = spv.constant 4 : i32
+    %scalar = ... : f32
+    %vector = ... : vector<4xf32>
+    %0 = spv.GroupNonUniformFAdd "Workgroup" "Reduce" %scalar : f32
+    %1 = spv.GroupNonUniformFAdd "Subgroup" "ClusteredReduce" %vector cluster_size(%four) : vector<4xf32>
+    ```
+  }];
+}
+
+// -----
+
+def SPV_GroupNonUniformFMulOp :
+    SPV_GroupNonUniformArithmeticOp<"GroupNonUniformFMul", SPV_Float, []> {
+  let summary = [{
+    A floating point multiply group operation of all Value operands
+    contributed by active invocations in the group.
+  }];
+
+  let description = [{
+    Result Type  must be a scalar or vector of floating-point type.
+
+    Execution must be Workgroup or Subgroup Scope.
+
+    The identity I for Operation is 1. If Operation is ClusteredReduce,
+    ClusterSize must be specified.
+
+     The type of Value must be the same as Result Type.  The method used to
+    perform the group operation on the contributed Value(s) from active
+    invocations is implementation defined.
+
+    ClusterSize is the size of cluster to use. ClusterSize must be a scalar
+    of integer type, whose Signedness operand is 0. ClusterSize must come
+    from a constant instruction. ClusterSize must be at least 1, and must be
+    a power of 2. If ClusterSize is greater than the declared SubGroupSize,
+    executing this instruction results in undefined behavior.
+
+    ### Custom assembly form
+
+    ```
+    scope ::= `"Workgroup"` | `"Subgroup"`
+    operation ::= `"Reduce"` | `"InclusiveScan"` | `"ExclusiveScan"` | ...
+    float-scalar-vector-type ::= float-type |
+                                 `vector<` integer-literal `x` float-type `>`
+    non-uniform-fmul-op ::= ssa-id `=` `spv.GroupNonUniformFMul` scope operation
+                            ssa-use ( `cluster_size` `(` ssa_use `)` )?
+                            `:` float-scalar-vector-type
+    ```
+
+    For example:
+
+    ```
+    %four = spv.constant 4 : i32
+    %scalar = ... : f32
+    %vector = ... : vector<4xf32>
+    %0 = spv.GroupNonUniformFMul "Workgroup" "Reduce" %scalar : f32
+    %1 = spv.GroupNonUniformFMul "Subgroup" "ClusteredReduce" %vector cluster_size(%four) : vector<4xf32>
+    ```
+  }];
+}
+
+// -----
+
+def SPV_GroupNonUniformIAddOp :
+    SPV_GroupNonUniformArithmeticOp<"GroupNonUniformIAdd", SPV_Integer, []> {
   let summary = [{
     An integer add group operation of all Value operands contributed active
     by invocations in the group.
   }];
 
   let description = [{
     Result Type  must be a scalar or vector of integer type.
 
let description = [{
     ```
     %four = spv.constant 4 : i32
     %scalar = ... : i32
     %vector = ... : vector<4xi32>
     %0 = spv.GroupNonUniformIAdd "Workgroup" "Reduce" %scalar : i32
     %1 = spv.GroupNonUniformIAdd "Subgroup" "ClusteredReduce" %vector cluster_size(%four) : vector<4xi32>
     ```
   }];
+}
 
-  let availability = [
-    MinVersion<SPV_V_1_3>,
-    MaxVersion<SPV_V_1_5>,
-    Extension<[]>,
-    Capability<[SPV_C_GroupNonUniformArithmetic, SPV_C_GroupNonUniformClustered, SPV_C_GroupNonUniformPartitionedNV]>
-  ];
+// -----
 
-  let arguments = (ins
-    SPV_ScopeAttr:$execution_scope,
-    SPV_GroupOperationAttr:$group_operation,
-    SPV_ScalarOrVectorOf<SPV_Integer>:$value,
-    SPV_Optional<SPV_Integer>:$cluster_size
-  );
+def SPV_GroupNonUniformIMulOp :
+    SPV_GroupNonUniformArithmeticOp<"GroupNonUniformIMul", SPV_Integer, []> {
+  let summary = [{
+    An integer multiply group operation of all Value operands contributed by
+    active invocations in the group.
+  }];
 
-  let results = (outs
-    SPV_ScalarOrVectorOf<SPV_Integer>:$result
-  );
+  let description = [{
+    Result Type  must be a scalar or vector of integer type.
+
+    Execution must be Workgroup or Subgroup Scope.
+
+    The identity I for Operation is 1. If Operation is ClusteredReduce,
+    ClusterSize must be specified.
+
+     The type of Value must be the same as Result Type.
+
+    ClusterSize is the size of cluster to use. ClusterSize must be a scalar
+    of integer type, whose Signedness operand is 0. ClusterSize must come
+    from a constant instruction. ClusterSize must be at least 1, and must be
+    a power of 2. If ClusterSize is greater than the declared SubGroupSize,
+    executing this instruction results in undefined behavior.
+
+    ### Custom assembly form
+
+    ```
+    scope ::= `"Workgroup"` | `"Subgroup"`
+    operation ::= `"Reduce"` | `"InclusiveScan"` | `"ExclusiveScan"` | ...
+    integer-scalar-vector-type ::= integer-type |
+                                 `vector<` integer-literal `x` integer-type `>`
+    non-uniform-imul-op ::= ssa-id `=` `spv.GroupNonUniformIMul` scope operation
+                            ssa-use ( `cluster_size` `(` ssa_use `)` )?
+                            `:` integer-scalar-vector-type
+    ```
+
+    For example:
+
+    ```
+    %four = spv.constant 4 : i32
+    %scalar = ... : i32
+    %vector = ... : vector<4xi32>
+    %0 = spv.GroupNonUniformIMul "Workgroup" "Reduce" %scalar : i32
+    %1 = spv.GroupNonUniformIMul "Subgroup" "ClusteredReduce" %vector cluster_size(%four) : vector<4xi32>
+    ```
+  }];
 }
 
 // -----
 
 #endif // SPIRV_NON_UNIFORM_OPS

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

if (op->getNumOperands() > 1) {
     if (valueType != elementType)
       return op->emitOpError("expected value to have the same type as the "
                              "pointer operand's pointee type ")
              << elementType << ", but found " << valueType;
   }
   return success();
 }
 
+static ParseResult parseGroupNonUniformArithmeticOp(OpAsmParser &parser,
+                                                    OperationState &state) {
+  spirv::Scope executionScope;
+  spirv::GroupOperation groupOperation;
+  OpAsmParser::OperandType valueInfo;
+  if (parseEnumAttribute(executionScope, parser, state,
+                         kExecutionScopeAttrName) ||
+      parseEnumAttribute(groupOperation, parser, state,
+                         kGroupOperationAttrName) ||
+      parser.parseOperand(valueInfo))
+    return failure();
+
+  Optional<OpAsmParser::OperandType> clusterSizeInfo;
+  if (succeeded(parser.parseOptionalKeyword(kClusterSize))) {
+    clusterSizeInfo = OpAsmParser::OperandType();
+    if (parser.parseLParen() || parser.parseOperand(*clusterSizeInfo) ||
+        parser.parseRParen())
+      return failure();
+  }
+
+  Type resultType;
+  if (parser.parseColonType(resultType))
+    return failure();
+
+  if (parser.resolveOperand(valueInfo, resultType, state.operands))
+    return failure();
+
+  if (clusterSizeInfo.hasValue()) {
+    Type i32Type = parser.getBuilder().getIntegerType(32);
+    if (parser.resolveOperand(*clusterSizeInfo, i32Type, state.operands))
+      return failure();
+  }
+
+  return parser.addTypeToList(resultType, state.types);
+}
+
+static void printGroupNonUniformArithmeticOp(Operation *groupOp,
+                                             OpAsmPrinter &printer) {
+  printer << groupOp->getName() << " \""
+          << stringifyScope(static_cast<spirv::Scope>(
+                 groupOp->getAttrOfType<IntegerAttr>(kExecutionScopeAttrName)
+                     .getInt()))
+          << "\" \""
+          << stringifyGroupOperation(static_cast<spirv::GroupOperation>(
+                 groupOp->getAttrOfType<IntegerAttr>(kGroupOperationAttrName)
+                     .getInt()))
+          << "\" " << groupOp->getOperand(0);
+
+  if (groupOp->getNumOperands() > 1)
+    printer << " " << kClusterSize << '(' << groupOp->getOperand(1) << ')';
+  printer << " : " << groupOp->getResult(0).getType();
+}
+
+static LogicalResult verifyGroupNonUniformArithmeticOp(Operation *groupOp) {
+  spirv::Scope scope = static_cast<spirv::Scope>(
+      groupOp->getAttrOfType<IntegerAttr>(kExecutionScopeAttrName).getInt());
+  if (scope != spirv::Scope::Workgroup && scope != spirv::Scope::Subgroup)
+    return groupOp->emitOpError(
+        "execution scope must be 'Workgroup' or 'Subgroup'");
+
+  spirv::GroupOperation operation = static_cast<spirv::GroupOperation>(
+      groupOp->getAttrOfType<IntegerAttr>(kGroupOperationAttrName).getInt());
+  if (operation == spirv::GroupOperation::ClusteredReduce &&
+      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();
+    int32_t clusterSize = 0;
+
+    // TODO(antiagainst): support specialization constant here.
+    if (failed(extractValueFromConstOp(sizeOp, clusterSize)))
+      return groupOp->emitOpError(
+          "cluster size operand must come from a constant op");
+
+    if (!llvm::isPowerOf2_32(clusterSize))
+      return groupOp->emitOpError(
+          "cluster size operand must be a power of two");
+  }
+  return success();
+}
+
 // Parses an op that has no inputs and no outputs.
 static ParseResult parseNoIOOp(OpAsmParser &parser, OperationState &state) {
   if (parser.parseOptionalAttrDict(state.attributes))
     return failure();
   return success();
 }
 
 // Prints an op that has no inputs and no outputs.
static LogicalResult verify(spirv::GroupNonUniformElectOp groupOp) {
   spirv::Scope scope = groupOp.execution_scope();
   if (scope != spirv::Scope::Workgroup && scope != spirv::Scope::Subgroup)
     return groupOp.emitOpError(
         "execution scope must be 'Workgroup' or 'Subgroup'");
 
   return success();
 }
 
-//===----------------------------------------------------------------------===//
-// spv.GroupNonUniformIAddOp
-//===----------------------------------------------------------------------===//
 
-static ParseResult parseGroupNonUniformIAddOp(OpAsmParser &parser,
-                                              OperationState &state) {
-  spirv::Scope executionScope;
-  spirv::GroupOperation groupOperation;
-  OpAsmParser::OperandType valueInfo;
-  if (parseEnumAttribute(executionScope, parser, state,
-                         kExecutionScopeAttrName) ||
-      parseEnumAttribute(groupOperation, parser, state,
-                         kGroupOperationAttrName) ||
-      parser.parseOperand(valueInfo))
-    return failure();
-
-  Optional<OpAsmParser::OperandType> clusterSizeInfo;
-  if (succeeded(parser.parseOptionalKeyword(kClusterSize))) {
-    clusterSizeInfo = OpAsmParser::OperandType();
-    if (parser.parseLParen() || parser.parseOperand(*clusterSizeInfo) ||
-        parser.parseRParen())
-      return failure();
-  }
-
-  Type resultType;
-  if (parser.parseColonType(resultType))
-    return failure();
-
-  if (parser.resolveOperand(valueInfo, resultType, state.operands))
-    return failure();
-
-  if (clusterSizeInfo.hasValue()) {
-    Type i32Type = parser.getBuilder().getIntegerType(32);
-    if (parser.resolveOperand(*clusterSizeInfo, i32Type, state.operands))
-      return failure();
-  }
-
-  return parser.addTypeToList(resultType, state.types);
-}
-
-static void print(spirv::GroupNonUniformIAddOp groupOp, OpAsmPrinter &printer) {
-  printer << spirv::GroupNonUniformIAddOp::getOperationName() << " \""
-          << stringifyScope(groupOp.execution_scope()) << "\" \""
-          << stringifyGroupOperation(groupOp.group_operation()) << "\" "
-          << groupOp.value();
-  if (!groupOp.cluster_size().empty())
-    printer << " " << kClusterSize << '(' << groupOp.cluster_size() << ')';
-  printer << " : " << groupOp.getType();
-}
-
-static LogicalResult verify(spirv::GroupNonUniformIAddOp groupOp) {
-  spirv::Scope scope = groupOp.execution_scope();
-  if (scope != spirv::Scope::Workgroup && scope != spirv::Scope::Subgroup)
-    return groupOp.emitOpError(
-        "execution scope must be 'Workgroup' or 'Subgroup'");
-
-  spirv::GroupOperation operation = groupOp.group_operation();
-  if (operation == spirv::GroupOperation::ClusteredReduce &&
-      groupOp.cluster_size().empty())
-    return groupOp.emitOpError("cluster size operand must be provided for "
-                               "'ClusteredReduce' group operation");
-
-  if (!groupOp.cluster_size().empty()) {
-    Operation *sizeOp = (*groupOp.cluster_size().begin()).getDefiningOp();
-    int32_t clusterSize = 0;
-
-    // TODO(antiagainst): support specialization constant here.
-    if (failed(extractValueFromConstOp(sizeOp, clusterSize)))
-      return groupOp.emitOpError(
-          "cluster size operand must come from a constant op");
-
-    if (!llvm::isPowerOf2_32(clusterSize))
-      return groupOp.emitOpError("cluster size operand must be a power of two");
-  }
-
-  return success();
-}
 
 //===----------------------------------------------------------------------===//
 // spv.IAdd
 //===----------------------------------------------------------------------===//
 
 OpFoldResult spirv::IAddOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 2 && "spv.IAdd expects two operands");
   // x + 0 = x

mlir/test/Dialect/SPIRV/Serialization/non-uniform-ops.mlir

spv.module "Logical" "GLSL450" {
 
   // CHECK-LABEL: @group_non_uniform_elect
   func @group_non_uniform_elect() -> i1 {
     // CHECK: %{{.+}} = spv.GroupNonUniformElect "Workgroup" : i1
     %0 = spv.GroupNonUniformElect "Workgroup" : i1
     spv.ReturnValue %0: i1
   }
 
+  // CHECK-LABEL: @group_non_uniform_fadd_reduce
+  func @group_non_uniform_fadd_reduce(%val: f32) -> f32 {
+    // CHECK: %{{.+}} = spv.GroupNonUniformFAdd "Workgroup" "Reduce" %{{.+}} : f32
+    %0 = spv.GroupNonUniformFAdd "Workgroup" "Reduce" %val : f32
+    spv.ReturnValue %0: f32
+  }
+
+  // CHECK-LABEL: @group_non_uniform_fmul_reduce
+  func @group_non_uniform_fmul_reduce(%val: f32) -> f32 {
+    // CHECK: %{{.+}} = spv.GroupNonUniformFMul "Workgroup" "Reduce" %{{.+}} : f32
+    %0 = spv.GroupNonUniformFMul "Workgroup" "Reduce" %val : f32
+    spv.ReturnValue %0: f32
+  }
+
   // CHECK-LABEL: @group_non_uniform_iadd_reduce
   func @group_non_uniform_iadd_reduce(%val: i32) -> i32 {
     // CHECK: %{{.+}} = spv.GroupNonUniformIAdd "Workgroup" "Reduce" %{{.+}} : i32
     %0 = spv.GroupNonUniformIAdd "Workgroup" "Reduce" %val : i32
     spv.ReturnValue %0: i32
   }
 
   // CHECK-LABEL: @group_non_uniform_iadd_clustered_reduce
   func @group_non_uniform_iadd_clustered_reduce(%val: vector<2xi32>) -> vector<2xi32> {
     %four = spv.constant 4 : i32
     // CHECK: %{{.+}} = spv.GroupNonUniformIAdd "Workgroup" "ClusteredReduce" %{{.+}} cluster_size(%{{.+}}) : vector<2xi32>
     %0 = spv.GroupNonUniformIAdd "Workgroup" "ClusteredReduce" %val cluster_size(%four) : vector<2xi32>
     spv.ReturnValue %0: vector<2xi32>
   }
+
+  // CHECK-LABEL: @group_non_uniform_imul_reduce
+  func @group_non_uniform_imul_reduce(%val: i32) -> i32 {
+    // CHECK: %{{.+}} = spv.GroupNonUniformIMul "Workgroup" "Reduce" %{{.+}} : i32
+    %0 = spv.GroupNonUniformIMul "Workgroup" "Reduce" %val : i32
+    spv.ReturnValue %0: i32
+  }
+
 }

mlir/test/Dialect/SPIRV/non-uniform-ops.mlir

func @group_non_uniform_elect() -> i1 {
   // expected-error @+1 {{execution scope must be 'Workgroup' or 'Subgroup'}}
   %0 = spv.GroupNonUniformElect "CrossDevice" : i1
   return %0: i1
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
+// spv.GroupNonUniformFAdd
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @group_non_uniform_fadd_reduce
+func @group_non_uniform_fadd_reduce(%val: f32) -> f32 {
+  // CHECK: %{{.+}} = spv.GroupNonUniformFAdd "Workgroup" "Reduce" %{{.+}} : f32
+  %0 = spv.GroupNonUniformFAdd "Workgroup" "Reduce" %val : f32
+  return %0: f32
+}
+
+// CHECK-LABEL: @group_non_uniform_fadd_clustered_reduce
+func @group_non_uniform_fadd_clustered_reduce(%val: vector<2xf32>) -> vector<2xf32> {
+  %four = spv.constant 4 : i32
+  // CHECK: %{{.+}} = spv.GroupNonUniformFAdd "Workgroup" "ClusteredReduce" %{{.+}} cluster_size(%{{.+}}) : vector<2xf32>
+  %0 = spv.GroupNonUniformFAdd "Workgroup" "ClusteredReduce" %val cluster_size(%four) : vector<2xf32>
+  return %0: vector<2xf32>
+}
+
+//===----------------------------------------------------------------------===//
+// spv.GroupNonUniformFMul
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @group_non_uniform_fmul_reduce
+func @group_non_uniform_fmul_reduce(%val: f32) -> f32 {
+  // CHECK: %{{.+}} = spv.GroupNonUniformFMul "Workgroup" "Reduce" %{{.+}} : f32
+  %0 = spv.GroupNonUniformFMul "Workgroup" "Reduce" %val : f32
+  return %0: f32
+}
+
+// CHECK-LABEL: @group_non_uniform_fmul_clustered_reduce
+func @group_non_uniform_fmul_clustered_reduce(%val: vector<2xf32>) -> vector<2xf32> {
+  %four = spv.constant 4 : i32
+  // CHECK: %{{.+}} = spv.GroupNonUniformFMul "Workgroup" "ClusteredReduce" %{{.+}} cluster_size(%{{.+}}) : vector<2xf32>
+  %0 = spv.GroupNonUniformFMul "Workgroup" "ClusteredReduce" %val cluster_size(%four) : vector<2xf32>
+  return %0: vector<2xf32>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
 // spv.GroupNonUniformIAdd
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @group_non_uniform_iadd_reduce
 func @group_non_uniform_iadd_reduce(%val: i32) -> i32 {
   // CHECK: %{{.+}} = spv.GroupNonUniformIAdd "Workgroup" "Reduce" %{{.+}} : i32
   %0 = spv.GroupNonUniformIAdd "Workgroup" "Reduce" %val : i32
   return %0: i32

 // -----
 
 func @group_non_uniform_iadd_clustered_reduce(%val: vector<2xi32>) -> vector<2xi32> {
   %five = spv.constant 5 : i32
   // expected-error @+1 {{cluster size operand must be a power of two}}
   %0 = spv.GroupNonUniformIAdd "Workgroup" "ClusteredReduce" %val cluster_size(%five) : vector<2xi32>
   return %0: vector<2xi32>
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GroupNonUniformIMul
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @group_non_uniform_imul_reduce
+func @group_non_uniform_imul_reduce(%val: i32) -> i32 {
+  // CHECK: %{{.+}} = spv.GroupNonUniformIMul "Workgroup" "Reduce" %{{.+}} : i32
+  %0 = spv.GroupNonUniformIMul "Workgroup" "Reduce" %val : i32
+  return %0: i32
+}
+
+// CHECK-LABEL: @group_non_uniform_imul_clustered_reduce
+func @group_non_uniform_imul_clustered_reduce(%val: vector<2xi32>) -> vector<2xi32> {
+  %four = spv.constant 4 : i32
+  // CHECK: %{{.+}} = spv.GroupNonUniformIMul "Workgroup" "ClusteredReduce" %{{.+}} cluster_size(%{{.+}}) : vector<2xi32>
+  %0 = spv.GroupNonUniformIMul "Workgroup" "ClusteredReduce" %val cluster_size(%four) : vector<2xi32>
+  return %0: vector<2xi32>
+}