2258 lines
97 KiB
TableGen
2258 lines
97 KiB
TableGen
//===-- AMDGPUOps.td - AMDGPU dialect operations *- 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
|
||
//
|
||
//===----------------------------------------------------------------------===//
|
||
|
||
#ifndef MLIR_DIALECT_AMDGPU_IR_AMDGPUOPS_TD
|
||
#define MLIR_DIALECT_AMDGPU_IR_AMDGPUOPS_TD
|
||
|
||
include "mlir/Dialect/AMDGPU/IR/AMDGPUBase.td"
|
||
include "mlir/Dialect/AMDGPU/IR/AMDGPUAttrs.td"
|
||
include "mlir/Dialect/AMDGPU/IR/AMDGPUTypes.td"
|
||
|
||
include "mlir/Interfaces/InferTypeOpInterface.td"
|
||
include "mlir/Interfaces/SideEffectInterfaces.td"
|
||
include "mlir/Interfaces/ViewLikeInterface.td"
|
||
include "mlir/IR/OpBase.td"
|
||
|
||
//===----------------------------------------------------------------------===//
|
||
// AMDGPU common type constraints
|
||
//===----------------------------------------------------------------------===//
|
||
|
||
class AMDGPU_ConcreteVector<Type elem, int length> :
|
||
FixedVectorOfLengthAndType<[length], [elem]>,
|
||
BuildableType<
|
||
"::mlir::VectorType::get({" # length # "} ,"
|
||
# elem.builderCall # ")">;
|
||
|
||
def AnyIntegerOrFloat : AnyTypeOf<[AnySignlessInteger, AnyFloat], "Integer or Float">;
|
||
|
||
def AnyIntegerOrFloatOr1DVector :
|
||
AnyTypeOf<[AnyIntegerOrFloat, FixedVectorOfRankAndType<[1], [AnyIntegerOrFloat]>]>;
|
||
|
||
// Types with element width up to 64 bits, used to keep dpp operands legal.
|
||
def AMDGPU_IntOrFloatWidthLeq64 : Type<
|
||
CPred<"$_self.isIntOrFloat() && $_self.getIntOrFloatBitWidth() <= 64">,
|
||
"integer or float with element bitwidth <= 64">;
|
||
|
||
def AMDGPU_IntOrFloatOr1DVectorWidthLeq64 :
|
||
AnyTypeOf<[
|
||
AMDGPU_IntOrFloatWidthLeq64,
|
||
FixedVectorOfRankAndType<[1], [AMDGPU_IntOrFloatWidthLeq64]>
|
||
]>;
|
||
|
||
//===----------------------------------------------------------------------===//
|
||
// AMDGPU Op definitions
|
||
//===----------------------------------------------------------------------===//
|
||
|
||
class AMDGPU_Op<string mnemonic, list<Trait> traits = []> :
|
||
Op<AMDGPU_Dialect, mnemonic, traits> {}
|
||
|
||
def AMDGPU_ExtPackedFp8Op :
|
||
AMDGPU_Op<"ext_packed_fp8", [Pure]>,
|
||
Arguments<(ins AnyTypeOf<[F8E5M2FNUZ, F8E4M3FNUZ, F8E5M2, F8E4M3FN,
|
||
VectorOfLengthAndType<[1, 2, 3, 4], [F8E5M2FNUZ, F8E4M3FNUZ, F8E5M2, F8E4M3FN]>]>:$source,
|
||
ConfinedAttr<I32Attr, [IntNonNegative, IntMaxValue<3>]>:$index)>,
|
||
Results<(outs AnyTypeOf<[F32, FixedVectorOfLengthAndType<[2], [F32]>]>:$res)> {
|
||
let summary = "Extend a fp8 value to a float or a vector of packed fp8 values to two floats";
|
||
|
||
let description = [{
|
||
Extend one or two 8-bit floats in `source[index]` to a 32-bit float or
|
||
two floats and return them.
|
||
|
||
This rather unusual signature arises from the fact that AMD GPUs cannot
|
||
easily work with sub 32-bit quantities, so the compiler intrinsics for
|
||
extending 8-bit floats (which are, currently, the only way to work with
|
||
this operation) take packed vectors of 4 such floats.
|
||
|
||
If the passed-in vector has fewer than four elements, or the input is scalar,
|
||
the remaining values in the <4 x i8> will be filled with
|
||
undefined values as needed.
|
||
}];
|
||
let assemblyFormat = [{
|
||
attr-dict $source `[` $index `]` `:` type($source) `to` type($res)
|
||
}];
|
||
}
|
||
|
||
def AMDGPU_ScaledExtPackedMatrixOp
|
||
: AMDGPU_Op<"scaled_ext_packed_matrix", [Pure, AllShapesMatch<["source", "res"]>]>,
|
||
Arguments<(
|
||
ins AnyTypeOf<[FixedVectorOfShapeAndType<[8], F4E2M1FN>,
|
||
FixedVectorOfShapeAndType<[8], F8E4M3FN>,
|
||
FixedVectorOfShapeAndType<[8], F8E5M2>,
|
||
FixedVectorOfShapeAndType<[16], F6E2M3FN>,
|
||
FixedVectorOfShapeAndType<[16], F6E3M2FN>]>:$source,
|
||
FixedVectorOfShapeAndType<[4], F8E8M0FNU>:$scale,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[16, 32]>]>:$blockSize,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[0, 16]>]>:$firstScaleLane,
|
||
ConfinedAttr<I32Attr, [IntMinValue<0>, IntMaxValue<3>]>:$firstScaleByte)>,
|
||
Results<(
|
||
outs AnyTypeOf<[FixedVectorOfShapeAndType<[8], F32>,
|
||
FixedVectorOfShapeAndType<[8], F16>,
|
||
FixedVectorOfShapeAndType<[8], BF16>,
|
||
FixedVectorOfShapeAndType<[16], F32>,
|
||
FixedVectorOfShapeAndType<[16], F16>,
|
||
FixedVectorOfShapeAndType<[16], BF16>]>:$res)> {
|
||
|
||
let summary = "Extend a wave-wide matrix of packed floating point values";
|
||
|
||
let description = [{
|
||
Extend matrix of microfloats (8 or 16 elements per lane) using a set of scales
|
||
that may be stored on other lanes.
|
||
|
||
The scales applied to the input microfloats are stored in bytes which
|
||
come from the `scales` input provided in a *half* of the wave identified
|
||
by `firstScaleLane`. The bytes used is selected by `firstScaleByte` and depends
|
||
on the type of `source`. The 16 vectors in consecutive lanes starting from
|
||
`firstScaleLane` (which we'll call the scale vectors) will be used by both
|
||
halves of the wave (with lane L reading from L % 16'th scale vector).
|
||
|
||
When `source` is either F4E2M1FN, F6E2M3FN, or F6E3M2FN each half of the
|
||
wave will use a different byte. The first one being `firstScaleByte` and
|
||
the second one being `firstScaleByte` + 1. When the block size is 32,
|
||
`firstScaleByte` can be either 0 or 2, selecting halves of the scale vectors.
|
||
Lanes 0-15 will read from `firstScaleByte` and lanes 16-31 will read
|
||
from `firstScaleByte` + 1.
|
||
|
||
|
||
For example:
|
||
```mlir
|
||
// Input: 8-element vector of F8E4M3FN, converting to F32
|
||
// Lanes 0-15 read from byte 0, lanes 16-31 read from byte 1
|
||
%result = amdgpu.scaled_ext_packed_matrix %source scale(%scales)
|
||
blockSize(32) firstScaleLane(0) firstScaleByte(0)
|
||
: vector<8xf8E4M3FN>, vector<4xf8E8M0FNU> -> vector<8xf32>
|
||
|
||
// Input: 16-element vector of F6E2M3FN, converting to F16
|
||
// Lanes 0-15 read from byte 2, lanes 16-31 read from byte 3
|
||
%result = amdgpu.scaled_ext_packed_matrix %source scale(%scales)
|
||
blockSize(32) firstScaleLane(16) firstScaleByte(2)
|
||
: vector<16xf6E2M3FN>, vector<4xf8E8M0FNU> -> vector<16xf16>
|
||
```
|
||
|
||
When `source` is either F4E2M1FN, F6E2M3FN, or F6E3M2FN and
|
||
the block size is 16, `firstScaleByte` can be 0 or 1.
|
||
Lanes 0-15 read from the `firstScaleByte`th element of the scale vectors,
|
||
while lanes 16-31 read from `firstScaleByte` + 2.
|
||
For example:
|
||
```mlir
|
||
// Input: 8-element vector of F8E5M2, converting to BF16
|
||
// Lanes 0-15 read from byte 0, lanes 16-31 read from byte 2 (0+2)
|
||
%result = amdgpu.scaled_ext_packed_matrix %source scale(%scales)
|
||
blockSize(16) firstScaleLane(0) firstScaleByte(0)
|
||
: vector<8xf8E5M2>, vector<4xf8E8M0FNU> -> vector<8xbf16>
|
||
|
||
// Input: 16-element vector of F6E3M2FN, converting to F32
|
||
// Lanes 0-15 read from byte 1, lanes 16-31 read from byte 3 (1+2)
|
||
%result = amdgpu.scaled_ext_packed_matrix %source scale(%scales)
|
||
blockSize(16) firstScaleLane(16) firstScaleByte(1)
|
||
: vector<16xf6E3M2FN>, vector<4xf8E8M0FNU> -> vector<16xf32>
|
||
```
|
||
|
||
Note: the layout for the scales generally mirrors how the WMMA
|
||
instructions use for matrix scales. These selection operands allows
|
||
one to choose portions of the matrix to convert.
|
||
|
||
When `source` is either F8E4M3FN or F8E5M2 and `blockSize` is 32,
|
||
then the same byte will be used by both halves of the wave.
|
||
In this case, `firstScaleByte` can be any value from 0 to 3.
|
||
|
||
When `source` is either F8E4M3FN or F8E5M2 and `blockSize` is 16,
|
||
following combinations are allowed:
|
||
* `firstScaleLane(0), firstScaleByte(0)`
|
||
* `firstScaleLane(16), firstScaleByte(2)`
|
||
all other combinations are reserved.
|
||
|
||
Available on gfx1250+.
|
||
}];
|
||
|
||
let assemblyFormat = [{
|
||
attr-dict $source
|
||
`scale` `(` $scale `)`
|
||
`blockSize` `(` $blockSize `)`
|
||
`firstScaleLane` `(` $firstScaleLane`)`
|
||
`firstScaleByte` `(` $firstScaleByte `)`
|
||
`:` type($source) `,` type($scale) `->` type($res)
|
||
}];
|
||
|
||
let hasVerifier = 1;
|
||
|
||
}
|
||
|
||
def AMDGPU_ScaledExtPackedOp
|
||
: AMDGPU_Op<"scaled_ext_packed", [Pure]>,
|
||
Arguments<(
|
||
ins AnyTypeOf<[VectorOfLengthAndType<[1, 2, 3, 4], [F8E5M2, F8E4M3FN]>,
|
||
VectorOfLengthAndType<[1, 2, 3, 4, 5, 6, 7, 8],
|
||
[F4E2M1FN]>]>:$source,
|
||
F32:$scale,
|
||
ConfinedAttr<I32Attr, [IntNonNegative, IntMaxValue<7>]>:$index)>,
|
||
Results<(
|
||
outs AnyTypeOf<[FixedVectorOfLengthAndType<[2], [F32]>,
|
||
FixedVectorOfLengthAndType<[2], [F16]>,
|
||
FixedVectorOfLengthAndType<[2], [BF16]>]>:$res)> {
|
||
let summary = "Extend a vector of packed floating point values";
|
||
|
||
let description = [{
|
||
Extend and scale two packed floats in `source[index]` to two floats and
|
||
return them.
|
||
|
||
This rather unusual signature arises from the fact that AMD GPUs cannot
|
||
easily work with sub 32-bit quantities, so the compiler intrinsics for
|
||
extending 8-bit floats (which are, currently, the only way to work with
|
||
this operation) take packed vectors of 2 such floats.
|
||
|
||
If the passed-in vector has fewer than two elements, or the input is scalar,
|
||
the remaining values in the <2 x i8> will be filled with
|
||
undefined values as needed.
|
||
}];
|
||
let assemblyFormat = [{
|
||
attr-dict $source `[` $index `]` `,` $scale `:` type($source) `to` type($res)
|
||
}];
|
||
}
|
||
|
||
def AMDGPU_PackedTrunc2xFp8Op :
|
||
AMDGPU_Op<"packed_trunc_2xfp8", [Pure, AttrSizedOperandSegments]>,
|
||
Arguments<(ins F32:$sourceA,
|
||
Optional<F32>:$sourceB,
|
||
ConfinedAttr<I32Attr, [IntNonNegative, IntMaxValue<1>]>:$wordIndex,
|
||
Optional<FixedVectorOfLengthAndType<[4], [F8E4M3FNUZ, F8E5M2FNUZ, F8E4M3FN, F8E5M2]>>:$existing)>,
|
||
Results<(outs FixedVectorOfLengthAndType<[4], [F8E4M3FNUZ, F8E5M2FNUZ, F8E4M3FN, F8E5M2]>:$res)> {
|
||
let summary = "Round two floats into a packed vector of 8-bit floats";
|
||
let description = [{
|
||
Round the inputs `sourceA` and `sourceB` (which is undefined if not
|
||
specified) into the low or high word (bottom two or top two) elements
|
||
of the returned vector, keeping the other two elements of `existing`
|
||
unchanged if present (or undefined if it was not passed in).
|
||
|
||
The reason for this odd signature is that AMD GPUs cannot easily work with
|
||
sub-registers, and so the conversion intrinsics (which are currently the
|
||
only way to work with 8-bit float types) take packed vectors of 4 8-bit
|
||
values.
|
||
}];
|
||
let assemblyFormat = [{
|
||
attr-dict $sourceA `,` ($sourceB^):(`undef`)?
|
||
`into` ($existing^):(`undef`)? `[` `word` $wordIndex `]`
|
||
`:` type($sourceA) `to` type($res) (`into` type($existing)^)?
|
||
}];
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
def AMDGPU_PackedScaledTruncOp
|
||
: AMDGPU_Op<"packed_scaled_trunc", [Pure]>,
|
||
Arguments<(ins VectorOfLengthAndType<[1, 2], [F32, F16, BF16]>:$source,
|
||
F32:$scale,
|
||
ConfinedAttr<I32Attr, [IntNonNegative, IntMaxValue<7>]>:$index,
|
||
Optional<AnyTypeOf<
|
||
[FixedVectorOfLengthAndType<[4], [F8E5M2, F8E4M3FN]>,
|
||
FixedVectorOfLengthAndType<[8], [F4E2M1FN]>]>>:$existing)>,
|
||
Results<(
|
||
outs AnyTypeOf<[FixedVectorOfLengthAndType<[4], [F8E5M2, F8E4M3FN]>,
|
||
FixedVectorOfLengthAndType<[8], [F4E2M1FN]>]>:$res)> {
|
||
let summary = "Round two floats into a packed vector of floats";
|
||
let description = [{
|
||
Scale and round the inputs `source` (which is undefined if not
|
||
specified) into the low or high word (bottom two or top two) elements
|
||
of the returned vector, keeping the other two elements of `existing`
|
||
unchanged if present (or undefined if it was not passed in).
|
||
|
||
The reason for this odd signature is that AMD GPUs cannot easily work with
|
||
sub-registers, and so the conversion intrinsics take 32-bit wide
|
||
packed vectors of float values.
|
||
}];
|
||
let assemblyFormat = [{
|
||
attr-dict $source `into` ($existing^):(`undef`)? `[` $index `]`
|
||
`,` $scale
|
||
`:` type($source) `to` type($res) (`into` type($existing)^)?
|
||
}];
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
def AMDGPU_PackedStochRoundFp8Op :
|
||
AMDGPU_Op<"packed_stoch_round_fp8", [Pure]>,
|
||
Arguments<(ins F32:$source,
|
||
I32:$stochiasticParam,
|
||
ConfinedAttr<I32Attr, [IntNonNegative, IntMaxValue<3>]>:$storeIndex,
|
||
Optional<FixedVectorOfLengthAndType<[4], [F8E4M3FNUZ, F8E5M2FNUZ, F8E4M3FN, F8E5M2]>>:$existing)>,
|
||
Results<(outs FixedVectorOfLengthAndType<[4], [F8E4M3FNUZ, F8E5M2FNUZ, F8E4M3FN, F8E5M2]>:$res)> {
|
||
let summary = "Round float stochiastically into a packed vector of 8-bit floats";
|
||
let description = [{
|
||
Round the input `source`, adding in `stochiasticParam`, and place it into
|
||
the `storeIndex`th element of `res`.
|
||
|
||
If `existing` is passed in, elements of `res` other than the one at `storeIndex`
|
||
are copied from `existing`.
|
||
|
||
The reason for this odd signature is that AMD GPUs cannot easily work with
|
||
sub-registers, and so the conversion intrinsics (which are currently the
|
||
only way to work with 8-bit float types) take packed vectors of 4 8-bit
|
||
values.
|
||
}];
|
||
let assemblyFormat = [{
|
||
attr-dict $source `+` $stochiasticParam
|
||
`into` ($existing^):(`undef`)? `[` $storeIndex `]`
|
||
`:` type($source) `to` type($res) (`into` type($existing)^)?
|
||
}];
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
def AMDGPU_FatRawBufferCastOp :
|
||
AMDGPU_Op<"fat_raw_buffer_cast",
|
||
[Pure,
|
||
DeclareOpInterfaceMethods<InferTypeOpInterface>,
|
||
DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface,
|
||
["reifyDimOfResult"]>,
|
||
ViewLikeOpInterface, AttrSizedOperandSegments]>,
|
||
Arguments<(ins AnyMemRef:$source,
|
||
Optional<I64>:$validBytes,
|
||
Optional<I<14>>:$cacheSwizzleStride,
|
||
DefaultValuedAttr<BoolAttr, "true">:$boundsCheck,
|
||
UnitAttr:$resetOffset)>,
|
||
Results<(outs AnyMemRef:$result)> {
|
||
// TODO: Set `resetOffset` and `boundsCheck` to use `Property` once
|
||
// we implemented pythonic binding for `Property`.
|
||
let summary = "Create a raw buffer fat pointer that matches `memref`";
|
||
let description = [{
|
||
Wraps the memory pointed to by `source` as a raw buffer fat pointer, or,
|
||
in LLVM terms, a `ptr addrspace(7)`, returning a memref that has the same
|
||
sizes and layout but the `#amdgpu.address_space<fat_raw_buffer>`
|
||
address space.
|
||
|
||
This memref can be used with standard memref operations like `memref.load`,
|
||
`memref.store`, and `memref.atomicrmw`, which will be lowered to the relevant
|
||
buffer intrinsics. (`vector.masked_load/store` will work once there's backend
|
||
support for lowering them, and then this document will be updated)
|
||
|
||
If `validBytes` is given, it is the number of bytes that will be valid as
|
||
an offset to `out`. If it is not provided, this will be inferred from
|
||
the size of the memref during lowering. This size is
|
||
max_{d = 0 upto rank(source)} (sizes[d] * strides[d]) * sizeof(element type).
|
||
|
||
The flags of the buffer descriptor will be set up to enable raw usage -
|
||
for example, stride = 0, add_tid = 0, and so on. The `boundsCheck`
|
||
property determines if bounds checking is enabled or not (on architectures
|
||
where this can be controlled - that is, on RDNA chips).
|
||
|
||
If `cacheSwizzleStride` is provided, L1 cache swizzling will be enabled
|
||
on architectures that support it. This swizzling, unlike the main swizzling
|
||
mode (whose usage makes a buffer non-raw) does not affect index calculation,
|
||
but does affect cache behavior. Mixing access between cache-swizzled raw
|
||
buffers and other forms of memory access, like ordinary pointer loads or
|
||
unswizzled buffer pointers can cause incorrect behavior and must be avoided.
|
||
|
||
This operation preserves the sizes, strides, and offset of the input
|
||
memref - they'll be added in by `memref.load` later. However, if
|
||
`resetOffset` is set, that offset will be added to the base pointer.
|
||
If the value of the memref's offset is not uniform (independent of the lane/thread ID),
|
||
this will lead to substantially decreased performance due to the need for
|
||
a waterfall loop on the base address of the buffer resource.
|
||
}];
|
||
|
||
let extraClassDeclaration = [{
|
||
Value getViewSource() { return getSource(); }
|
||
}];
|
||
|
||
let assemblyFormat = [{
|
||
$source oilist (`validBytes` `(` $validBytes `)`
|
||
| `cacheSwizzleStride` `(` $cacheSwizzleStride `)`
|
||
| `boundsCheck` `(` $boundsCheck `)`
|
||
| `resetOffset` $resetOffset )
|
||
attr-dict `:` type($source) `to` type($result)
|
||
}];
|
||
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
/// Raw buffer load
|
||
def AMDGPU_RawBufferLoadOp :
|
||
AMDGPU_Op<"raw_buffer_load", [AllElementTypesMatch<["value", "memref"]>,
|
||
AttrSizedOperandSegments]>,
|
||
Arguments<(ins Arg<AnyMemRef, "buffer to load from", [MemRead]>:$memref,
|
||
Variadic<I32>:$indices,
|
||
DefaultValuedAttr<BoolAttr, "true">:$boundsCheck,
|
||
OptionalAttr<I32Attr>:$indexOffset,
|
||
Optional<I32>:$sgprOffset)>,
|
||
Results<(outs AnyType:$value)> {
|
||
|
||
let summary = "Raw Buffer load, exposing GCN features";
|
||
let description = [{
|
||
The `amdgpu.raw_buffer_load` op is a wrapper around the buffer load intrinsics
|
||
available on AMD GPUs, including extensions in newer GPUs.
|
||
|
||
The index into the buffer is computed as for `memref.load` with the additon
|
||
of `indexOffset` and `sgprOffset` (which **may or may not** be considered
|
||
in bounds checks and includes any offset present on the memref type if it's
|
||
non-zero).
|
||
|
||
All indices and offsets are in units of the memref's data type and are
|
||
converted to bytes during lowering.
|
||
|
||
When a load is out of bounds, the instruction returns zero.
|
||
Partially-out of bounds have chipset-dependent behavior: whether reading
|
||
2 elements starting at index 7 of a `memref<8xf32>` returns the last element
|
||
in the first vector component depends on the architecture.
|
||
|
||
The memref struct is converted into a buffer resource (a V#) and the arguments
|
||
are translated to intrinsic arguments as follows:
|
||
- The base address of the buffer is the base address of the memref
|
||
- The stride is 0 to enable raw mode
|
||
- The number of records is the size of the memref, in bytes
|
||
In the case of dynamically-shaped memrefs, this is computed at runtime
|
||
as max_d (size(d) * stride(d)) * sizeof(elementType(memref))
|
||
- The offset enable bit is 1, the index enable bit is 0.
|
||
- The thread ID addition bit is off
|
||
- If `boundsCheck` is false and the target chipset is RDNA, OOB_SELECT is set
|
||
to 2 to disable bounds checks, otherwise it is 3
|
||
- The cache coherency bits are off
|
||
}];
|
||
let assemblyFormat = [{
|
||
attr-dict $memref `[` $indices `]`
|
||
(`sgprOffset` $sgprOffset^)? `:`
|
||
type($memref) (`,` type($indices)^)? `->` type($value)
|
||
}];
|
||
let hasCanonicalizer = 1;
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
/// Raw buffer store
|
||
def AMDGPU_RawBufferStoreOp :
|
||
AMDGPU_Op<"raw_buffer_store", [AllElementTypesMatch<["value", "memref"]>,
|
||
AttrSizedOperandSegments]>,
|
||
Arguments<(ins AnyType:$value,
|
||
Arg<AnyMemRef, "buffer to store to", [MemWrite]>:$memref,
|
||
Variadic<I32>:$indices,
|
||
DefaultValuedAttr<BoolAttr, "true">:$boundsCheck,
|
||
OptionalAttr<I32Attr>:$indexOffset,
|
||
Optional<I32>:$sgprOffset)> {
|
||
|
||
let summary = "Raw Buffer Store, exposing GCN features";
|
||
let description = [{
|
||
The `amdgpu.raw_buffer_store` op is a wrapper around the buffer store
|
||
intrinsics available on AMD GPUs, including extensions in newer GPUs.
|
||
|
||
The store index is computed as in `memref.store` with the addition of
|
||
`indexOffset` (which is included for uniformity with atomics and may be useful
|
||
when writing vectorized code) and `sgprOffset` (which is added after bounds
|
||
checks and implicitly includes the offset of the memref type if non-zero).
|
||
All index components are in terms of the elements of the memref, not bytes,
|
||
and are scaled up appropriately.
|
||
|
||
Out of bounds stores are ignored in hardware.
|
||
Wthether a vector write that includes some in-bounds and soeme out-of-bounds
|
||
components is partically completed is chipset-dependent.
|
||
|
||
See `amdgpu.raw_buffer_load` for a description of how the underlying
|
||
instruction is constructed.
|
||
}];
|
||
let assemblyFormat = [{
|
||
attr-dict $value `->` $memref `[` $indices `]`
|
||
(`sgprOffset` $sgprOffset^)? `:`
|
||
type($value) `->` type($memref) (`,` type($indices)^)?
|
||
}];
|
||
let hasCanonicalizer = 1;
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
// Raw buffer atomic compare-and-swap
|
||
def AMDGPU_RawBufferAtomicCmpswapOp :
|
||
AMDGPU_Op<"raw_buffer_atomic_cmpswap", [
|
||
AttrSizedOperandSegments,
|
||
AllTypesMatch<["src", "cmp", "value"]>,
|
||
AllElementTypesMatch<["value", "memref"]>]>,
|
||
Arguments<(ins AnyType:$src,
|
||
AnyType:$cmp,
|
||
Arg<AnyMemRef, "buffer to operate on", [MemRead, MemWrite]>:$memref,
|
||
Variadic<I32>:$indices,
|
||
DefaultValuedAttr<BoolAttr, "true">:$boundsCheck,
|
||
OptionalAttr<I32Attr>:$indexOffset,
|
||
Optional<I32>:$sgprOffset)>,
|
||
Results<(outs AnyType:$value)> {
|
||
|
||
let summary = "Raw Buffer Atomic compare-and-swap";
|
||
let description = [{
|
||
The `amdgpu.raw_buffer_atomic_cmpswap` op is a wrapper around the
|
||
buffer-based atomic compare-and-swap min available on AMD GPUs.
|
||
|
||
The index into the buffer is computed as for `memref.store` with the addition
|
||
of `indexOffset` (which is used to aid in emitting vectorized code) and,
|
||
if present `sgprOffset` (which is added after bounds checks and includes
|
||
any non-zero offset on the memref type).
|
||
|
||
All indexing components are given in terms of the memref's element size, not
|
||
the byte lengths required by the intrinsic.
|
||
|
||
Out of bounds atomic operations are ignored in hardware.
|
||
|
||
See `amdgpu.raw_buffer_load` for a description of how the underlying
|
||
instruction is constructed.
|
||
}];
|
||
let assemblyFormat = [{
|
||
attr-dict $src `,` $cmp `->` $memref `[` $indices `]`
|
||
(`sgprOffset` $sgprOffset^)? `:`
|
||
type($value) `->` type($memref) `,` type($indices)
|
||
}];
|
||
let hasCanonicalizer = 1;
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
// Raw buffer atomic floating point add
|
||
def AMDGPU_RawBufferAtomicFaddOp :
|
||
AMDGPU_Op<"raw_buffer_atomic_fadd", [
|
||
AllElementTypesMatch<["value", "memref"]>,
|
||
AllTypesMatch<["oldValue", "value"]>,
|
||
AttrSizedOperandSegments]>,
|
||
Arguments<(ins AnyTypeOf<[F32, VectorOfLengthAndType<[2], [F16, BF16]>]>:$value,
|
||
Arg<AnyMemRef, "buffer to operate on", [MemRead, MemWrite]>:$memref,
|
||
Variadic<I32>:$indices,
|
||
DefaultValuedAttr<BoolAttr, "true">:$boundsCheck,
|
||
OptionalAttr<I32Attr>:$indexOffset,
|
||
Optional<I32>:$sgprOffset)>,
|
||
Results<(outs AnyTypeOf<[F32, VectorOfLengthAndType<[2], [F16, BF16]>]>:$oldValue)> {
|
||
|
||
let summary = "Raw Buffer Floating-point Atomic Add (MI-* only)";
|
||
let description = [{
|
||
The `amdgpu.raw_buffer_atomic_fadd` op is a wrapper around the
|
||
buffer-based atomic floating point addition available on the MI-* series
|
||
of AMD GPUs.
|
||
|
||
The index into the buffer is computed as for `memref.store` with the addition
|
||
of `indexOffset` (which is used to aid in emitting vectorized code) and,
|
||
if present `sgprOffset` (which is added after bounds checks and includes
|
||
any non-zero offset on the memref type).
|
||
|
||
All indexing components are given in terms of the memref's element size, not
|
||
the byte lengths required by the intrinsic.
|
||
|
||
The op returns the value observed before the atomic update.
|
||
|
||
See `amdgpu.raw_buffer_load` for a description of how the underlying
|
||
instruction is constructed.
|
||
}];
|
||
let assemblyFormat = [{
|
||
attr-dict $value `->` $memref `[` $indices `]`
|
||
(`sgprOffset` $sgprOffset^)? `:`
|
||
type($value) `->` type($memref) `,` type($indices)
|
||
}];
|
||
let hasCanonicalizer = 1;
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
// Raw buffer atomic floating point max
|
||
def AMDGPU_RawBufferAtomicFmaxOp :
|
||
AMDGPU_Op<"raw_buffer_atomic_fmax", [
|
||
AllElementTypesMatch<["value", "memref"]>,
|
||
AllTypesMatch<["oldValue", "value"]>,
|
||
AttrSizedOperandSegments]>,
|
||
Arguments<(ins AnyTypeOf<[F32, F64]>:$value,
|
||
Arg<AnyMemRef, "buffer to operate on", [MemRead, MemWrite]>:$memref,
|
||
Variadic<I32>:$indices,
|
||
DefaultValuedAttr<BoolAttr, "true">:$boundsCheck,
|
||
OptionalAttr<I32Attr>:$indexOffset,
|
||
Optional<I32>:$sgprOffset)>,
|
||
Results<(outs AnyTypeOf<[F32, F64]>:$oldValue)> {
|
||
|
||
let summary = "Raw Buffer Floating-point Atomic Max (non-GFX9)";
|
||
let description = [{
|
||
The `amdgpu.raw_buffer_atomic_fmax` op is a wrapper around the
|
||
buffer-based atomic floating point max available on AMD GPUs (except GFX9).
|
||
|
||
The index into the buffer is computed as for `memref.store` with the addition
|
||
of `indexOffset` (which is used to aid in emitting vectorized code) and,
|
||
if present `sgprOffset` (which is added after bounds checks and includes
|
||
any non-zero offset on the memref type).
|
||
|
||
All indexing components are given in terms of the memref's element size, not
|
||
the byte lengths required by the intrinsic.
|
||
|
||
The op returns the value observed before the atomic update. Out of bounds
|
||
atomic operations are ignored in hardware.
|
||
|
||
See `amdgpu.raw_buffer_load` for a description of how the underlying
|
||
instruction is constructed.
|
||
}];
|
||
let assemblyFormat = [{
|
||
attr-dict $value `->` $memref `[` $indices `]`
|
||
(`sgprOffset` $sgprOffset^)? `:`
|
||
type($value) `->` type($memref) `,` type($indices)
|
||
}];
|
||
let hasCanonicalizer = 1;
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
// Raw buffer atomic signed integer max
|
||
def AMDGPU_RawBufferAtomicSmaxOp :
|
||
AMDGPU_Op<"raw_buffer_atomic_smax", [
|
||
AllTypesMatch<["oldValue", "value"]>,
|
||
AttrSizedOperandSegments]>,
|
||
Arguments<(ins I32:$value,
|
||
Arg<AnyMemRef, "buffer to operate on", [MemRead, MemWrite]>:$memref,
|
||
Variadic<I32>:$indices,
|
||
DefaultValuedAttr<BoolAttr, "true">:$boundsCheck,
|
||
OptionalAttr<I32Attr>:$indexOffset,
|
||
Optional<I32>:$sgprOffset)>,
|
||
Results<(outs I32:$oldValue)> {
|
||
|
||
let summary = "Raw Buffer Signed Integer Atomic Max";
|
||
let description = [{
|
||
The `amdgpu.raw_buffer_atomic_smax` op is a wrapper around the
|
||
buffer-based atomic signed integer max available on AMD GPUs.
|
||
|
||
The index into the buffer is computed as for `memref.store` with the addition
|
||
of `indexOffset` (which is used to aid in emitting vectorized code) and,
|
||
if present `sgprOffset` (which is added after bounds checks and includes
|
||
any non-zero offset on the memref type).
|
||
|
||
All indexing components are given in terms of the memref's element size, not
|
||
the byte lengths required by the intrinsic.
|
||
|
||
The op returns the value observed before the atomic update. Out of bounds
|
||
atomic operations are ignored in hardware.
|
||
|
||
See `amdgpu.raw_buffer_load` for a description of how the underlying
|
||
instruction is constructed.
|
||
}];
|
||
let assemblyFormat = [{
|
||
attr-dict $value `->` $memref `[` $indices `]`
|
||
(`sgprOffset` $sgprOffset^)? `:`
|
||
type($value) `->` type($memref) `,` type($indices)
|
||
}];
|
||
let hasCanonicalizer = 1;
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
// Raw buffer atomic unsigned integer min
|
||
def AMDGPU_RawBufferAtomicUminOp :
|
||
AMDGPU_Op<"raw_buffer_atomic_umin", [
|
||
AllTypesMatch<["oldValue", "value"]>,
|
||
AttrSizedOperandSegments]>,
|
||
Arguments<(ins I32:$value,
|
||
Arg<AnyMemRef, "buffer to operate on", [MemRead, MemWrite]>:$memref,
|
||
Variadic<I32>:$indices,
|
||
DefaultValuedAttr<BoolAttr, "true">:$boundsCheck,
|
||
OptionalAttr<I32Attr>:$indexOffset,
|
||
Optional<I32>:$sgprOffset)>,
|
||
Results<(outs I32:$oldValue)> {
|
||
|
||
let summary = "Raw Buffer Unsigned Integer Atomic Min";
|
||
let description = [{
|
||
The `amdgpu.raw_buffer_atomic_umin` op is a wrapper around the
|
||
buffer-based atomic signed integer min available on AMD GPUs.
|
||
|
||
The index into the buffer is computed as for `memref.store` with the addition
|
||
of `indexOffset` (which is used to aid in emitting vectorized code) and,
|
||
if present `sgprOffset` (which is added after bounds checks and includes
|
||
any non-zero offset on the memref type).
|
||
|
||
All indexing components are given in terms of the memref's element size, not
|
||
the byte lengths required by the intrinsic.
|
||
|
||
The op returns the value observed before the atomic update. Out of bounds
|
||
atomic operations are ignored in hardware.
|
||
|
||
See `amdgpu.raw_buffer_load` for a description of how the underlying
|
||
instruction is constructed.
|
||
}];
|
||
let assemblyFormat = [{
|
||
attr-dict $value `->` $memref `[` $indices `]`
|
||
(`sgprOffset` $sgprOffset^)? `:`
|
||
type($value) `->` type($memref) `,` type($indices)
|
||
}];
|
||
let hasCanonicalizer = 1;
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
|
||
def AMDGPU_DPPOp : AMDGPU_Op<"dpp",
|
||
[Pure, SameTypeOperands, AllTypesMatch<["result", "old", "src"]>]>,
|
||
Arguments<(ins AMDGPU_IntOrFloatOr1DVectorWidthLeq64:$old,
|
||
AMDGPU_IntOrFloatOr1DVectorWidthLeq64:$src,
|
||
AMDGPU_DPPPermAttr:$kind,
|
||
OptionalAttr<AnyAttrOf<[I32Attr, ArrayAttr, UnitAttr]>>:$permArgument,
|
||
DefaultValuedAttr<I32Attr, "0xf">:$row_mask,
|
||
DefaultValuedAttr<I32Attr, "0xf">:$bank_mask,
|
||
DefaultValuedAttr<BoolAttr, "false">:$bound_ctrl)> {
|
||
let summary = "AMDGPU DPP operation";
|
||
let description = [{
|
||
The `amdgpu.dpp` op performs a Data Parallel Primitives (DPP) lane
|
||
permutation on a source value within a wavefront. Each lane reads its
|
||
source data from another lane according to the permutation mode specified
|
||
by `kind`. DPP operates at dword (32-bit) granularity: sub-32-bit types
|
||
(e.g., f16, i16) are packed into an i32 during lowering, permuted, and
|
||
extracted back.
|
||
|
||
- Lanes are organized into rows of 16.
|
||
- A Wave64 wavefront has 4 rows of 16 lanes each: row 0 = lanes 0-15,
|
||
row 1 = lanes 16-31, row 2 = lanes 32-47, row 3 = lanes 48-63.
|
||
- Similarly, a Wave32 wavefront has two rows of 16 lanes each, organized
|
||
in the same fashion.
|
||
- Each row is divided into 4 banks of 4 consecutive lanes: bank 0 =
|
||
lanes 0-3, bank 1 = lanes 4-7, bank 2 = lanes 8-11, bank 3 =
|
||
lanes 12-15 (lane numbers shown for row 0; add 16/32/48 for other rows).
|
||
|
||
The `kind` attribute selects the permutation. Some modes require a
|
||
`permArgument`; others take no argument.
|
||
|
||
Quad permutation:
|
||
- `quad_perm([a, b, c, d])`: Full permute within each group of 4
|
||
consecutive lanes (a quad). Each element is in [0, 3] and selects which
|
||
lane within the quad to read from. Lane 4k+i reads from lane 4k+perm[i].
|
||
For example, `quad_perm([1, 0, 3, 2])` swaps adjacent pairs within
|
||
every quad.
|
||
|
||
Row shifts and rotates (operate within each 16-lane row independently):
|
||
- `row_shl(N)`: Shift left by N (1-15) within the row. Lane n reads from
|
||
lane (n % 16) + N in the same row. Lanes where the source index exceeds
|
||
15 are out of bounds (see `bound_ctrl`).
|
||
- `row_shr(N)`: Shift right by N (1-15) within the row. Lane n reads from
|
||
lane (n % 16) - N in the same row. Lanes where the source index is
|
||
negative are out of bounds.
|
||
- `row_ror(N)`: Rotate right by N (1-15) within the row. Lane n reads from
|
||
lane ((n % 16) - N) mod 16 in the same row. Always in bounds.
|
||
|
||
Wavefront shifts and rotates (not available on RDNA):
|
||
- `wave_shl`: Shift left by 1. Lane n reads from lane n + 1. The last lane
|
||
in the wavefront is out of bounds.
|
||
- `wave_shr`: Shift right by 1. Lane n reads from lane n - 1. Lane 0 is
|
||
out of bounds.
|
||
- `wave_rol`: Rotate left by 1. Lane n reads from lane (n + 1) mod W, where
|
||
W is the wavefront size.
|
||
- `wave_ror`: Rotate right by 1. Lane n reads from lane (n - 1) mod W, where
|
||
W is the wavefront size.
|
||
|
||
Row mirrors:
|
||
- `row_mirror`: Reverse lanes within each 16-lane row. Lane n reads from
|
||
lane 15 - (n % 16) within its row.
|
||
- `row_half_mirror`: Reverse within each 8-lane half-row. Lane n reads
|
||
from lane 7 - (n % 8) within its half-row.
|
||
|
||
Row broadcasts (not available on RDNA):
|
||
- `row_bcast_15`: Lane 15 of each row broadcasts to all lanes of the next
|
||
row. Lanes in row 0 are not affected (retain `old`).
|
||
- `row_bcast_31`: Lane 31 broadcasts to all lanes in rows 2 and 3.
|
||
Lanes in rows 0 and 1 are not affected (retain `old`).
|
||
|
||
Example:
|
||
```mlir
|
||
// Swap adjacent pairs within each quad (lanes 0<->1, 2<->3, etc.)
|
||
%0 = amdgpu.dpp %old %src quad_perm( [1, 0, 3, 2] ) : i32
|
||
|
||
// Shift right by 1 lane within each 16-lane row.
|
||
// bound_ctrl=true -> lanes that would read past the row return 0.
|
||
// row_mask=0x5 (0b0101) -> only rows 0 and 2 apply the shift;
|
||
// rows 1 and 3 pass through %old unchanged.
|
||
%1 = amdgpu.dpp %old %src row_shr( 0x1 : i32 )
|
||
{ row_mask = 0x5 : i32, bound_ctrl = true } : f32
|
||
|
||
// Rotate left across the full wavefront by 1 lane
|
||
%2 = amdgpu.dpp %old %src wave_rol : i32
|
||
```
|
||
|
||
Operands:
|
||
* `$old`: Fallback value. Lanes that are masked off by `row_mask` /
|
||
`bank_mask` retain `old`. For lanes with an out-of-bounds source, behavior
|
||
depends on `bound_ctrl`.
|
||
* `$src`: Source value to be permuted across lanes.
|
||
* `$kind`: A `#amdgpu.dpp_perm` enum selecting the permutation mode.
|
||
* `$permArgument`: Mode-specific argument. Required for `quad_perm`
|
||
(array of 4 integers in [0, 3]) and `row_shl`/`row_shr`/`row_ror`
|
||
(integer in [1, 15]). Absent for all other modes.
|
||
* `$row_mask` (default 0xf): 4-bit mask controlling which rows write
|
||
results. Bit i enables row i (bit 0 = lanes 0-15, bit 1 = lanes
|
||
16-31, etc.). Disabled lanes retain `old`.
|
||
* `$bank_mask` (default 0xf): 4-bit mask controlling which banks write
|
||
results. Bit i enables bank i (bit 0 = lanes 0-3, 16-19, etc. across all rows).
|
||
Disabled lanes retain `old`.
|
||
* `$bound_ctrl` (default false): When false, out of bounds lanes retain
|
||
`old`. When true, out-of-bounds lanes receive zero.
|
||
}];
|
||
let results = (outs AnyType:$result);
|
||
let assemblyFormat = [{
|
||
$old $src $kind (`(` $permArgument^ `)`)? attr-dict `:` type($result)
|
||
}];
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
def AMDGPU_SwizzleBitModeOp : AMDGPU_Op<"swizzle_bitmode",
|
||
[Pure, AllTypesMatch<["result", "src"]>]>,
|
||
Arguments<(ins AnyIntegerOrFloatOr1DVector:$src,
|
||
I32Attr:$and_mask,
|
||
I32Attr:$or_mask,
|
||
I32Attr:$xor_mask
|
||
)> {
|
||
let summary = "AMDGPU ds_swizzle op, bitmode variant";
|
||
let description = [{
|
||
High-level wrapper on bitmode `rocdl.ds_swizzle` op, masks are represented
|
||
as separate fields so user won't need to do manual bitpacking.
|
||
|
||
Supports arbitrary int/float/vector types, which will be repacked to i32 and
|
||
one or more `rocdl.ds_swizzle` ops during lowering.
|
||
}];
|
||
let results = (outs AnyIntegerOrFloatOr1DVector:$result);
|
||
let assemblyFormat = [{
|
||
$src $and_mask $or_mask $xor_mask attr-dict `:` type($result)
|
||
}];
|
||
}
|
||
|
||
def AMDGPU_PermlaneSwapOp : AMDGPU_Op<"permlane_swap", [Pure, AllTypesMatch<["result", "src"]>]> {
|
||
let summary = "AMDGPU permlane swap op";
|
||
let description = [{
|
||
High-level wrapper on `rocdl.permlane{16,32}.swap` variants for permutations
|
||
on rows of lanes in a subgroup.
|
||
|
||
Supports arbitrary int/float/vector types, which will be repacked to i32 and
|
||
one or more `rocdl.permlane_swap` ops during lowering.
|
||
Supported lane permutations:
|
||
- Swap the data between odd and even rows of 16 lanes
|
||
- Swap the data between the first 32 lanes and the last 32 lanes
|
||
|
||
Example:
|
||
```mlir
|
||
%0 = amdgpu.permlane_swap %src 16 : f16
|
||
%1 = amdgpu.permlane_swap %src 32 { fetch_inactive = true, bound_ctrl = true } : f16
|
||
```
|
||
|
||
Operands:
|
||
* `$src`: Vector register to permute across lanes of the subgroup.
|
||
* `$row_length`: The length of a row to permute in number of lanes (valid values are 16 and 32).
|
||
* `$fetch_inactive`: Optional. Used to dertermine behavior of a fetch from a disabled lane.
|
||
`fetch_inactive = false`: If the source lane is disabled, use `bound_ctrl` to determine the source value.
|
||
`fetch_inactive = true`: If the source lane is disabled, fetch the source value anyway (ignoring `bound_ctrl`).
|
||
* `$bound_ctrl`: Optional. Used to determine what a thread should do if its source operand is from
|
||
a disabled lane: use the value zero, or disable the write.
|
||
`bound_ctrl = false`: Do not write when source is from a disabled lane
|
||
`bound_ctrl = true`: Use zero as input if source is from a disabled lane
|
||
|
||
Note: Lowering is only supported on gfx950 and up.
|
||
}];
|
||
let arguments = (ins AnyIntegerOrFloatOr1DVector:$src,
|
||
I32Attr:$row_length,
|
||
DefaultValuedAttr<BoolAttr, "false">:$fetch_inactive,
|
||
DefaultValuedAttr<BoolAttr, "false">:$bound_ctrl);
|
||
let results = (outs AnyIntegerOrFloatOr1DVector:$result);
|
||
let assemblyFormat = [{
|
||
$src $row_length attr-dict `:` type($result)
|
||
}];
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
def AMDGPU_PermlaneVarOp : AMDGPU_Op<"permlane_var",
|
||
[Pure, AllTypesMatch<["result", "src"]>]> {
|
||
let summary = "AMDGPU variable-selector permlane op (GFX12+)";
|
||
let description = [{
|
||
High-level wrapper on `rocdl.permlane16.var` and `rocdl.permlanex16.var`
|
||
for per-lane variable-selector permutations within a wave32 subgroup.
|
||
|
||
Supports arbitrary int/float/vector types, which will be repacked to i32
|
||
and one or more ROCDL intrinsic calls during lowering.
|
||
|
||
- `cross = false`: intra-row permutation (each lane in a 16-lane row reads
|
||
from a lane in the same row, selected by `$selector`). Maps to
|
||
`rocdl.permlane16.var`.
|
||
- `cross = true`: cross-row permutation (each lane reads from the opposite
|
||
16-lane row, selected by `$selector`). Maps to `rocdl.permlanex16.var`.
|
||
|
||
`$selector` is an i32 VGPR providing the per-lane source-lane index.
|
||
|
||
Example:
|
||
```mlir
|
||
%0 = amdgpu.permlane_var %src, %sel { cross = false } : f16
|
||
%1 = amdgpu.permlane_var %src, %sel { cross = true } : f32
|
||
```
|
||
|
||
Note: Lowering is only supported on GFX12+.
|
||
}];
|
||
let arguments = (ins AnyIntegerOrFloatOr1DVector:$src,
|
||
I32:$selector,
|
||
DefaultValuedAttr<BoolAttr, "false">:$cross,
|
||
DefaultValuedAttr<BoolAttr, "false">:$fetch_inactive,
|
||
DefaultValuedAttr<BoolAttr, "false">:$bound_ctrl);
|
||
let results = (outs AnyIntegerOrFloatOr1DVector:$result);
|
||
let assemblyFormat = [{
|
||
$src `,` $selector attr-dict `:` type($result)
|
||
}];
|
||
}
|
||
|
||
def AMDGPU_LDSBarrierOp : AMDGPU_Op<"lds_barrier"> {
|
||
let summary = "Barrier that includes a wait for LDS memory operations.";
|
||
let description = [{
|
||
**DEPRECATION NOTICE**: Unless you need the inline-assembly-based workaround
|
||
for gfx908/MI-100, you should represent this pattern with the equivalent
|
||
|
||
```mlir
|
||
gpu.barrier memfence [#gpu.address_space<workgroup>]
|
||
```
|
||
|
||
instead.
|
||
|
||
`amdgpu.lds_barrier` is both a barrier (all workitems in a workgroup must reach
|
||
the barrier before any of them may proceed past it) and a wait for all
|
||
operations that affect the Local Data Store (LDS) issued from that workgroup
|
||
to complete before the workgroup may continue. Since the LDS is per-workgroup
|
||
memory, this barrier may be used, for example, to ensure all workitems have
|
||
written data to LDS before any workitem attempts to read from it.
|
||
|
||
Note that `lds_barrier` does **not** force reads to or from global memory
|
||
to complete before execution continues. Therefore, it should be used when
|
||
operations on global memory can be issued far in advance of when their results
|
||
are used (for example, by writing them to LDS).
|
||
|
||
WARNING: On architectures that do not support the BackOffBarrier feature,
|
||
(those which will implement this barrier by emitting inline assembly),
|
||
use of this operation will impede the usabiliity of memory watches (including
|
||
breakpoints set on variables) when debugging.
|
||
}];
|
||
let assemblyFormat = "attr-dict";
|
||
let hasCanonicalizer = 1;
|
||
}
|
||
|
||
def AMDGPU_SchedBarrierOp :
|
||
AMDGPU_Op<"sched_barrier">,
|
||
Arguments<(ins ROCDL_SchedGroupMaskAttr:$opts)>
|
||
{
|
||
let summary = "Barrier that limits the backend scheduler of instruction movement";
|
||
let description = [{
|
||
`amdgpu.sched_barrier` serves as a barrier that could be
|
||
configured to restrict movements of instructions through it as
|
||
defined by the ROCDL scheduling group mask enum.
|
||
}];
|
||
let assemblyFormat = [{
|
||
`allow` `=` $opts attr-dict
|
||
}];
|
||
}
|
||
|
||
def AMDGPU_MemoryCounterWaitOp :
|
||
AMDGPU_Op<"memory_counter_wait">,
|
||
Arguments<(ins
|
||
OptionalAttr<I32Attr>:$load,
|
||
OptionalAttr<I32Attr>:$store,
|
||
OptionalAttr<I32Attr>:$ds,
|
||
OptionalAttr<I32Attr>:$exp,
|
||
OptionalAttr<I32Attr>:$tensor
|
||
)>
|
||
{
|
||
let summary = "Wait for specified hardware counters";
|
||
let description = [{
|
||
Wait for the specified counters to be less-than or equal-to the provided
|
||
values before continuing.
|
||
|
||
Counters can lower to different instructions on different architectires,
|
||
including clamping to the some HW supported max value or combining multiple
|
||
counters into one.
|
||
}];
|
||
let assemblyFormat = [{
|
||
oilist( `load` `(` $load `)` | `store` `(` $store `)` | `ds` `(` $ds `)` | `exp` `(` $exp `)` | `tensor` `(` $tensor `)` ) attr-dict
|
||
}];
|
||
|
||
let hasCanonicalizer = 1;
|
||
}
|
||
|
||
|
||
// mfma
|
||
def MFMAInTypes : AnyTypeOf<[F32, F64, I32, I64,
|
||
VectorOfLengthAndType<[2], [F32]>,
|
||
VectorOfLengthAndType<[4, 8], [F16]>,
|
||
VectorOfLengthAndType<[2, 4, 8], [BF16]>,
|
||
VectorOfLengthAndType<[4, 8, 16], [I8]>,
|
||
VectorOfLengthAndType<[8], [F8E5M2FNUZ, F8E4M3FNUZ]>,
|
||
VectorOfLengthAndType<[8, 32], [F8E5M2, F8E4M3FN]>,
|
||
VectorOfLengthAndType<[32], [F6E2M3FN, F6E3M2FN, F4E2M1FN]>]>;
|
||
def MFMAOutTypes : AnyTypeOf<[F64,
|
||
VectorOfLengthAndType<[4, 16, 32], [F32]>,
|
||
VectorOfLengthAndType<[4, 16, 32], [I32]>,
|
||
VectorOfLengthAndType<[4], [F64]>]>;
|
||
|
||
// sparse_mfma (smfmac)
|
||
def SMFMACSparseInTypes : AnyTypeOf<[
|
||
VectorOfLengthAndType<[4, 8], [F16]>,
|
||
VectorOfLengthAndType<[4, 8], [BF16]>,
|
||
VectorOfLengthAndType<[8, 16], [I8]>,
|
||
VectorOfLengthAndType<[8, 16], [F8E4M3FN, F8E5M2]>,
|
||
VectorOfLengthAndType<[8, 16], [F8E4M3FNUZ, F8E5M2FNUZ]>
|
||
]>;
|
||
|
||
def SMFMACDenseInTypes : AnyTypeOf<[
|
||
VectorOfLengthAndType<[8, 16], [F16]>,
|
||
VectorOfLengthAndType<[8, 16], [BF16]>,
|
||
VectorOfLengthAndType<[16, 32], [I8]>,
|
||
VectorOfLengthAndType<[16, 32], [F8E4M3FN, F8E5M2]>,
|
||
VectorOfLengthAndType<[16, 32], [F8E4M3FNUZ, F8E5M2FNUZ]>
|
||
]>;
|
||
|
||
def SMFMACOutTypes : AnyTypeOf<[
|
||
VectorOfLengthAndType<[4, 16], [F32]>,
|
||
VectorOfLengthAndType<[4, 16], [I32]>
|
||
]>;
|
||
|
||
def SMFMACIdxTypes : AnyTypeOf<[
|
||
FixedVectorOfLengthAndType<[4], [I8]>,
|
||
FixedVectorOfLengthAndType<[2], [I16]>,
|
||
I32
|
||
]>;
|
||
|
||
// scaled_mfma
|
||
def ScaledMFMAInTypes : AnyTypeOf<[VectorOfLengthAndType<[32], [F8E5M2, F8E4M3FN]>,
|
||
VectorOfLengthAndType<[32], [F6E2M3FN, F6E3M2FN, F4E2M1FN]>]>;
|
||
def ScaledMFMAOutTypes : AnyTypeOf<[VectorOfLengthAndType<[4, 16], [F32]>]>;
|
||
|
||
// scaled_wmma
|
||
def ScaledWMMAInTypes
|
||
: AnyTypeOf<[VectorOfLengthAndType<[64], [F8E5M2, F8E4M3FN]>,
|
||
VectorOfLengthAndType<[64], [F6E2M3FN, F6E3M2FN]>,
|
||
VectorOfLengthAndType<[64, 128], [F4E2M1FN]>]>;
|
||
|
||
def ScaledWMMAOutTypes : AnyTypeOf<[VectorOfLengthAndType<[8, 16], [F32]>]>;
|
||
|
||
// wmma
|
||
def WMMAInTypes : AnyTypeOf<[VectorOfLengthAndType<[2], [F32]>,
|
||
VectorOfLengthAndType<[4, 8, 16], [F16, BF16]>,
|
||
VectorOfLengthAndType<[4, 8, 16, 32], [I8, SI8, UI8]>,
|
||
VectorOfLengthAndType<[4, 8, 32, 64], [F8E4M3FN, F8E5M2]>,
|
||
VectorOfLengthAndType<[4, 8, 16], [I<4>, SI<4>, UI<4>]>]>;
|
||
def WMMAOutTypes : AnyTypeOf<[VectorOfLengthAndType<[4, 8], [F32, I32]>,
|
||
VectorOfLengthAndType<[4, 8, 16], [F16, BF16]>]>;
|
||
|
||
def AMDGPU_MFMAOp :
|
||
AMDGPU_Op<"mfma", [AllTypesMatch<["destC", "destD"]>,
|
||
Pure]>,
|
||
Arguments<(ins
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[4, 16, 32]>]>:$m,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[4, 16, 32]>]>:$n,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[1, 2, 4, 8, 16, 32, 64, 128]>]>:$k,
|
||
DefaultValuedAttr<ConfinedAttr<I32Attr, [IntIsOneOf<[1, 2, 4, 16]>]>, "1">:$blocks,
|
||
MFMAInTypes:$sourceA,
|
||
MFMAInTypes:$sourceB,
|
||
MFMAOutTypes:$destC,
|
||
DefaultValuedAttr<I32Attr, "0">:$cbsz,
|
||
DefaultValuedAttr<I32Attr, "0">:$abid,
|
||
DefaultValuedAttr<ROCDL_MFMAPermBAttr,
|
||
"::mlir::ROCDL::MFMAPermB::none">:$blgp,
|
||
UnitAttr:$reducePrecision,
|
||
UnitAttr:$negateA,
|
||
UnitAttr:$negateB,
|
||
UnitAttr:$negateC)>,
|
||
Results<(outs MFMAOutTypes: $destD)> {
|
||
let summary = "MLIR wrapper for CDNA mfma instructions";
|
||
let description = [{
|
||
The `amdgpu.mfma` op is an MLIR wrapper around intrinsics
|
||
for various `mfma` instructions in the CDNA architecture, which perform
|
||
multiple outer products in order to allow fast matrix multiplication.
|
||
|
||
The wrapper will select an appropriate `mfma` instruction, if one is available,
|
||
based on the provided `m`, `k`, `n`, and `nBlks` attributes, along with the
|
||
types of the source and destination arguments.
|
||
|
||
For information on the layouts of the input and output matrices (which are stored
|
||
in `sourceA`, `sourceB`, `destC`, and `destD`), see the CDNA ISA documentation.
|
||
|
||
The `cbsz`, `abid`, and `blgp` parameters control how the lanes of the wave
|
||
are permuted when matrix data is being loaded: `blgp` can be any number of
|
||
fixed permutations, `cbsz` specifies the log_2 of the number of chunks the lanes
|
||
holding sourceA are split into, and `abid` selects one of those chunks.
|
||
|
||
Note, this wrapper allows specifying `vector<4Kxi8>` arguments to MFMA
|
||
intrinsics that take an integer type of width `4K`. For example,
|
||
one can provide a vector<4xi8> as an argument to an MFMA instruction that
|
||
logically takes 4 i8s but whose intrinsics are specified to take an i32.
|
||
In these cases, the bytes in the vector will be concatenated in little-endian
|
||
order (that is, v[0] will go to arg[7:0], v[1] to arg[15:8] and so on).
|
||
|
||
The negateA, negateB, and negateC flags are only supported for double-precision
|
||
operations on gfx94x.
|
||
|
||
Example:
|
||
```mlir
|
||
%0 = amdgpu.mfma 16x16x16 %matA * %matB + %matC
|
||
: vector<4xf16>, vector<4xf16>, vector<4xf32>
|
||
|
||
%1 = amdgpu.mfma 32x32x1 %matD * %matE + %matF
|
||
{ abid = 1 : i32, cbsz = 1 : i32, blocks = 2 : i32 }
|
||
blgp = bcast_second_32 : f32, f32, vector<32xf32>
|
||
```
|
||
}];
|
||
let assemblyFormat = [{
|
||
custom<MNKDimensionList>($m, $n, $k) $sourceA `*` $sourceB `+` $destC
|
||
attr-dict
|
||
`blgp` `=` $blgp
|
||
`:` type($sourceA) `,` type($sourceB) `,` type($destC)
|
||
}];
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
def AMDGPU_WMMAOp :
|
||
AMDGPU_Op<"wmma", [AllTypesMatch<["destC", "destD"]>,
|
||
Pure]>,
|
||
Arguments<(ins
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[16]>]>:$m,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[16]>]>:$n,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[4, 16, 32, 64, 128]>]>:$k,
|
||
WMMAInTypes:$sourceA,
|
||
WMMAInTypes:$sourceB,
|
||
WMMAOutTypes:$destC,
|
||
DefaultValuedAttr<ConfinedAttr<I32Attr, [IntIsOneOf<[0, 1]>]>, "0">:$subwordOffset,
|
||
UnitAttr:$unsignedA,
|
||
UnitAttr:$unsignedB,
|
||
UnitAttr:$clamp)>,
|
||
Results<(outs WMMAOutTypes: $destD)> {
|
||
let summary = "MLIR wrapper for wmma instructions";
|
||
let description = [{
|
||
The `amdgpu.wmma` op is an MLIR wrapper around intrinsics for various `wmma`
|
||
instructions in the AMDGPU architecture, which perform matrix multiplication.
|
||
|
||
On gfx11/RDNA3, wmma intrinsics have M=N=K=16 dimensions.
|
||
|
||
On gfx12/RDNA4, wmma intrinsics have M=N=16 dimensions and support K=16 for
|
||
all element types, and K=32 for i4 sources.
|
||
|
||
On gfx1250, wmma intrinsics have M=N=16 and K dimensions of 4, 32, 64, or 128,
|
||
depending on the element types.
|
||
|
||
On gfx11/RDNA3, emitting f16->f16 (or bf16->bf16) wmma the output is a 16xf16
|
||
(or 16xbf16) vector containing only 8 valid values:
|
||
- If `subwordOffset` is 0, then the output is stored at indices 0, 2, 4, ..., 14.
|
||
- If `subwordOffset` is 1, then the output is stored at indices 1, 3, 5, ..., 15.
|
||
On gfx12/RDNA4 and gfx1250, the result is instead returned as vector where all
|
||
the values are valid and the `subwordOffset` must be `0`, as it cannot be used.
|
||
|
||
`unsignedA` and `unsignedB` flag that the `int8` LLVM inputs are unsigned.
|
||
|
||
The `clamp` flag is used to saturate the output of type T to `numeric_limits<T>::max()`
|
||
in case of overflow.
|
||
|
||
The `wave64`attribute indicates whether an op is designed for 64 threads wavefont.
|
||
|
||
Example:
|
||
```mlir
|
||
%0 = amdgpu.wmma 16x16x16 %matA * %matB + %matC : vector<8xf16>, vector<8xf16>, vector<8xf16>
|
||
|
||
%1 = amdgpu.wmma 16x16x64 %matD * %matE + %matF : vector<32xi8>, vector<8xf32>, vector<8xf32>
|
||
|
||
%2 = amdgpu.wmma 16x16x128 %matG * %matH + %matI : vector<64xf4E2M1FN>, vector<64xf4E2M1FN>, vector<8xf32>
|
||
|
||
%3 = amdgpu.wmma 16x16x4 %matJ * %matK + %matL : vector<2xf32>, vector<2xf32>, vector<8xf32>
|
||
```
|
||
}];
|
||
let assemblyFormat = [{
|
||
custom<MNKDimensionList>($m, $n, $k) $sourceA `*` $sourceB `+` $destC
|
||
attr-dict
|
||
`:` type($sourceA) `,` type($sourceB) `,` type($destC)
|
||
}];
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
def DotInTypes : AnyTypeOf<[VectorOfLengthAndType<[2], [F16, BF16, I16]>,
|
||
VectorOfLengthAndType<[4], [I8, F8E4M3FN, F8E5M2]>,
|
||
VectorOfLengthAndType<[8], [I<4>]>]>;
|
||
|
||
def DotOutTypes : AnyTypeOf<[F32, F16, BF16, I32]>;
|
||
|
||
def AMDGPU_DotOp :
|
||
AMDGPU_Op<"dot", [AllTypesMatch<["destC", "destD"]>, Pure]>,
|
||
Arguments<(ins DotInTypes:$sourceA,
|
||
DotInTypes:$sourceB,
|
||
DotOutTypes:$destC,
|
||
UnitAttr:$unsignedA,
|
||
UnitAttr:$unsignedB,
|
||
UnitAttr:$clamp)>,
|
||
Results<(outs DotOutTypes:$destD)> {
|
||
let summary = "MLIR wrapper for AMDGPU v_dot* intrinsics";
|
||
let description = [{
|
||
The `amdgpu.dot` op is an MLIR wrapper over the `v_dot*` family of intrinsics,
|
||
which compute `D = sum_i A[i] * B[i] + C`.
|
||
|
||
Variants (source, dest, signedness, chipset -> intrinsic).
|
||
|
||
```text
|
||
| A elem | B elem | destC | signedness | chipset | ROCDL op |
|
||
|----------|----------|-------|------------|---------------------------|------------------------------|
|
||
| f16 | f16 | f32 | n/a | gfx906+ | fdot2 |
|
||
| f16 | f16 | f16 | n/a | gfx11+ | fdot2.f16.f16 |
|
||
| bf16 | bf16 | f32 | n/a | gfx11+, gfx950+ | fdot2.f32.bf16 |
|
||
| bf16 | bf16 | bf16 | n/a | gfx11+ | fdot2.bf16.bf16 |
|
||
| i16 | i16 | i32 | s / u | gfx906+, no gfx11+/gfx12+ | sdot2 / udot2 |
|
||
| i8 | i8 | i32 | s / u | gfx906+ | sdot4 / udot4 |
|
||
| i8 | i8 | i32 | mixed | gfx11+ | sudot4 |
|
||
| i4 | i4 | i32 | s / u | gfx906+ | sdot8 / udot8 |
|
||
| i4 | i4 | i32 | mixed | gfx11+ | sudot8 |
|
||
| fp8/bf8 | fp8/bf8 | f32 | n/a | gfx11.7, gfx12+ | dot4.f32.{fp8,bf8}.{fp8,bf8} |
|
||
```
|
||
|
||
Example:
|
||
```mlir
|
||
%r0 = amdgpu.dot %a * %b + %c : vector<4xi8>, vector<4xi8>, i32
|
||
%r1 = amdgpu.dot %a * %b + %c {unsignedA, unsignedB, clamp}
|
||
: vector<8xi4>, vector<8xi4>, i32
|
||
%r2 = amdgpu.dot %a * %b + %c {unsignedB}
|
||
: vector<4xi8>, vector<4xi8>, i32
|
||
%r3 = amdgpu.dot %a * %b + %c : vector<2xf16>, vector<2xf16>, f32
|
||
%r4 = amdgpu.dot %a * %b + %c : vector<2xf16>, vector<2xf16>, f16
|
||
%r5 = amdgpu.dot %a * %b + %c
|
||
: vector<4xf8E4M3FN>, vector<4xf8E5M2>, f32
|
||
```
|
||
}];
|
||
let assemblyFormat = [{
|
||
$sourceA `*` $sourceB `+` $destC attr-dict
|
||
`:` type($sourceA) `,` type($sourceB) `,` type($destC)
|
||
}];
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
def AMDGPU_SparseMFMAOp :
|
||
AMDGPU_Op<"sparse_mfma", [AllTypesMatch<["destC", "destD"]>,
|
||
Pure]>,
|
||
Arguments<(ins
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[16, 32]>]>:$m,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[16, 32]>]>:$n,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[16, 32, 64, 128]>]>:$k,
|
||
SMFMACSparseInTypes:$sourceA,
|
||
SMFMACDenseInTypes:$sourceB,
|
||
SMFMACOutTypes:$destC,
|
||
SMFMACIdxTypes:$sparseIdx,
|
||
DefaultValuedAttr<I32Attr, "0">:$cbsz,
|
||
DefaultValuedAttr<I32Attr, "0">:$abid)>,
|
||
Results<(outs SMFMACOutTypes: $destD)> {
|
||
let summary = "MLIR wrapper for CDNA sparse mfma (smfmac) instructions";
|
||
let description = [{
|
||
The `amdgpu.sparse_mfma` op is an MLIR wrapper around intrinsics for various
|
||
`smfmac` instructions in the AMDGPU architecture, which perform matrix
|
||
multiply-accumulate operations using 2:4 structured sparsity on matrix A
|
||
with dense matrices B, C, and D.
|
||
|
||
On gfx942, smfmac intrinsics support:
|
||
- M=N=16, K=32 and M=N=32, K=16 for f16 and bf16 sources
|
||
- M=N=16, K=64 and M=N=32, K=32 for i8 and fp8 sources
|
||
|
||
On gfx950, smfmac intrinsics additionally support:
|
||
- M=N=16, K=64 and M=N=32, K=32 for f16 and bf16 sources
|
||
- M=N=16, K=128 and M=N=32, K=64 for i8 and fp8 sources
|
||
|
||
The `sparseIdx` parameter contains packed 2-bit indices identifying which
|
||
of every 4 dense-K positions are non-zero in the 2:4 sparse matrix A.
|
||
The required `sparseIdx` type depends on the variant:
|
||
- gfx942 16-bit (`(m,k)` in `{(16,32), (32,16)}`): 8 bits per lane,
|
||
carried as `vector<4xi8>` (one 8-bit set per i8 element).
|
||
- gfx942 8-bit (`(m,k)` in `{(16,64), (32,32)}`) and gfx950 16-bit
|
||
(`(m,k)` in `{(16,64), (32,32)}`): 16 bits per lane, carried as
|
||
`vector<2xi16>` (one 16-bit set per i16 element).
|
||
- gfx950 8-bit (`(m,k)` in `{(16,128), (32,64)}`): 32 bits per lane (a
|
||
full VGPR with no internal set structure), carried as `i32`.
|
||
|
||
The `cbsz` and `abid` parameters select which index set within the VGPR is
|
||
used:
|
||
- gfx942 16-bit: `cbsz == 0` selects one of four 8-bit sets via
|
||
`abid[1:0]` (range `[0, 3]`); `cbsz != 0` selects the first set.
|
||
- gfx942 8-bit and gfx950 16-bit: `cbsz == 0` selects one of two 16-bit
|
||
sets via `abid[0]` (range `[0, 1]`); `cbsz != 0` selects the first set.
|
||
- gfx950 8-bit: hardware ignores both `cbsz` and `abid`; both must be 0.
|
||
|
||
Example:
|
||
```mlir
|
||
%0 = amdgpu.sparse_mfma 16x16x32 %matA * %matB + %matC sparse(%idx : vector<4xi8>)
|
||
: vector<4xf16>, vector<8xf16>, vector<4xf32>
|
||
|
||
%1 = amdgpu.sparse_mfma 16x16x64 %matA * %matB + %matC sparse(%idx : vector<2xi16>)
|
||
: vector<8xi8>, vector<16xi8>, vector<4xi32>
|
||
|
||
%2 = amdgpu.sparse_mfma 16x16x64 %matA * %matB + %matC sparse(%idx : vector<2xi16>)
|
||
{ cbsz = 0 : i32, abid = 1 : i32 }
|
||
: vector<8xf8E4M3FNUZ>, vector<16xf8E4M3FNUZ>, vector<4xf32>
|
||
|
||
%3 = amdgpu.sparse_mfma 16x16x128 %matA * %matB + %matC sparse(%idx : i32)
|
||
: vector<16xf8E4M3FN>, vector<32xf8E4M3FN>, vector<4xf32>
|
||
```
|
||
}];
|
||
let assemblyFormat = [{
|
||
custom<MNKDimensionList>($m, $n, $k) $sourceA `*` $sourceB `+` $destC
|
||
`sparse` `(` $sparseIdx `:` type($sparseIdx) `)`
|
||
attr-dict
|
||
`:` type($sourceA) `,` type($sourceB) `,` type($destC)
|
||
}];
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
// sparse_wmma (swmmac)
|
||
def SWMMACSparseInTypes : AnyTypeOf<[
|
||
VectorOfLengthAndType<[4, 8, 16], [F16]>,
|
||
VectorOfLengthAndType<[4, 8, 16], [BF16]>,
|
||
VectorOfLengthAndType<[4, 8, 32], [I8]>,
|
||
VectorOfLengthAndType<[8, 16], [I<4>]>,
|
||
VectorOfLengthAndType<[4, 8, 16, 32], [F8E4M3FN, F8E5M2]>,
|
||
VectorOfLengthAndType<[4, 8, 16, 32], [F8E4M3FNUZ, F8E5M2FNUZ]>
|
||
]>;
|
||
|
||
def SWMMACDenseInTypes : AnyTypeOf<[
|
||
VectorOfLengthAndType<[8, 16, 32], [F16]>,
|
||
VectorOfLengthAndType<[8, 16, 32], [BF16]>,
|
||
VectorOfLengthAndType<[4, 8, 16, 64], [I8]>,
|
||
VectorOfLengthAndType<[8, 16, 32], [I<4>]>,
|
||
VectorOfLengthAndType<[4, 8, 16, 64], [F8E4M3FN, F8E5M2]>,
|
||
VectorOfLengthAndType<[4, 8, 16, 64], [F8E4M3FNUZ, F8E5M2FNUZ]>
|
||
]>;
|
||
|
||
def SWMMACOutTypes : AnyTypeOf<[
|
||
VectorOfLengthAndType<[4, 8, 16], [F32]>,
|
||
VectorOfLengthAndType<[4, 8], [F16]>,
|
||
VectorOfLengthAndType<[4, 8], [BF16]>,
|
||
VectorOfLengthAndType<[4, 8], [I32]>
|
||
]>;
|
||
|
||
def SWMMACIdxTypes : AnyTypeOf<[
|
||
FixedVectorOfLengthAndType<[4], [I8]>,
|
||
]>;
|
||
|
||
|
||
def AMDGPU_SparseWMMAOp :
|
||
AMDGPU_Op<"sparse_wmma", [AllTypesMatch<["destC", "destD"]>,
|
||
Pure]>,
|
||
Arguments<(ins
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[16]>]>:$m,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[16]>]>:$n,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[32, 64, 128]>]>:$k,
|
||
SWMMACSparseInTypes:$sourceA,
|
||
SWMMACDenseInTypes:$sourceB,
|
||
SWMMACOutTypes:$destC,
|
||
SWMMACIdxTypes:$sparseIdx,
|
||
UnitAttr:$unsignedA,
|
||
UnitAttr:$unsignedB,
|
||
UnitAttr:$reuseA,
|
||
UnitAttr:$reuseB,
|
||
UnitAttr:$clamp,
|
||
UnitAttr:$wave64)>,
|
||
Results<(outs SWMMACOutTypes: $destD)> {
|
||
let summary = "MLIR wrapper for gfx12+ sparse wmma instructions";
|
||
let description = [{
|
||
The `amdgpu.sparse_wmma` op is an MLIR wrapper around intrinsics for various
|
||
`swmmac` instructions in the AMDGPU architecture, which perform matrix
|
||
multiply-accumulate operations using 2:4 structured sparsity on matrix A
|
||
with dense matrices B, C, and D.
|
||
|
||
On gfx12, swmmac intrinsics support:
|
||
- M=N=16, K=32 and M=N=32, K=16 for f16, bf16, i8 and i4 sources
|
||
- M=N=16, K=64 for i4 sources
|
||
|
||
On gfx1250, swmmac intrinsics additionally support:
|
||
- M=N=16, K=64 for f16 and bf16 sources
|
||
- M=N=16, K=128 for f16, bf16 and i8 sources
|
||
|
||
The `sparseIdx` parameter contains packed indices identifying the positions
|
||
of non-zero elements in the 2:4 sparse matrix A. For 16-bit source data,
|
||
use `vector<4xi8>` (four 8-bit indices). For 8-bit source data, use
|
||
`vector<2xi16>` (two 16-bit indices).
|
||
|
||
`unsignedA` and `unsignedB` flag that the `int8` LLVM inputs are unsigned.
|
||
|
||
The `clamp` flag is used to saturate the output of type T to `numeric_limits<T>::max()`
|
||
in case of overflow.
|
||
|
||
Example:
|
||
```mlir
|
||
%0 = amdgpu.sparse_wmma 16x16x32 %matA * %matB + %matC sparse(%idx : vector<4xi8>)
|
||
: vector<4xf16>, vector<8xf16>, vector<4xf32>
|
||
|
||
%1 = amdgpu.sparse_wmma 16x16x64 %matA * %matB + %matC sparse(%idx : vector<2xi16>)
|
||
: vector<8xi8>, vector<16xi8>, vector<4xi32>
|
||
|
||
%2 = amdgpu.sparse_wmma 16x16x64 %matA * %matB + %matC sparse(%idx : vector<2xi16>)
|
||
{ unsignedA = 0 : i1, unsignedB = 1 : i1, clamp = 0 : i1 }
|
||
: vector<8xf8E4M3FNUZ>, vector<16xf8E4M3FNUZ>, vector<4xf32>
|
||
```
|
||
}];
|
||
let assemblyFormat = [{
|
||
custom<MNKDimensionList>($m, $n, $k) $sourceA `*` $sourceB `+` $destC
|
||
`sparse` `(` $sparseIdx `:` type($sparseIdx) `)`
|
||
attr-dict
|
||
`:` type($sourceA) `,` type($sourceB) `,` type($destC)
|
||
}];
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
// Promises IndexedMemCopyOpInterface.
|
||
def AMDGPU_GatherToLDSOp :
|
||
AMDGPU_Op<"gather_to_lds", [AttrSizedOperandSegments]>,
|
||
Arguments<(ins
|
||
Arg<AnyMemRef, "buffer to gather from", [MemRead]>:$src,
|
||
Variadic<Index>:$srcIndices,
|
||
Arg<AnyMemRef, "buffer to write to", [MemWrite]>:$dst,
|
||
Variadic<Index>:$dstIndices,
|
||
TypeAttr:$transferType,
|
||
UnitAttr:$async
|
||
)>,
|
||
Results<(outs)> {
|
||
let summary = "MLIR wrapper for CDNA Gather to LDS instructions";
|
||
let description = [{
|
||
The `amdgpu.gather_to_lds` op is a wrapper around the `global_load_lds` instructions.
|
||
|
||
Operands:
|
||
* `$src`: global memory (including fat buffer) memref to read from.
|
||
* `$srcIndices`: indices into `$src` to read from for this thread.
|
||
These indices must be non-negative and in-bounds when `$src` is not a
|
||
fat raw buffer. Fat raw buffer sources permit out-of-bounds indices with
|
||
raw buffer semantics.
|
||
* `$dst`: LDS memory memref to write to.
|
||
* `$dstIndices`: base indices into `$dst` to write to for the subgroup of this thread.
|
||
These indices must be non-negative and in-bounds.
|
||
The elements gathered by the subgroup will be written contiguously in order of lane ID
|
||
starting at `$dst[$dstIndices]`. Byte-sized (ex. i8) or short-sized (ex. i16)
|
||
types will be zero-padded/extended to 32 bits before being written. 96-bit types
|
||
(ex. vector<3xf32>) will be zero-padded to 128 bits before being written. Only the
|
||
offsets held by lane 0 are used.
|
||
* `$transferType`: type of the data to be transferred by each thread. This is used to determine
|
||
the size of the data to be transferred and the number of threads in the subgroup.
|
||
The transfer type must be a scalar type or a vector type with a single element type.
|
||
* If `$async` is set, the compiler will not attempt to infer the
|
||
memory waits needed to ensure that the DMA operation has succeeded
|
||
before a load that might access the stored-to LDS is performed.
|
||
Instead, the `rocdl.asyncmark` and `rocdl.wait.asyncmark N`
|
||
operations must be used to explicitly indicate the desired completion
|
||
behavior. This enables more precise calculation of these waits at the
|
||
cost of requiring user management of asynchrony.
|
||
|
||
The `$dst`, along with its indices, points to the memory location the subgroup of this thread
|
||
will write to.
|
||
|
||
Note: only supported on gfx9 and gfx10.
|
||
}];
|
||
let assemblyFormat = [{
|
||
(`async` $async^)? $src `[` $srcIndices `]` `,` $dst `[` $dstIndices `]` attr-dict `:` $transferType `,` type($src) `,` type($dst)
|
||
}];
|
||
let hasVerifier = 1;
|
||
let hasCanonicalizer = 1;
|
||
}
|
||
|
||
// Promises IndexedMemCopyOpInterface.
|
||
def AMDGPU_GlobalLoadAsyncToLDSOp :
|
||
AMDGPU_Op<"global_load_async_to_lds", [AttrSizedOperandSegments]>,
|
||
Arguments<(ins
|
||
Arg<AnyMemRef, "global memory to load from", [MemRead]>:$src,
|
||
Variadic<Index>:$srcIndices,
|
||
Arg<AnyMemRef, "LDS memory to write to", [MemWrite]>:$dst,
|
||
Variadic<Index>:$dstIndices,
|
||
TypeAttr:$transferType,
|
||
Optional<I1>: $mask
|
||
)>,
|
||
Results<(outs)> {
|
||
let summary = "MLIR wrapper for async global load to lds instructions";
|
||
let description = [{
|
||
AMDGPU wrapper for `global.load.async.to.lds` instructions, which performs
|
||
asynchronous load of data from global memory into LDS while bypassing VGPRs.
|
||
|
||
* `$src`: global memory memref to read from (global addrspace only, no fat buffer).
|
||
* `$srcIndices`: indices into `$src` for this thread's global read location.
|
||
These indices must be non-negative and in-bounds.
|
||
* `$dst`: LDS memref to write to (workgroup addrspace).
|
||
* `$dstIndices`: indices into `$dst` for this thread's LDS write location.
|
||
These indices must be non-negative and in-bounds when `$mask` is not
|
||
provided. When `$mask` is provided, the destination indices are not
|
||
guaranteed to be in-bounds because masked-off lanes may carry invalid
|
||
destination indices.
|
||
* `$transferType`: type of data to be transferred. Must be 8, 32, 64 or 128 bit scalar
|
||
or vector type.
|
||
* `$mask`: optional per-thread mask. When false, the thread's LDS write
|
||
is masked off. The global read still occurs for all threads regardless of mask.
|
||
|
||
Note: only supported on gfx1250 and later.
|
||
|
||
Examples:
|
||
```mlir
|
||
amdgpu.global_load_async_to_lds %src[%i, %j], %dst[%k, %l]
|
||
: f32, memref<128x64xf32, #gpu.address_space<global>>,
|
||
memref<64x64xf32, #gpu.address_space<workgroup>>
|
||
|
||
amdgpu.global_load_async_to_lds %src[%i, %j], %dst[%k, %l]
|
||
: vector<4xf32>, memref<128x64xf32, #gpu.address_space<global>>,
|
||
memref<64x64xf32, #gpu.address_space<workgroup>>
|
||
|
||
amdgpu.global_load_async_to_lds %src[%i], %dst[%j]
|
||
: i8, memref<512xi8, #gpu.address_space<global>>,
|
||
memref<256xi8, #gpu.address_space<workgroup>>
|
||
```
|
||
}];
|
||
let assemblyFormat = [{
|
||
$src `[` $srcIndices `]` `,` $dst `[` $dstIndices `]` (`,` $mask^)?
|
||
attr-dict `:` $transferType `,` type($src) `,` type($dst)
|
||
}];
|
||
let hasVerifier = 1;
|
||
let hasCanonicalizer = 1;
|
||
}
|
||
|
||
// Promises IndexedAccessOpInterface.
|
||
def AMDGPU_TransposeLoadOp :
|
||
AMDGPU_Op<"transpose_load", [SameVariadicOperandSize]>,
|
||
Arguments<(ins Arg<AnyMemRef, "buffer to transpose load from", [MemRead]>:$src, Variadic<Index>:$srcIndices)>,
|
||
Results<(outs AnyTypeOf<[AnyVectorOfNonZeroRank]>:$result)> {
|
||
let summary = "MLIR wrapper for CDNA transpose Load instructions";
|
||
let description = [{
|
||
The `amdgpu.transpose_load` op is a wrapper around the `ds_read_tr` instructions
|
||
on gfx9 and the `ds_load_tr` family of instructions on gfx1250.
|
||
|
||
The transpose load op represents a subgroup load from LDS memory,
|
||
where the subgroup of threads collectively reads a matrix from the source
|
||
memref, with each thread reading a vector of the matrix, and gets a transposed matrix
|
||
in as the result. That is, each thread reads a vector of the col-major matrix at different
|
||
indices, and the thread's read result is a vector of the corresponding row of the transposed
|
||
matrix.
|
||
|
||
This op is a direct wrapper around the ROCDL `ds_read_tr` family intrinsics on
|
||
gfx950 and the `ds_load_tr` family of instructions on gfx1250. Please refer
|
||
to the respective ISA documentation for more details about its exact semantics.
|
||
|
||
Format example:
|
||
```
|
||
%0 = amdgpu.transpose_load %src[%srcIndices] : memref<128x256xf16> -> vector<4xf16>
|
||
```
|
||
Operands:
|
||
* `$src`: LDS memref to read from.
|
||
* `$srcIndices`: indices into `$src` to read from for this thread. Indices
|
||
must be non-negative and in-bounds for the corresponding dimension of
|
||
`$src`, matching the constraints of `memref.load`.
|
||
* `$result`: target register this transpose load instruction will write to.
|
||
|
||
Note: Lowering is only supported on gfx950 and gfx1250, with different
|
||
permitted load types.
|
||
}];
|
||
let assemblyFormat = [{
|
||
$src `[` $srcIndices `]` attr-dict `:` type($src) `->` type($result)
|
||
}];
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
// Promises IndexedAccessOpInterface.
|
||
def AMDGPU_GlobalTransposeLoadOp :
|
||
AMDGPU_Op<"global_transpose_load", [SameVariadicOperandSize]>,
|
||
Arguments<(ins Arg<AnyMemRef, "buffer to transpose load from", [MemRead]>:$src,
|
||
Variadic<Index>:$srcIndices)>,
|
||
Results<(outs AnyTypeOf<[
|
||
FixedVectorOfLengthAndType<[8], [I8, F8E5M2FNUZ, F8E4M3FNUZ, F8E5M2,
|
||
F8E4M3FN, F16, BF16, I16]>,
|
||
FixedVectorOfLengthAndType<[16], [I<4>, F4E2M1FN, I<6>, F6E2M3FN, F6E3M2FN]>
|
||
]>:$result)> {
|
||
let summary = "MLIR wrapper for global memory transpose load instructions";
|
||
let description = [{
|
||
The `amdgpu.global_transpose_load` op is a wrapper around the
|
||
`global_load_tr` family of instructions introduced in gfx1200.
|
||
|
||
Each thread reads a column of a matrix stored in global memory and receives
|
||
the corresponding row of the transposed matrix in its result register.
|
||
The subgroup collectively performs a transpose of the tile.
|
||
|
||
This op is a direct wrapper around the ROCDL `global.load.tr` family
|
||
intrinsics. Refer to the ISA manual for exact semantics.
|
||
|
||
Format example:
|
||
```mlir
|
||
%0 = amdgpu.global_transpose_load %src[%i, %j]
|
||
: memref<128x256xf16, #gpu.address_space<global>> -> vector<8xf16>
|
||
```
|
||
|
||
Operands:
|
||
* `$src`: Global address space memref to read from.
|
||
* `$srcIndices`: indices into `$src` for this thread. Indices must be
|
||
non-negative and in-bounds for the corresponding dimension of `$src`,
|
||
matching the constraints of `memref.load`.
|
||
* `$result`: register this transpose load instruction writes to.
|
||
|
||
Valid (element bits, num elements) pairs:
|
||
* (4, 16) -> global_load_tr4_b64 (gfx1250+)
|
||
* (6, 16) -> global_load_tr6_b96 (gfx1250+)
|
||
* (8, 8) -> global_load_tr_b64 (gfx1200+)
|
||
* (16, 8) -> global_load_tr_b128 (gfx1200+)
|
||
|
||
Note: 8-bit and 16-bit element lowering requires gfx1200+.
|
||
4-bit and 6-bit element lowering requires gfx1250+.
|
||
}];
|
||
let assemblyFormat = [{
|
||
$src `[` $srcIndices `]` attr-dict `:` type($src) `->` type($result)
|
||
}];
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
def AMDGPU_ScaledMFMAOp :
|
||
AMDGPU_Op<"scaled_mfma", [AllTypesMatch<["destC", "destD"]>,
|
||
Pure]>,
|
||
Arguments<(ins
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[16, 32]>]>:$m,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[16, 32]>]>:$n,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[64, 128]>]>:$k,
|
||
ScaledMFMAInTypes:$sourceA,
|
||
ScaledMFMAInTypes:$sourceB,
|
||
ScaledMFMAOutTypes:$destC,
|
||
AnyTypeOf<[F8E8M0FNU, FixedVectorOfLengthAndType<[4], [F8E8M0FNU]>]>:$scalesA,
|
||
AnyTypeOf<[F8E8M0FNU, FixedVectorOfLengthAndType<[4], [F8E8M0FNU]>]>:$scalesB,
|
||
ConfinedAttr<I32Attr, [IntNonNegative, IntMaxValue<3>]>:$scalesIdxA,
|
||
ConfinedAttr<I32Attr, [IntNonNegative, IntMaxValue<3>]>:$scalesIdxB
|
||
)>,
|
||
Results<(outs ScaledMFMAOutTypes: $destD)> {
|
||
let summary = "MLIR wrapper for CDNA scaled mfma instructions";
|
||
let description = [{
|
||
The `amdgpu.scaled_mfma` op is an MLIR wrapper around intrinsics
|
||
for various scaled versions of `mfma` instructions in the CDNA architecture, which
|
||
perform multiple outer products in order to allow fast matrix multiplication.
|
||
|
||
The wrapper will select an appropriate `mfma` instruction, if one is available,
|
||
based on the provided `m`, `k`, `n`, and `nBlks` attributes, along with the
|
||
types of the source and destination arguments.
|
||
|
||
Note, this wrapper allows specifying `vector<4Kxi8>` arguments to MFMA
|
||
intrinsics that take an integer type of width `4K`. For example,
|
||
one can provide a `vector<4xi8>` as an argument to an MFMA instruction that
|
||
logically takes 4 i8s but whose intrinsics are specified to take an i32.
|
||
In these cases, the bytes in the vector will be concatenated in little-endian
|
||
order (that is, v[0] will go to arg[7:0], v[1] to arg[15:8] and so on).
|
||
|
||
This wrapper takes inspiration from `amdgpu.mfma`, but has some key differences:
|
||
- `amdgpu.scaled_mfma` operates on fp4 (f4E2M1FN), fp6 (f6E2M3FN and f6E3M2FN) and
|
||
fp8 (f8E4M3FN and f8E5M2) types using either M=N=16, K=128 or M=N=32, K=64 as
|
||
their tile size.
|
||
- `amdgpu.scaled_mfma` does not support broadcasting. So, `cbsz`, `abid`, and `blgp`
|
||
are omitted from this wrapper.
|
||
- The `negateA`, `negateB`, and `negateC` flags in `amdgpu.mfma` are only supported
|
||
for double-precision operations on gfx94x and so are not included here.
|
||
|
||
Example:
|
||
```mlir
|
||
%0 = amdgpu.scaled_mfma 32x32x64 (%arg0[0] * %arg1) * (%arg0[1] * %arg1) + %arg2
|
||
: vector<4xf8E8M0FNU>, vector<32xf6E2M3FN>, f8E8M0FNU, vector<32xf6E2M3FN>, vector<16xf32>
|
||
```
|
||
}];
|
||
let assemblyFormat = [{
|
||
custom<MNKDimensionList>($m, $n, $k) ` `
|
||
`(` $scalesA `[` $scalesIdxA `]` `*` $sourceA `)` `*`
|
||
`(` $scalesB `[` $scalesIdxB `]` `*` $sourceB `)` `+` $destC
|
||
attr-dict
|
||
`:` type($scalesA) `,` type($sourceA) `,` type($scalesB) `,` type($sourceB) `,` type($destC)
|
||
}];
|
||
let hasCanonicalizer = 1;
|
||
}
|
||
|
||
def AMDGPU_ScaledWMMAOp
|
||
: AMDGPU_Op<"scaled_wmma", [AllTypesMatch<["destC", "destD"]>, Pure]>,
|
||
Arguments<(ins ConfinedAttr<I32Attr, [IntIsOneOf<[16, 32]>]>:$m,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[16]>]>:$n,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[128]>]>:$k,
|
||
ScaledWMMAInTypes:$sourceA, ScaledWMMAInTypes:$sourceB,
|
||
ScaledWMMAOutTypes:$destC,
|
||
VectorOfLengthAndType<[4, 8], [F8E8M0FNU, F8E4M3FN]>:$scaleA,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[0, 16]>]>:$a_first_scale_lane,
|
||
VectorOfLengthAndType<[4, 8], [F8E8M0FNU, F8E4M3FN]>:$scaleB,
|
||
ConfinedAttr<I32Attr, [IntIsOneOf<[0, 16]>]>:$b_first_scale_lane)>,
|
||
Results<(outs ScaledWMMAOutTypes:$destD)> {
|
||
// TODO: E5M3FNU scales are supported, but there is not yet MLIR support for
|
||
// this datatype. Once we have support for that, update the scaleA and scaleB
|
||
// types here.
|
||
let summary = "MLIR wrapper for scaled wmma instructions";
|
||
let description = [{
|
||
The `amdgpu.scaled_wmma` op is an MLIR wrapper around intrinsics for scaled
|
||
`wmma` instructions. These instructions perform matrix multiplication with
|
||
per-block scaling of inputs, supporting fp4, fp6, and fp8 data formats.
|
||
|
||
The scale instructions support a block size of 16 or 32 and two tile sizes:
|
||
- 16x16x128 with mixed f8/f6/f4 formats (output: vector<8xf32>)
|
||
- 32x16x128 with f4 format only (output: vector<16xf32>)
|
||
|
||
Scale parameters (`scaleA`, `scaleB`) are small vectors of f8 scale values
|
||
(either f8E8M0FNU, or f8E4M3FN) that are packed into i32/i64 values during
|
||
lowering. Each lane can operate on 4 bytes (4 scale values), and the
|
||
number of scales required for each matrix is determined by:
|
||
num_scales_A = (M × K) / block_size
|
||
num_scales_B = (N × K) / block_size
|
||
|
||
The index attributes (`a_first_scale_lane`, `b_first_scale_lane`) select
|
||
which lane to start reading scale values from (0 or 16):
|
||
- For block size 32, 32 lanes across a single wave are used for the scale
|
||
values. If the number of scales (num_scales_A or num_scales_B) can fit
|
||
into half of the available lanes
|
||
(i.e., num_scales / scales_per_lane == 16 (num_lanes)),
|
||
then then first_scale_lane can be either 0 or 16. If all lanes are required
|
||
for storing the scale values (num_scales / scales_per_lane == 32 (num_lanes)),
|
||
then the first_scale_lane must be 0.
|
||
- For block size 16, the same rules apply as above except that there are 64
|
||
lanes across two waves that are used for the scale values. When
|
||
num_scales / scales_per_lane == 32 (num lanes), then 16 lanes from each wave are used.
|
||
first_scale_lane of 0 or 16 will decide which lanes are used for this. When
|
||
num_scales / scales_per_lane == 64 (num_lanes), then first_scale_lane must
|
||
be set to 0.
|
||
|
||
Example:
|
||
```mlir
|
||
// 16x16x128: fp8 inputs
|
||
%0 = amdgpu.scaled_wmma 16x16x128 (%scaleVecA * %matA) * (%scaleVecB * %matB) + %matC
|
||
{a_first_scale_lane = 0 : i32, b_first_scale_lane = 0 : i32}
|
||
: vector<4xf8E8M0FNU>, vector<64xf8E4M3FN>,
|
||
vector<4xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<8xf32>
|
||
|
||
// 32x16x128: fp4 inputs with different scale lanes
|
||
%1 = amdgpu.scaled_wmma 32x16x128 (%scaleVecD * %matD) * (%scaleVecE * %matE) + %matF
|
||
{a_first_scale_lane = 0 : i32, b_first_scale_lane = 16 : i32}
|
||
: vector<8xf8E4M3FN>, vector<128xf4E2M1FN>,
|
||
vector<8xf8E4M3FN>, vector<64xf4E2M1FN>, vector<16xf32>
|
||
```
|
||
}];
|
||
let assemblyFormat = [{
|
||
custom<MNKDimensionList>($m, $n, $k) ` `
|
||
`(` $scaleA `*` $sourceA `)` `*`
|
||
`(` $scaleB `*` $sourceB `)` `+` $destC
|
||
attr-dict
|
||
`:` type($scaleA) `,` type($sourceA) `,` type($scaleB) `,` type($sourceB) `,` type($destC)
|
||
}];
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
class AMDGPU_DmaBaseOp<string mnemonic, Type outType> :
|
||
AMDGPU_Op<mnemonic, [Pure, AttrSizedOperandSegments, AllElementTypesMatch<["global", "lds"]>]>,
|
||
Arguments<(ins Arg<AnyMemRef>:$global,
|
||
Variadic<Index>:$global_indices,
|
||
Arg<AnyMemRef>:$lds,
|
||
Variadic<Index>:$lds_indices)>,
|
||
Results<(outs outType: $base)> {
|
||
|
||
let assemblyFormat = [{
|
||
$global `[` $global_indices `]` `,` $lds `[` $lds_indices `]` attr-dict `:` type($global) `,` type($lds) `->` type(results)
|
||
}];
|
||
}
|
||
|
||
// Promises IndexedMemCopyOpInterface.
|
||
def AMDGPU_MakeGatherDmaBaseOp : AMDGPU_DmaBaseOp<"make_gather_dma_base", AMDGPU_TDMGatherBaseType> {
|
||
let summary = "Pair of based addresses used when moving tiles between LDS and global memory.";
|
||
|
||
let description = [{
|
||
This operation creates a pair of addresses that will be used by `tensor_load_to_lds`
|
||
and `tensor_store_from_lds`.
|
||
|
||
The global and LDS indices must be non-negative and in-bounds for the
|
||
corresponding dimensions of their memrefs.
|
||
|
||
This operation creates a value corresponding to the tensor descriptor (D#) group 0
|
||
found in TensorLoadToLDSOp and TensorStoreFromLDSOp in the rocdl dialect.
|
||
|
||
Unlike `make_dma_base`, this operation returns `!amdgpu.tdm_gather_base<$element_type, $index_type>`
|
||
which is only compatible with `make_gather_dma_descriptor`. Using the descriptor returned
|
||
by `make_gather_dma_descriptor` will set the `tensor_load_to_lds` and `tensor_store_from_lds` to gather mode.
|
||
|
||
```mlir
|
||
%base = amdgpu.make_gather_dma_base %global[%idx0, %idx1], %lds[%idx2, %idx3] : memref<64x64xi32>, memref<64x64xi32, #gpu.address_space<workgroup>> -> !amdgpu.tdm_gather_base<i32, i16>
|
||
// %indices : i16
|
||
%descriptor = amdgpu.make_gather_dma_descriptor %base[%indices] globalSize [2, 2] globalStride [2, 1] sharedSize [2, 2] : !amdgpu.tdm_gather_base<i32, i16>, i16 -> !amdgpu.tdm_descriptor
|
||
amdgpu.tensor_load_to_lds %descriptor : !amdgpu.tdm_descriptor
|
||
```
|
||
}];
|
||
|
||
let hasVerifier = 1;
|
||
|
||
let extraClassDeclaration = [{
|
||
static constexpr bool isGather() {
|
||
return true;
|
||
}
|
||
}];
|
||
}
|
||
|
||
// Promises IndexedMemCopyOpInterface.
|
||
def AMDGPU_MakeDmaBaseOp : AMDGPU_DmaBaseOp<"make_dma_base", AMDGPU_TDMBaseType> {
|
||
|
||
let summary = "Pair of based addresses used when moving tiles between LDS and global memory.";
|
||
let description = [{
|
||
This operation creates a pair of addresses that will be used by tensor_load_to_lds
|
||
and tensor_store_from_lds.
|
||
|
||
The global and LDS indices must be non-negative and in-bounds for the
|
||
corresponding dimensions of their memrefs.
|
||
|
||
This operation creates a value corresponding to the tensor descriptor (D#) group 0
|
||
found in TensorLoadToLDSOp and TensorStoreFromLDSOp in the rocdl dialect.
|
||
|
||
For example:
|
||
|
||
```mlir
|
||
%base = amdgpu.make_dma_base %global[%idx0, %idx1], %lds[%idx2, %idx3] : memref<64x64xi32>, memref<64x64xi32, #gpu.address_space<workgroup>> -> !amdgpu.tdm_base<i32>
|
||
%descriptor = amdgpu.make_dma_descriptor %base globalSize [2, 2] globalStride [2, 1] sharedSize [2, 2] : !amdgpu.tdm_base<i32> -> !amdgpu.tdm_descriptor
|
||
amdgpu.tensor_load_to_lds %descriptor : !amdgpu.tdm_descriptor
|
||
```
|
||
|
||
to
|
||
|
||
```mlir
|
||
// pseudo-code
|
||
%global_base = llvm.extractvalue %global_memref[1]
|
||
%global_address = llvm.get_element_ptr ...
|
||
|
||
%lds_base = llvm.extractvalue %lds_memref[1]
|
||
%lds_address = llvm.get_element_ptr ...
|
||
|
||
// Definition of %base
|
||
%undef = llvm.mlir.undef : vector<4xi32>
|
||
%v0 = llvm.insertelement %15, %undef[0] : vector<4xi32>
|
||
%v1 = llvm.insertelement %lds_address, %v0[1] : vector<4xi32>
|
||
%v2 = llvm.insertelement %global_address_low, %v1[2] : vector<4xi32>
|
||
%base = llvm.insertelement %global_address_high, %v2[3] : vector<4xi32>
|
||
|
||
rocdl.tensor.load.to.lds %base, %dgroup1, %dgroup2, %dgroup3 cachepolicy 0 : vector<4xi32>, vector<8xi32>
|
||
```
|
||
|
||
These tensor DMA operations were introduced in gfx1250.
|
||
}];
|
||
|
||
let hasVerifier = 1;
|
||
|
||
let extraClassDeclaration = [{
|
||
static constexpr bool isGather() {
|
||
return false;
|
||
}
|
||
}];
|
||
}
|
||
|
||
class AMDGPU_MakeDescriptorOp<string mnemonic> :
|
||
AMDGPU_Op<mnemonic, [Pure, AttrSizedOperandSegments]>,
|
||
Results<(outs AMDGPU_TDMDescriptorType: $desc)> {
|
||
|
||
dag baseArgs = (ins
|
||
Variadic<Index>: $global_dynamic_sizes,
|
||
DenseI64ArrayAttr: $global_static_sizes,
|
||
Variadic<Index>: $global_dynamic_strides,
|
||
DenseI64ArrayAttr: $global_static_strides,
|
||
Variadic<Index>: $shared_dynamic_sizes,
|
||
DenseI64ArrayAttr: $shared_static_sizes,
|
||
Optional<AMDGPU_ConcreteVector<I1, 16>>: $workgroup_mask,
|
||
Optional<I1>: $early_timeout,
|
||
Optional<I32>: $pad_amount,
|
||
Optional<I32>: $pad_interval,
|
||
Optional<MemRefOf<[AMDGPU_DsBarrierStateType]>>: $atomic_barrier_address,
|
||
Variadic<Index>: $atomic_barrier_indices,
|
||
Optional<Index>: $global_increment,
|
||
Optional<I32>: $lds_increment,
|
||
Optional<Index>: $iteration_count);
|
||
|
||
code extraClassDeclarationBase = [{
|
||
int64_t getRank() {
|
||
return getGlobalStaticSizes().size();
|
||
}
|
||
|
||
unsigned getElementTypeWidth() {
|
||
return getBase().getType().getElementType().getIntOrFloatBitWidth();
|
||
}
|
||
|
||
SmallVector<OpFoldResult> getMixedGlobalSizes() {
|
||
return getMixedValues(getGlobalStaticSizes(), getGlobalDynamicSizes(), getContext());
|
||
}
|
||
|
||
SmallVector<OpFoldResult> getMixedGlobalStrides() {
|
||
return getMixedValues(getGlobalStaticStrides(), getGlobalDynamicStrides(), getContext());
|
||
}
|
||
|
||
SmallVector<OpFoldResult> getMixedSharedSizes() {
|
||
return getMixedValues(getSharedStaticSizes(), getSharedDynamicSizes(), getContext());
|
||
}
|
||
|
||
}];
|
||
|
||
}
|
||
|
||
// Promises IndexedAccessOpInterface for the optional atomic barrier address.
|
||
def AMDGPU_MakeGatherDmaDescriptorOp : AMDGPU_MakeDescriptorOp<"make_gather_dma_descriptor"> {
|
||
dag args = (ins AMDGPU_TDMGatherBaseType: $base,
|
||
AnyTypeOf<[VectorOfMinMaxLengthAndType<1, 8, [I32]>,
|
||
VectorOfMinMaxLengthAndType<1, 16, [I16]>]>: $indices);
|
||
let arguments = !con(args, baseArgs);
|
||
let summary = "Make all descriptor groups needed by TensorLoadToLDS/TensorStoreFromLDS.";
|
||
let description = [{
|
||
Make all descriptor groups needed by tensor memory operations in gather
|
||
mode.
|
||
|
||
If an atomic barrier is provided, its indices must be non-negative and
|
||
in-bounds for the corresponding dimensions of the barrier memref.
|
||
}];
|
||
|
||
let assemblyFormat = [{
|
||
$base `[` $indices `]`
|
||
`globalSize` custom<DynamicIndexList>($global_dynamic_sizes, $global_static_sizes)
|
||
`globalStride` custom<DynamicIndexList>($global_dynamic_strides, $global_static_strides)
|
||
`sharedSize` custom<DynamicIndexList>($shared_dynamic_sizes, $shared_static_sizes)
|
||
( `padShared` `(` $pad_amount^ `every` $pad_interval `)` )?
|
||
( `workgroupMask` $workgroup_mask^ ( `earlyTimeout` $early_timeout^)?)?
|
||
( `atomicBarrier` `(` $atomic_barrier_address^ `[` $atomic_barrier_indices `]`
|
||
`:` type($atomic_barrier_address) `)`)?
|
||
( `iterate` $global_increment^ `,` $lds_increment `,` $iteration_count )?
|
||
attr-dict `:` qualified(type($base)) `,` type($indices) `->` type(results)
|
||
}];
|
||
|
||
let hasVerifier = 1;
|
||
let hasFolder = 1;
|
||
|
||
let extraClassDeclaration = extraClassDeclarationBase # [{
|
||
static constexpr bool isGather() {
|
||
return true;
|
||
}
|
||
}];
|
||
}
|
||
|
||
// Promises IndexedAccessOpInterface for the optional atomic barrier address.
|
||
def AMDGPU_MakeDmaDescriptorOp : AMDGPU_MakeDescriptorOp<"make_dma_descriptor"> {
|
||
dag args = (ins AMDGPU_TDMBaseType: $base);
|
||
let arguments = !con(args, baseArgs);
|
||
let summary = "Make all descriptor groups needed by TensorLoadToLDS/TensorStoreFromLDS.";
|
||
let description = [{
|
||
Make all descriptor groups needed by tensor memory operations.
|
||
|
||
The $base operand corresponds to the base pair addresses, one must be an address in LDS
|
||
while the other must be a global memory location.
|
||
|
||
$global_{static/dynamic}_sizes determine the size of the tensor.
|
||
$global_{static/dynamic}_strides determine the strides of the tensor.
|
||
$shared_{static/dynamic}_sizes determines the size of the tile.
|
||
|
||
$workgroup_mask broadcast load to workgroups inside of a workgroup cluster
|
||
(0 = do not broadcast result to workgroup, 1 = broadcast result to workgroup). Ignored for stores.
|
||
An all zeros mask is interpreted as a non-broadcasted load.
|
||
|
||
$early_timeout return data to requesters as soon as cache supplies it.
|
||
|
||
Padding can be applied to the LDS address when copying from memory to LDS,
|
||
but not when copying from LDS to memory.
|
||
The values in the padded target addresses remain the same as before the operation was applied.
|
||
$pad_interval must be a power of two contained in [2, 256].
|
||
$pad_amount must be a value contained in [1, 128].
|
||
|
||
If an atomic barrier is provided, it will be arrived at **once** after
|
||
each load/store using this descriptor is completed.
|
||
Its indices must be non-negative and in-bounds for the corresponding
|
||
dimensions of the barrier memref.
|
||
|
||
2D and 3D tensors may be iterated over by setting $global_increment, $lds_increment, and $iteration_count.
|
||
$global_increment determines how much to increment the starting global memory address per iteration in units of the $base's element type.
|
||
$lds_increment determines how much to increment the starting LDS address per iteration in units of the $base's element type.
|
||
$iterate_count determines how many times to iterate, it must be a value in the inclusive interval [1, 256].
|
||
|
||
```mlir
|
||
// Example of moving a two-dimensional tensor to LDS.
|
||
%base = amdgpu.make_dma_base %global[0, 0], %lds[0, 0] : memref<64x64xi32>, memref<64x64xi32, #gpu.address_space<workgroup>> -> !amdgpu.tdm_base<i32>
|
||
%descriptor = amdgpu.make_dma_descriptor %base globalSize [64, 64] globalStride [64, 1] sharedSize [64, 64] : !amdgpu.tdm_base<i32> -> !amdgpu.tdm_descriptor
|
||
amdgpu.tensor_load_to_lds %descriptor : !amdgpu.tdm_descriptor
|
||
|
||
// Example of moving a two dimension tensor to LDS where padding is applied after every integer.
|
||
%base = amdgpu.make_dma_base %global[0, 0], %lds[0, 0] : memref<32x32xi32>, memref<64x64xi32, #gpu.address_space<workgroup>> -> !amdgpu.tdm_base<i32>
|
||
%descriptor = amdgpu.make_dma_descriptor %base globalSize [32, 32] globalStride [32, 1] sharedSize [64, 64] padShared(%pad_amount every %pad_interval) : !amdgpu.tdm_base<i32> -> !amdgpu.tdm_descriptor
|
||
amdgpu.tensor_load_to_lds %descriptor : !amdgpu.tdm_descriptor
|
||
```
|
||
}];
|
||
|
||
let assemblyFormat = [{
|
||
$base
|
||
`globalSize` custom<DynamicIndexList>($global_dynamic_sizes, $global_static_sizes)
|
||
`globalStride` custom<DynamicIndexList>($global_dynamic_strides, $global_static_strides)
|
||
`sharedSize` custom<DynamicIndexList>($shared_dynamic_sizes, $shared_static_sizes)
|
||
( `padShared` `(` $pad_amount^ `every` $pad_interval `)` )?
|
||
( `workgroupMask` $workgroup_mask^ ( `earlyTimeout` $early_timeout^)?)?
|
||
( `atomicBarrier` `(` $atomic_barrier_address^ `[` $atomic_barrier_indices `]`
|
||
`:` type($atomic_barrier_address) `)`)?
|
||
( `iterate` $global_increment^ `,` $lds_increment `,` $iteration_count )?
|
||
attr-dict `:` qualified(type($base)) `->` type(results)
|
||
}];
|
||
|
||
let hasVerifier = 1;
|
||
let hasFolder = 1;
|
||
|
||
let extraClassDeclaration = extraClassDeclarationBase # [{
|
||
static constexpr bool isGather() {
|
||
return false;
|
||
}
|
||
}];
|
||
|
||
}
|
||
|
||
def AMDGPU_TensorLoadToLDSOp :
|
||
AMDGPU_Op<"tensor_load_to_lds", [MemoryEffects<[MemWrite, MemRead]>]>,
|
||
Arguments<(ins AMDGPU_TDMDescriptorType: $desc)> {
|
||
let summary = "Load tensors from global memory to LDS.";
|
||
let description = [{
|
||
Load tensors of up to five dimensions from global memory to LDS.
|
||
|
||
This operation was introduced in gfx1250.
|
||
}];
|
||
|
||
let assemblyFormat = [{
|
||
$desc attr-dict `:` qualified(type($desc))
|
||
}];
|
||
}
|
||
|
||
def AMDGPU_TensorStoreFromLDSOp :
|
||
AMDGPU_Op<"tensor_store_from_lds", [MemoryEffects<[MemWrite, MemRead]>]>,
|
||
Arguments<(ins AMDGPU_TDMDescriptorType: $desc)> {
|
||
|
||
let summary = "Store tensors from LDS to global memory.";
|
||
let description = [{
|
||
Store tensors of up to five dimensions from LDS to global memory.
|
||
|
||
This operation was introduced in gfx1250.
|
||
}];
|
||
|
||
let assemblyFormat = [{
|
||
$desc attr-dict `:` qualified(type($desc))
|
||
}];
|
||
}
|
||
|
||
//===----------------------------------------------------------------------===//
|
||
// In-LDS Barrier Operations
|
||
//
|
||
// General note: any of these operations that impact memory have read and write
|
||
// effects as a crude model of their atomic nature - we don't want "reads"
|
||
// being hoisted out of loops.
|
||
//===----------------------------------------------------------------------===//
|
||
|
||
// Promises IndexedAccessOpInterface.
|
||
def AMDGPU_DsBarrierInitOp :
|
||
AMDGPU_Op<"ds_barrier_init">,
|
||
Arguments<(ins Arg<MemRefOf<[AMDGPU_DsBarrierStateType]>, "barrier(s)",
|
||
[MemRead, MemWrite]>:$base,
|
||
Variadic<Index>:$indices,
|
||
I32:$participants)> {
|
||
let summary = "Initialize an in-LDS barrier.";
|
||
let description = [{
|
||
Given the location `!amdgpu.ds_barrier_state` in LDS (as specified by `base` and `indices`),
|
||
initialize the barrier structure so that the pending and init counts are equal to
|
||
`participants - 1`, which will have its high bits masked off, and its phase is equal to 0.
|
||
|
||
The indices must be non-negative and in-bounds for the corresponding
|
||
dimensions of `base`.
|
||
|
||
Note that we subtract 1 from `participants` when constructing the barrier state
|
||
to provide clearer high-level semantics.
|
||
|
||
The subtraction means that, when the `participant`th arrival occurs, the phase will change.
|
||
In practical terms, this means that you can use (for example) the number of subgroups or
|
||
waves per workgroup as `participants`, instead of manually needing to remove one.
|
||
|
||
While the write of the initial state will be performed atomically, no synchronization
|
||
between waves will be performed by this operation.
|
||
|
||
Example:
|
||
```mlir
|
||
amdgpu.ds_barrier_init %barrier[], %c32 : memref<!amdgpu.ds_barrier_state, #gpu.address_space<workgroup>>, i32
|
||
```
|
||
|
||
This operation is only available on gfx1250+.
|
||
}];
|
||
|
||
let assemblyFormat = [{
|
||
$base `[` $indices `]` `,` $participants attr-dict `:` type($base) `,` type($participants)
|
||
}];
|
||
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
// Promises IndexedAccessOpInterface.
|
||
def AMDGPU_DsBarrierPollStateOp :
|
||
AMDGPU_Op<"ds_barrier_poll_state">,
|
||
Arguments<(ins Arg<MemRefOf<[AMDGPU_DsBarrierStateType]>, "barrier(s)",
|
||
[MemRead, MemWrite]>:$base,
|
||
Variadic<Index>:$indices)>,
|
||
Results<(outs AMDGPU_DsBarrierStateType:$out)> {
|
||
let summary = "Atomically read the state of an in-LDS barrier.";
|
||
let description = [{
|
||
Atomically read and return the state of the barrier at `base[indices...]`.
|
||
|
||
This will ultimately act like a `memref.load`, but this operation will ensure
|
||
that appropriate atomic orderings and syncscopes are set.
|
||
The indices must be non-negative and in-bounds for the corresponding
|
||
dimensions of `base`.
|
||
|
||
Example:
|
||
```mlir
|
||
%state = amdgpu.ds_barrier_poll_state %barrier[] : memref<!amdgpu.ds_barrier_state, #gpu.address_space<workgroup>> -> !amdgpu.ds_barrier_state
|
||
```
|
||
|
||
This operation is only available on gfx1250+.
|
||
}];
|
||
|
||
let assemblyFormat = [{
|
||
$base `[` $indices `]` attr-dict `:` type($base) `->` type($out)
|
||
}];
|
||
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
// Promises IndexedAccessOpInterface.
|
||
def AMDGPU_DsAsyncBarrierArriveOp :
|
||
AMDGPU_Op<"ds_async_barrier_arrive">,
|
||
Arguments<(ins Arg<MemRefOf<[AMDGPU_DsBarrierStateType]>, "barrier(s)",
|
||
[MemRead, MemWrite]>:$base,
|
||
Variadic<Index>:$indices)> {
|
||
let summary = "Asynchronously arrive at an in-LDS barrier.";
|
||
let description = [{
|
||
Add a arrival to the LDS barrier at `base[indices]` to the sequence of pending
|
||
asynchronous memory operations.
|
||
|
||
The indices must be non-negative and in-bounds for the corresponding
|
||
dimensions of `base`.
|
||
|
||
This will add an "asynchronous memory operation" to the in-order list of pending
|
||
asynchronous loads from global memory to LDS. When the queue of such operations
|
||
issued before this operation is complete, the specified barrier will be arrived at,
|
||
decrementing the pending count by 1 **per lane that executes it** and rolling
|
||
over the phase if applicable.
|
||
|
||
This operation does not return the old barrier state.
|
||
|
||
Example:
|
||
```mlir
|
||
amdgpu.ds_async_barrier_arrive %barrier[] : memref<!amdgpu.ds_barrier_state, #gpu.address_space<workgroup>>
|
||
```
|
||
|
||
This operation is only available on gfx1250+.
|
||
}];
|
||
|
||
let assemblyFormat = [{
|
||
$base `[` $indices `]` attr-dict `:` type($base)
|
||
}];
|
||
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
// Promises IndexedAccessOpInterface.
|
||
def AMDGPU_DsBarrierArriveOp :
|
||
AMDGPU_Op<"ds_barrier_arrive">,
|
||
Arguments<(ins Arg<MemRefOf<[AMDGPU_DsBarrierStateType]>, "barrier(s)",
|
||
[MemRead, MemWrite]>:$base,
|
||
Variadic<Index>:$indices,
|
||
I64:$count)>,
|
||
Results<(outs AMDGPU_DsBarrierStateType:$out)> {
|
||
let summary = "Arrive at an in-LDS barrier and return old state.";
|
||
let description = [{
|
||
Atomically arrive at the LDS barrier at `base[indices]` and decrement it by `count`,
|
||
rolling over the phase if needed and returning the old barrier state.
|
||
|
||
The indices must be non-negative and in-bounds for the corresponding
|
||
dimensions of `base`.
|
||
|
||
`count` is the number of participants that should be subtracted from the barrier's
|
||
pending count **per lane that executes the operation**.
|
||
|
||
Example:
|
||
```mlir
|
||
%old_state = amdgpu.ds_barrier_arrive %barrier[], %c1 : memref<!amdgpu.ds_barrier_state, #gpu.address_space<workgroup>>, i64 -> !amdgpu.ds_barrier_state
|
||
```
|
||
|
||
This operation is only available on gfx1250+.
|
||
}];
|
||
|
||
let assemblyFormat = [{
|
||
$base `[` $indices `]` `,` $count attr-dict `:` type($base) `,` type($count) `->` type($out)
|
||
}];
|
||
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
def AMDGPU_DsBarrierStatePhaseOp :
|
||
AMDGPU_Op<"ds_barrier_state_phase", [Pure]>,
|
||
Arguments<(ins AMDGPU_DsBarrierStateType:$state)>,
|
||
Results<(outs I32:$res)> {
|
||
let summary = "Extract the phase of a barrier state.";
|
||
let description = [{
|
||
Extract the phase of the `!amdgpu.ds_barrier_state` `state` as a 32-bit value.
|
||
|
||
Example:
|
||
```mlir
|
||
%phase = amdgpu.ds_barrier_state_phase %state : !amdgpu.ds_barrier_state -> i32
|
||
```
|
||
}];
|
||
|
||
let assemblyFormat = [{
|
||
$state attr-dict `:` type($state) `->` type($res)
|
||
}];
|
||
}
|
||
|
||
def AMDGPU_DsBarrierStatePendingCountOp :
|
||
AMDGPU_Op<"ds_barrier_state_pending_count", [Pure]>,
|
||
Arguments<(ins AMDGPU_DsBarrierStateType:$state)>,
|
||
Results<(outs I32:$res)> {
|
||
let summary = "Extract the pending count of a barrier state.";
|
||
let description = [{
|
||
Extract the pending count of the `!amdgpu.ds_barrier_state` `state` as a 32-bit value.
|
||
|
||
Example:
|
||
```mlir
|
||
%pending = amdgpu.ds_barrier_state_pending_count %state : !amdgpu.ds_barrier_state -> i32
|
||
```
|
||
}];
|
||
|
||
let assemblyFormat = [{
|
||
$state attr-dict `:` type($state) `->` type($res)
|
||
}];
|
||
}
|
||
|
||
def AMDGPU_DsBarrierStateInitCountOp :
|
||
AMDGPU_Op<"ds_barrier_state_init_count", [Pure]>,
|
||
Arguments<(ins AMDGPU_DsBarrierStateType:$state)>,
|
||
Results<(outs I32:$res)> {
|
||
let summary = "Extract the init count of a barrier state.";
|
||
let description = [{
|
||
Extract the init count of the `!amdgpu.ds_barrier_state` `state` as a 32-bit value.
|
||
|
||
Example:
|
||
```mlir
|
||
%init = amdgpu.ds_barrier_state_init_count %state : !amdgpu.ds_barrier_state -> i32
|
||
```
|
||
}];
|
||
|
||
let assemblyFormat = [{
|
||
$state attr-dict `:` type($state) `->` type($res)
|
||
}];
|
||
}
|
||
|
||
def AMDGPU_DsBarrierStatePhaseParity :
|
||
AMDGPU_Op<"ds_barrier_state_phase_parity", [Pure]>,
|
||
Arguments<(ins AMDGPU_DsBarrierStateType:$state)>,
|
||
Results<(outs I1:$res)> {
|
||
let summary = "Extract the phase parity of a barrier state.";
|
||
let description = [{
|
||
Return the parity of the phase of the `!amdgpu.ds_barrier_state` `state`.
|
||
|
||
This is intended to simplify the case where the barrier is being used to repeatedly
|
||
track completion of a task where the precise value of the phase won't mater, only that
|
||
it changed since (or as a result of) the arrival.
|
||
|
||
Example:
|
||
```mlir
|
||
%parity = amdgpu.ds_barrier_state_phase_parity %state : !amdgpu.ds_barrier_state -> i1
|
||
```
|
||
}];
|
||
|
||
let assemblyFormat = [{
|
||
$state attr-dict `:` type($state) `->` type($res)
|
||
}];
|
||
}
|
||
|
||
def AMDGPU_GlobalPrefetchOp :
|
||
AMDGPU_Op<"global_prefetch", [MemoryEffects<[MemWrite, MemRead]>]>,
|
||
Arguments<(ins AnyMemRef:$src,
|
||
Variadic<I64>:$indices,
|
||
AMDGPU_LoadTemporalHintAttr:$temporalHint,
|
||
AMDGPU_CacheScopeAttr:$cacheScope,
|
||
UnitAttr:$speculative)>,
|
||
Results<(outs)> {
|
||
|
||
let summary = "Prefetch data to caches.";
|
||
let description = [{
|
||
Prefetches a cache line to high-level caches using the aligned address of
|
||
the source `memref` and an offset provided by the indices of the element
|
||
containing the cache line. This provides temporal hints (e.g., regular
|
||
or high-priority). Note that out-of-bounds access is allowed in
|
||
speculative mode. The provided memref must be in the global address space
|
||
(`#gpu.address_space<global>` or 1).
|
||
|
||
This operation was introduced in gfx1250.
|
||
|
||
Example:
|
||
```mlir
|
||
amdgpu.global_prefetch %src[%i, %j] RT SE speculative : memref<64x64xf16, #gpu.address_space<global>>
|
||
```
|
||
}];
|
||
|
||
let assemblyFormat = [{
|
||
$src `[` $indices `]` $temporalHint $cacheScope (`speculative` $speculative^)? attr-dict `:` qualified(type($src))
|
||
}];
|
||
|
||
let hasVerifier = 1;
|
||
}
|
||
|
||
#endif // MLIR_DIALECT_AMDGPU_IR_AMDGPUOPS_TD
|