[mlir][spirv] Add OpGroupBroadcast

[abialas]

Summary

OpGroupBroadcast added to SPIRV dialect

[ThomasRaoux]

mlir/include/mlir/Dialect/SPIRV/SPIRVGroupOps.td

89

You can use AllTypesMatch<["value", "result"]> to enforce that the source and result have the same type and then you can remove it from the assembly format to avoid having redundant types.

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2020–2024

nit: Don't use '{' for single line 'if'

mlir/test/Dialect/SPIRV/group-ops.mlir

19–23

Would be good to have an extra positive test with vector of localid

[mravishankar]

We also need tests for roundtripping the serialization to SPIR-V binary. Please see this section : https://mlir.llvm.org/docs/Dialects/SPIR-V/#add-a-new-op

"Tests for the op’s custom assembly form and verification should be added to the proper file in test/Dialect/SPIRV/.

The generated op will automatically gain the logic for (de)serialization. However, tests still need to be coupled with the change to make sure no surprises. Serialization tests live in test/Dialect/SPIRV/Serialization."

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

2013

As Thomas suggested, if you add the AllTypesMatch this would not be needed.

[abialas]

Changes added after code review

[abialas]

In D85435#2200646, @mravishankar wrote:

We also need tests for roundtripping the serialization to SPIR-V binary. Please see this section : https://mlir.llvm.org/docs/Dialects/SPIR-V/#add-a-new-op

"Tests for the op’s custom assembly form and verification should be added to the proper file in test/Dialect/SPIRV/.

The generated op will automatically gain the logic for (de)serialization. However, tests still need to be coupled with the change to make sure no surprises. Serialization tests live in test/Dialect/SPIRV/Serialization."

Thank You for the review. Regarding Serialization tests, is there anything more needed than what I already added to mlir/test/Dialect/SPIRV/Serialization/group-ops.mlir?

[ThomasRaoux]

Thanks Artur. Looks good to me.

[mravishankar]

mlir/include/mlir/Dialect/SPIRV/SPIRVGroupOps.td

50

Nit: Refactor this to be within 80 chars.

[mravishankar]

mlir/test/Dialect/SPIRV/Serialization/group-ops.mlir

10

Could you also add one for vector<..xi32> for localID.

Nit: Prefer the CHECK-* to be formatted like this to make it easier to read

CHECK-LABEL: func @group_broadcast
      CHECK:   spv.GroupBroadcast ...

i.e.
Align the ':'
and indent the check statements to make what is being checked more readable.
Suport nit, you dont need the leadeing %{{.*}} = if you are not trying to capture the value.

[abialas]

Changes after code review

[abialas]

I added remaining changes requested. Regarding the CHECK-* format I believe it is consistent with other tests.
Since I do not have rw access to llvm repo, may I ask someone to push this change to trunk?

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

 def SPV_OC_OpLoopMerge                 : I32EnumAttrCase<"OpLoopMerge", 246>;
 def SPV_OC_OpSelectionMerge            : I32EnumAttrCase<"OpSelectionMerge", 247>;
 def SPV_OC_OpLabel                     : I32EnumAttrCase<"OpLabel", 248>;
 def SPV_OC_OpBranch                    : I32EnumAttrCase<"OpBranch", 249>;
 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_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>;
SPV_I32EnumAttr<"Opcode", "valid SPIR-V instructions", [
       SPV_OC_OpBitFieldSExtract, SPV_OC_OpBitFieldUExtract, SPV_OC_OpBitReverse,
       SPV_OC_OpBitCount, SPV_OC_OpControlBarrier, SPV_OC_OpMemoryBarrier,
       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_OpNoLine,
-      SPV_OC_OpModuleProcessed, SPV_OC_OpGroupNonUniformElect,
+      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,

mlir/include/mlir/Dialect/SPIRV/SPIRVGroupOps.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SPIRV_GROUP_OPS
 #define SPIRV_GROUP_OPS
 
 // -----
 
+def SPV_GroupBroadcastOp : SPV_Op<"GroupBroadcast",
+  [NoSideEffect, AllTypesMatch<["value", "result"]>]> {
+  let summary = [{
+    Return the Value of the invocation identified by the local id LocalId to
+    all invocations in the group.
+  }];
+
+  let description = [{
+    All invocations of this module within Execution must reach this point of
+    execution.
+
+    Behavior is undefined if this instruction is used in control flow that
+    is non-uniform within Execution.
+
+    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.
+
+    LocalId must be an integer datatype. It can be a scalar, or a vector
+    with 2 components or a vector with 3 components. LocalId must be the
+    same for all invocations in 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 `>`
+    localid-type ::= integer-type |
+                   `vector<` integer-literal `x` integer-type `>`
+    group-broadcast-op ::= ssa-id `=` `spv.GroupBroadcast` scope ssa_use, ssa_use
+                         `:` integer-float-scalar-vector-type `,` localid-type
+    ```mlir
+
+    #### Example:
+
+    ```
+    %scalar_value = ... : f32
+    %vector_value = ... : vector<4xf32>
+    %scalar_localid = ... : i32
+    %vector_localid = ... : vector<3xi32>
+    %0 = spv.GroupBroadcast "Subgroup" %scalar_value, %scalar_localid : f32, i32
+    %1 = spv.GroupBroadcast "Workgroup" %vector_value, %vector_localid :
+      vector<4xf32>, vector<3xi32>
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPV_V_1_0>,
+    MaxVersion<SPV_V_1_5>,
+    Extension<[]>,
+    Capability<[SPV_C_Groups]>
+  ];
+
+  let arguments = (ins
+    SPV_ScopeAttr:$execution_scope,
+    SPV_Type:$value,
+    SPV_ScalarOrVectorOf<SPV_Integer>:$localid
+  );
+
+  let results = (outs
+    SPV_Type:$result
+  );
+
+  let assemblyFormat = [{
+    $execution_scope operands attr-dict `:` type($value) `,` type($localid)
+  }];
+
+}
+
+// -----
+
 def SPV_SubgroupBallotKHROp : SPV_Op<"SubgroupBallotKHR", []> {
   let summary = "See extension SPV_KHR_shader_ballot";
 
   let description = [{
     Computes a bitfield value combining the Predicate value from all invocations
     in the current Subgroup that execute the same dynamic instance of this
     instruction. The bit is set to one if the corresponding invocation is active
     and the predicate is evaluated to true; otherwise, it is set to zero.

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

if (!initOp ||
                                "spv.specConstant or spv.globalVariable op");
     }
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// spv.GroupBroadcast
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(spirv::GroupBroadcastOp 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'");
+
+  if (auto localIdTy = broadcastOp.localid().getType().dyn_cast<VectorType>())
+    if (!(localIdTy.getNumElements() == 2 || localIdTy.getNumElements() == 3))
+      return broadcastOp.emitOpError("localid is a vector and can be with only "
+                                     " 2 or 3 components, actual number is ")
+             << localIdTy.getNumElements();
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // spv.GroupNonUniformBallotOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(spirv::GroupNonUniformBallotOp ballotOp) {
   spirv::Scope scope = ballotOp.execution_scope();
   if (scope != spirv::Scope::Workgroup && scope != spirv::Scope::Subgroup)
     return ballotOp.emitOpError(
         "execution scope must be 'Workgroup' or 'Subgroup'");
 
   return success();

mlir/test/Dialect/SPIRV/Serialization/group-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: @subgroup_ballot
   spv.func @subgroup_ballot(%predicate: i1) -> vector<4xi32> "None" {
     // CHECK: %{{.*}} = spv.SubgroupBallotKHR %{{.*}}: vector<4xi32>
     %0 = spv.SubgroupBallotKHR %predicate: vector<4xi32>
     spv.ReturnValue %0: vector<4xi32>
   }
+  // CHECK-LABEL: @group_broadcast
+  spv.func @group_broadcast(%value: f32, %localid: i32 ) -> f32 "None" {
+    // CHECK: %{{.*}} = spv.GroupBroadcast "Workgroup" %{{.*}}, %{{.*}}: f32, i32
+    %0 = spv.GroupBroadcast "Workgroup" %value, %localid : f32, i32
+    spv.ReturnValue %0: f32
+  }
 }

mlir/test/Dialect/SPIRV/group-ops.mlir

 // RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s
 
 //===----------------------------------------------------------------------===//
 // spv.SubgroupBallotKHR
 //===----------------------------------------------------------------------===//
 
 func @subgroup_ballot(%predicate: i1) -> vector<4xi32> {
   // CHECK: %{{.*}} = spv.SubgroupBallotKHR %{{.*}} : vector<4xi32>
   %0 = spv.SubgroupBallotKHR %predicate: vector<4xi32>
   return %0: vector<4xi32>
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.GroupBroadcast
+//===----------------------------------------------------------------------===//
+
+func @group_broadcast_scalar(%value: f32, %localid: i32 ) -> f32 {
+  // CHECK: %{{.*}} = spv.GroupBroadcast "Workgroup" %{{.*}}, %{{.*}} : f32, i32
+  %0 = spv.GroupBroadcast "Workgroup" %value, %localid : f32, i32
+  return %0: f32
+}
+
+// -----
+
+func @group_broadcast_scalar_vector(%value: f32, %localid: vector<3xi32> ) -> f32 {
+  // CHECK: %{{.*}} = spv.GroupBroadcast "Workgroup" %{{.*}}, %{{.*}} : f32, vector<3xi32>
+  %0 = spv.GroupBroadcast "Workgroup" %value, %localid : f32, vector<3xi32>
+  return %0: f32
+}
+
+// -----
+
+func @group_broadcast_vector(%value: vector<4xf32>, %localid: vector<3xi32> ) -> vector<4xf32> {
+  // CHECK: %{{.*}} = spv.GroupBroadcast "Subgroup" %{{.*}}, %{{.*}} : vector<4xf32>, vector<3xi32>
+  %0 = spv.GroupBroadcast "Subgroup" %value, %localid : vector<4xf32>, vector<3xi32>
+  return %0: vector<4xf32>
+}
+
+// -----
+
+func @group_broadcast_negative_scope(%value: f32, %localid: vector<3xi32> ) -> f32 {
+  // expected-error @+1 {{execution scope must be 'Workgroup' or 'Subgroup'}} 
+  %0 = spv.GroupBroadcast "Device" %value, %localid : f32, vector<3xi32>
+  return %0: f32
+}
+
+// -----
+
+func @group_broadcast_negative_locid_dtype(%value: f32, %localid: vector<3xf32> ) -> f32 {
+  // expected-error @+1 {{operand #1 must be 8/16/32/64-bit integer or vector of 8/16/32/64-bit integer values}}
+  %0 = spv.GroupBroadcast "Subgroup" %value, %localid : f32, vector<3xf32>
+  return %0: f32
+}
+
+// -----
+
+func @group_broadcast_negative_locid_vec4(%value: f32, %localid: vector<4xi32> ) -> f32 {
+  // expected-error @+1 {{localid is a vector and can be with only  2 or 3 components, actual number is 4}}
+  %0 = spv.GroupBroadcast "Subgroup" %value, %localid : f32, vector<4xi32>
+  return %0: f32
+}