Skip to content

[MLIR][SPIRV] Add definition for SPV_INTEL_split_barrier ops #115738

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Nov 18, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
Original file line number Diff line number Diff line change
Expand Up @@ -4512,6 +4512,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>;

Expand Down Expand Up @@ -4602,7 +4604,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
]>;

Expand Down
89 changes: 89 additions & 0 deletions mlir/include/mlir/Dialect/SPIRV/IR/SPIRVIntelExtOps.td
Original file line number Diff line number Diff line change
Expand Up @@ -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
14 changes: 14 additions & 0 deletions mlir/test/Dialect/SPIRV/IR/intel-ext-ops.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,20 @@ spirv.func @bf16_to_f32_vec_unsupported(%arg0 : vector<2xi16>) "None" {

// -----

//===----------------------------------------------------------------------===//
// 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
}

// -----

//===----------------------------------------------------------------------===//
// spirv.INTEL.CacheControls
//===----------------------------------------------------------------------===//
Expand Down
18 changes: 18 additions & 0 deletions mlir/test/Target/SPIRV/intel-ext-ops.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}
Loading