[MLIR][SPIRV] Add definition for SPV_INTEL_split_barrier ops

The [`SPV_INTEL_split_barrier`
extension](https://htmlpreview.github.io/?https://github.com/KhronosGroup/SPIRV-Registry/blob/main/extensions/INTEL/SPV_INTEL_split_barrier.html)
defines operations to split control barrier semantics in two
operations. Add support for these operations (arrive and wait) to the
dialect.

Signed-off-by: Victor Perez <[email protected]>

Author: victor-eds

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

@@ -4466,6 +4466,8 @@ def SPIRV_OC_OpAssumeTrueKHR                : I32EnumAttrCase<"OpAssumeTrueKHR",
 def SPIRV_OC_OpAtomicFAddEXT                : I32EnumAttrCase<"OpAtomicFAddEXT", 6035>;
 def SPIRV_OC_OpConvertFToBF16INTEL          : I32EnumAttrCase<"OpConvertFToBF16INTEL", 6116>;
 def SPIRV_OC_OpConvertBF16ToFINTEL          : I32EnumAttrCase<"OpConvertBF16ToFINTEL", 6117>;
+def SPIRV_OC_OpControlBarrierArriveINTEL    : I32EnumAttrCase<"OpControlBarrierArriveINTEL", 6142>;
+def SPIRV_OC_OpControlBarrierWaitINTEL      : I32EnumAttrCase<"OpControlBarrierWaitINTEL", 6143>;
 def SPIRV_OC_OpGroupIMulKHR                 : I32EnumAttrCase<"OpGroupIMulKHR", 6401>;
 def SPIRV_OC_OpGroupFMulKHR                 : I32EnumAttrCase<"OpGroupFMulKHR", 6402>;
 
@@ -4556,7 +4558,9 @@ def SPIRV_OpcodeAttr :
       SPIRV_OC_OpCooperativeMatrixLengthKHR, SPIRV_OC_OpSubgroupBlockReadINTEL,
       SPIRV_OC_OpSubgroupBlockWriteINTEL, SPIRV_OC_OpAssumeTrueKHR,
       SPIRV_OC_OpAtomicFAddEXT, SPIRV_OC_OpConvertFToBF16INTEL,
-      SPIRV_OC_OpConvertBF16ToFINTEL, SPIRV_OC_OpGroupIMulKHR,
+      SPIRV_OC_OpConvertBF16ToFINTEL,
+      SPIRV_OC_OpControlBarrierArriveINTEL, SPIRV_OC_OpControlBarrierWaitINTEL,
+      SPIRV_OC_OpGroupIMulKHR,
       SPIRV_OC_OpGroupFMulKHR
     ]>;
 

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVIntelExtOps.td

@@ -111,6 +111,95 @@ def SPIRV_INTELConvertBF16ToFOp : SPIRV_IntelVendorOp<"ConvertBF16ToF", []> {
 }
 
 
+// -----
+
+class SPIRV_IntelSplitBarrierOp<string mnemonic>
+    : SPIRV_IntelVendorOp<mnemonic, []> {
+  let availability = [
+    MinVersion<SPIRV_V_1_0>,
+    MaxVersion<SPIRV_V_1_6>,
+    Extension<[SPV_INTEL_split_barrier]>,
+    Capability<[SPIRV_C_SplitBarrierINTEL]>
+  ];
+
+  let arguments = (ins
+    SPIRV_ScopeAttr:$execution_scope,
+    SPIRV_ScopeAttr:$memory_scope,
+    SPIRV_MemorySemanticsAttr:$memory_semantics
+  );
+
+  let results = (outs);
+
+  let assemblyFormat = [{
+    $execution_scope `,` $memory_scope `,` $memory_semantics attr-dict
+  }];
+
+  let hasVerifier = 0;
+}
+
+def SPIRV_INTELControlBarrierArriveOp
+    : SPIRV_IntelSplitBarrierOp<"ControlBarrierArrive"> {
+  let summary = "See extension SPV_INTEL_split_barrier";
+
+  let description = [{
+    Indicates that an invocation has arrived at a split control barrier. This
+    may allow other invocations waiting on the split control barrier to continue
+    executing.
+
+    When `Execution` is `Workgroup` or larger, behavior is undefined unless all
+    invocations within `Execution` execute the same dynamic instance of this
+    instruction. When `Execution` is `Subgroup` or `Invocation`, the behavior of
+    this instruction in non-uniform control flow is defined by the client API.
+
+    If `Semantics` is not `None`, this instruction also serves as the start of a
+    memory barrier similar to an `OpMemoryBarrier` instruction with the same
+    `Memory` and `Semantics` operands. This allows atomically specifying both a
+    control barrier and a memory barrier (that is, without needing two
+    instructions). If `Semantics` is `None`, `Memory` is ignored.
+
+    #### Example:
+
+    ```mlir
+    spirv.ControlBarrierArrive <Workgroup>, <Device>, <Acquire|UniformMemory>
+    ```
+  }];
+}
+
+
+// -----
+
+def SPIRV_INTELControlBarrierWaitOp
+    : SPIRV_IntelSplitBarrierOp<"ControlBarrierWait"> {
+  let summary = "See extension SPV_INTEL_split_barrier";
+
+  let description = [{
+    Waits for other invocations of this module to arrive at a split control
+    barrier.
+
+    When `Execution` is `Workgroup` or larger, behavior is undefined unless all
+    invocations within `Execution` execute the same dynamic instance of this
+    instruction. When `Execution` is `Subgroup` or `Invocation`, the behavior of
+    this instruction in non-uniform control flow is defined by the client API.
+
+    If `Semantics` is not `None`, this instruction also serves as the end of a
+    memory barrier similar to an `OpMemoryBarrier` instruction with the same
+    `Memory` and `Semantics` operands. This ensures that memory accesses issued
+    before arriving at the split barrier are observed before memory accesses
+    issued after this instruction. This control is ensured only for memory
+    accesses issued by this invocation and observed by another invocation
+    executing within `Memory` scope. This allows atomically specifying both a
+    control barrier and a memory barrier (that is, without needing two
+    instructions). If `Semantics` is `None`, `Memory` is ignored.
+
+    #### Example:
+
+    ```mlir
+    spirv.ControlBarrierWait <Workgroup>, <Device>, <Acquire|UniformMemory>
+    ```
+  }];
+}
+
+
 // -----
 
 #endif // MLIR_DIALECT_SPIRV_IR_INTEL_EXT_OPS

mlir/test/Dialect/SPIRV/IR/intel-ext-ops.mlir

@@ -69,3 +69,17 @@ spirv.func @bf16_to_f32_vec_unsupported(%arg0 : vector<2xi16>) "None" {
   %0 = spirv.INTEL.ConvertBF16ToF %arg0 : vector<2xi16> to vector<3xf32>
   spirv.Return
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.INTEL.SplitBarrier
+//===----------------------------------------------------------------------===//
+
+spirv.func @split_barrier() "None" {
+  // CHECK: spirv.INTEL.ControlBarrierArrive <Workgroup>, <Device>, <Acquire|UniformMemory>
+  spirv.INTEL.ControlBarrierArrive <Workgroup>, <Device>, <Acquire|UniformMemory>
+  // CHECK: spirv.INTEL.ControlBarrierWait <Workgroup>, <Device>, <Acquire|UniformMemory>
+  spirv.INTEL.ControlBarrierWait <Workgroup>, <Device>, <Acquire|UniformMemory>
+  spirv.Return
+}

mlir/test/Target/SPIRV/intel-ext-ops.mlir

@@ -29,3 +29,21 @@ spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Bfloat16ConversionINTEL]
     spirv.Return
   }
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.INTEL.SplitBarrier
+//===----------------------------------------------------------------------===//
+
+// CHECK: spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [SplitBarrierINTEL], [SPV_INTEL_split_barrier]>
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [SplitBarrierINTEL], [SPV_INTEL_split_barrier]> {
+  // CHECK-LABEL: @split_barrier
+  spirv.func @split_barrier() "None" {
+    // CHECK: spirv.INTEL.ControlBarrierArrive <Workgroup>, <Device>, <Acquire|UniformMemory>
+    spirv.INTEL.ControlBarrierArrive <Workgroup>, <Device>, <Acquire|UniformMemory>
+    // CHECK: spirv.INTEL.ControlBarrierWait <Workgroup>, <Device>, <Acquire|UniformMemory>
+    spirv.INTEL.ControlBarrierWait <Workgroup>, <Device>, <Acquire|UniformMemory>
+    spirv.Return
+  }
+}