[mlir][spirv] Add GroupNonUniformBroadcastOp

[abialas]

Summary

Added GroupNonUniformBroadcastOp to spirv dialect.

[ThomasRaoux]

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2064

I think you can remove both cast.

[abialas]

removed static_casts

[abialas]

Thanks Thomas for the review, may I ask you to push this change to llvm master?
Btw, to not bother you again with this, is there a process to get RW access to the repo so I can push future changes on my own?

[ThomasRaoux]

In D87688#2278471, @abialas wrote:

Thanks Thomas for the review, may I ask you to push this change to llvm master?
Btw, to not bother you again with this, is there a process to get RW access to the repo so I can push future changes on my own?

Sure, I'll push it in a bit. You can obtain commit access by following the process described here:
https://llvm.org/docs/DeveloperPolicy.html#obtaining-commit-access

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_OpGroupBroadcast            : I32EnumAttrCase<"OpGroupBroadcast", 263>;
 def SPV_OC_OpNoLine                    : I32EnumAttrCase<"OpNoLine", 317>;
 def SPV_OC_OpModuleProcessed           : I32EnumAttrCase<"OpModuleProcessed", 330>;
 def SPV_OC_OpGroupNonUniformElect      : I32EnumAttrCase<"OpGroupNonUniformElect", 333>;
+def SPV_OC_OpGroupNonUniformBroadcast  : I32EnumAttrCase<"OpGroupNonUniformBroadcast", 337>;
 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_OpGroupNonUniformSMin       : I32EnumAttrCase<"OpGroupNonUniformSMin", 353>;
 def SPV_OC_OpGroupNonUniformUMin       : I32EnumAttrCase<"OpGroupNonUniformUMin", 354>;
 def SPV_OC_OpGroupNonUniformFMin       : I32EnumAttrCase<"OpGroupNonUniformFMin", 355>;
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_OpGroupBroadcast,
       SPV_OC_OpNoLine, SPV_OC_OpModuleProcessed, SPV_OC_OpGroupNonUniformElect,
-      SPV_OC_OpGroupNonUniformBallot, SPV_OC_OpGroupNonUniformIAdd,
-      SPV_OC_OpGroupNonUniformFAdd, SPV_OC_OpGroupNonUniformIMul,
-      SPV_OC_OpGroupNonUniformFMul, SPV_OC_OpGroupNonUniformSMin,
-      SPV_OC_OpGroupNonUniformUMin, SPV_OC_OpGroupNonUniformFMin,
-      SPV_OC_OpGroupNonUniformSMax, SPV_OC_OpGroupNonUniformUMax,
-      SPV_OC_OpGroupNonUniformFMax, SPV_OC_OpSubgroupBallotKHR,
-      SPV_OC_OpTypeCooperativeMatrixNV, SPV_OC_OpCooperativeMatrixLoadNV,
-      SPV_OC_OpCooperativeMatrixStoreNV, SPV_OC_OpCooperativeMatrixMulAddNV,
-      SPV_OC_OpCooperativeMatrixLengthNV, SPV_OC_OpSubgroupBlockReadINTEL,
-      SPV_OC_OpSubgroupBlockWriteINTEL
+      SPV_OC_OpGroupNonUniformBroadcast, SPV_OC_OpGroupNonUniformBallot,
+      SPV_OC_OpGroupNonUniformIAdd, SPV_OC_OpGroupNonUniformFAdd,
+      SPV_OC_OpGroupNonUniformIMul, SPV_OC_OpGroupNonUniformFMul,
+      SPV_OC_OpGroupNonUniformSMin, SPV_OC_OpGroupNonUniformUMin,
+      SPV_OC_OpGroupNonUniformFMin, SPV_OC_OpGroupNonUniformSMax,
+      SPV_OC_OpGroupNonUniformUMax, SPV_OC_OpGroupNonUniformFMax,
+      SPV_OC_OpSubgroupBallotKHR, SPV_OC_OpTypeCooperativeMatrixNV,
+      SPV_OC_OpCooperativeMatrixLoadNV, SPV_OC_OpCooperativeMatrixStoreNV,
+      SPV_OC_OpCooperativeMatrixMulAddNV, SPV_OC_OpCooperativeMatrixLengthNV,
+      SPV_OC_OpSubgroupBlockReadINTEL, SPV_OC_OpSubgroupBlockWriteINTEL
     ]>;
 
 // End opcode section. Generated from SPIR-V spec; DO NOT MODIFY!
 
 //===----------------------------------------------------------------------===//
 // SPIR-V op definitions
 //===----------------------------------------------------------------------===//
 

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

def SPV_GroupNonUniformBallotOp : SPV_Op<"GroupNonUniformBallot", []> {
 
   let assemblyFormat = [{
     $execution_scope $predicate attr-dict `:` type($result)
   }];
 }
 
 // -----
 
+def SPV_GroupNonUniformBroadcastOp : SPV_Op<"GroupNonUniformBroadcast",
+  [NoSideEffect, AllTypesMatch<["value", "result"]>]> {
+  let summary = [{
+    Return the Value of the invocation identified by the id Id to all active
+    invocations in the group.
+  }];
+
+  let description = [{
+    Result Type  must be a scalar or vector of floating-point type, integer
+    type, or Boolean type.
+
+    Execution must be Workgroup or Subgroup Scope.
+
+     The type of Value must be the same as Result Type.
+
+    Id  must be a scalar of integer type, whose Signedness operand is 0.
+
+    Before version 1.5, Id must come from a constant instruction. Starting
+    with version 1.5, Id must be dynamically uniform.
+
+    The resulting value is undefined if Id is an inactive invocation, or is
+    greater than or equal to the size of the group.
+
+    <!-- End of AutoGen section -->
+
+    ```
+    scope ::= `"Workgroup"` | `"Subgroup"`
+    integer-float-scalar-vector-type ::= integer-type | float-type |
+                               `vector<` integer-literal `x` integer-type `>` |
+                               `vector<` integer-literal `x` float-type `>`
+    group-non-uniform-broadcast-op ::= ssa-id `=` 
+              `spv.GroupNonUniformBroadcast` scope ssa_use,
+                ssa_use `:` integer-float-scalar-vector-type `,` integer-type
+    ```mlir
+
+    #### Example:
+
+    ```
+    %scalar_value = ... : f32
+    %vector_value = ... : vector<4xf32>
+    %id = ... : i32
+    %0 = spv.GroupNonUniformBroadcast "Subgroup" %scalar_value, %id : f32, i32
+    %1 = spv.GroupNonUniformBroadcast "Workgroup" %vector_value, %id :
+      vector<4xf32>, i32
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPV_V_1_3>,
+    MaxVersion<SPV_V_1_5>,
+    Extension<[]>,
+    Capability<[SPV_C_GroupNonUniformBallot]>
+  ];
+
+  let arguments = (ins
+    SPV_ScopeAttr:$execution_scope,
+    SPV_Type:$value,
+    SPV_Integer:$id
+  );
+
+  let results = (outs
+    SPV_Type:$result
+  );
+
+  let assemblyFormat = [{
+    $execution_scope operands attr-dict `:` type($value) `,` type($id)
+  }];
+}
+
+// -----
+
 def SPV_GroupNonUniformElectOp : SPV_Op<"GroupNonUniformElect", []> {
   let summary = [{
     Result is true only in the active invocation with the lowest id in the
     group, otherwise result is false.
   }];
 
   let description = [{
     Result Type must be a Boolean type.
def SPV_GroupNonUniformFMulOp :
   }];
 }
 
 // -----
 
 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.
+    An integer 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 integer type.
 
     Execution must be Workgroup or Subgroup Scope.
 
     The identity I for Operation is 0. If Operation is ClusteredReduce,

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
 
 #include "mlir/Dialect/SPIRV/ParserUtils.h"
 #include "mlir/Dialect/SPIRV/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/SPIRVTypes.h"
+#include "mlir/Dialect/SPIRV/TargetAndABI.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/bit.h"
static LogicalResult verify(spirv::GroupNonUniformBallotOp ballotOp) {
   if (scope != spirv::Scope::Workgroup && scope != spirv::Scope::Subgroup)
     return ballotOp.emitOpError(
         "execution scope must be 'Workgroup' or 'Subgroup'");
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// spv.GroupNonUniformBroadcast
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(spirv::GroupNonUniformBroadcastOp broadcastOp) {
+  spirv::Scope scope = broadcastOp.execution_scope();
+  if (scope != spirv::Scope::Workgroup && scope != spirv::Scope::Subgroup)
+    return broadcastOp.emitOpError(
+        "execution scope must be 'Workgroup' or 'Subgroup'");
+
+  // SPIR-V spec: "Before version 1.5, Id must come from a
+  // constant instruction.
+  auto targetEnv = spirv::getDefaultTargetEnv(broadcastOp.getContext());
+  if (auto spirvModule = broadcastOp.getParentOfType<spirv::ModuleOp>())
+    targetEnv = spirv::lookupTargetEnvOrDefault(spirvModule);
+
+  if (targetEnv.getVersion() < spirv::Version::V_1_5) {
+    auto *idOp = broadcastOp.id().getDefiningOp();
+    if (!idOp || !isa<spirv::ConstantOp,           // for normal constant
+                      spirv::ReferenceOfOp>(idOp)) // for spec constant
+      return broadcastOp.emitOpError("id must be the result of a constant op");
+  }
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // spv.SubgroupBlockReadINTEL
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseSubgroupBlockReadINTELOp(OpAsmParser &parser,
                                                  OperationState &state) {
   // Parse the storage class specification
   spirv::StorageClass storageClass;
   OpAsmParser::OperandType ptrInfo;

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

 // RUN: mlir-translate -test-spirv-roundtrip -split-input-file %s | FileCheck %s
 
 spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
   // CHECK-LABEL: @group_non_uniform_ballot
   spv.func @group_non_uniform_ballot(%predicate: i1) -> vector<4xi32> "None" {
     // CHECK: %{{.*}} = spv.GroupNonUniformBallot "Workgroup" %{{.*}}: vector<4xi32>
   %0 = spv.GroupNonUniformBallot "Workgroup" %predicate : vector<4xi32>
     spv.ReturnValue %0: vector<4xi32>
   }
 
+  // CHECK-LABEL: @group_non_uniform_broadcast
+  spv.func @group_non_uniform_broadcast(%value: f32) -> f32 "None" {
+    %one = spv.constant 1 : i32
+    // CHECK: spv.GroupNonUniformBroadcast "Subgroup" %{{.*}}, %{{.*}} : f32, i32
+    %0 = spv.GroupNonUniformBroadcast "Subgroup" %value, %one : f32, i32
+    spv.ReturnValue %0: f32
+  }
+
   // CHECK-LABEL: @group_non_uniform_elect
   spv.func @group_non_uniform_elect() -> i1 "None" {
     // CHECK: %{{.+}} = spv.GroupNonUniformElect "Workgroup" : i1
     %0 = spv.GroupNonUniformElect "Workgroup" : i1
     spv.ReturnValue %0: i1
   }
 
   // CHECK-LABEL: @group_non_uniform_fadd_reduce

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

func @group_non_uniform_ballot(%predicate: i1) -> vector<4xsi32> {
   // expected-error @+1 {{op result #0 must be vector of 8/16/32/64-bit signless/unsigned integer values of length 4, but got 'vector<4xsi32>'}}
   %0 = spv.GroupNonUniformBallot "Workgroup" %predicate : vector<4xsi32>
   return %0: vector<4xsi32>
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
+// spv.NonUniformGroupBroadcast
+//===----------------------------------------------------------------------===//
+
+func @group_non_uniform_broadcast_scalar(%value: f32) -> f32 {
+  %one = spv.constant 1 : i32
+  // CHECK: spv.GroupNonUniformBroadcast "Workgroup" %{{.*}}, %{{.*}} : f32, i32
+  %0 = spv.GroupNonUniformBroadcast "Workgroup" %value, %one : f32, i32
+  return %0: f32
+}
+
+// -----
+
+func @group_non_uniform_broadcast_vector(%value: vector<4xf32>) -> vector<4xf32> {
+  %one = spv.constant 1 : i32
+  // CHECK: spv.GroupNonUniformBroadcast "Subgroup" %{{.*}}, %{{.*}} : vector<4xf32>, i32
+  %0 = spv.GroupNonUniformBroadcast "Subgroup" %value, %one : vector<4xf32>, i32
+  return %0: vector<4xf32>
+}
+
+// -----
+
+func @group_non_uniform_broadcast_negative_scope(%value: f32, %localid: i32 ) -> f32 {
+  %one = spv.constant 1 : i32
+  // expected-error @+1 {{execution scope must be 'Workgroup' or 'Subgroup'}} 
+  %0 = spv.GroupNonUniformBroadcast "Device" %value, %one : f32, i32
+  return %0: f32
+}
+
+// -----
+
+func @group_non_uniform_broadcast_negative_non_const(%value: f32, %localid: i32) -> f32 {
+  // expected-error @+1 {{id must be the result of a constant op}}
+  %0 = spv.GroupNonUniformBroadcast "Subgroup" %value, %localid : f32, i32
+  return %0: f32
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
 // spv.GroupNonUniformElect
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @group_non_uniform_elect
 func @group_non_uniform_elect() -> i1 {
   // CHECK: %{{.+}} = spv.GroupNonUniformElect "Workgroup" : i1
   %0 = spv.GroupNonUniformElect "Workgroup" : i1
   return %0: i1