Beaver.MLIR.Dialect.NVGPU (beaver v0.4.8)

Copy Markdown

Summary

Functions

Return op name nvgpu.device_async_copy as a bitstring.

nvgpu.device_async_copy - device-side asynchronous copy

Return op name nvgpu.device_async_create_group as a bitstring.

nvgpu.device_async_create_group - device side asynchronous create group operation

Return op name nvgpu.device_async_wait as a bitstring.

nvgpu.device_async_wait - Wait for async gpu ops to complete.

Return op name nvgpu.extf as a bitstring.

nvgpu.extf - Extend floating-point to wider floating-point

Return op name nvgpu.ldmatrix as a bitstring.

nvgpu.ldmatrix

Return op name nvgpu.mbarrier.arrive as a bitstring.

nvgpu.mbarrier.arrive - Performs arrive operation on the nvgpu.mbarrier.arrive.

Return op name nvgpu.mbarrier.arrive.expect_tx as a bitstring.

nvgpu.mbarrier.arrive.expect_tx - Performs expect_tx operation on the nvgpu.mbarrier.arrive

Return op name nvgpu.mbarrier.arrive.nocomplete as a bitstring.

nvgpu.mbarrier.arrive.nocomplete - Performs arrive operation on the nvgpu.mbarrier.arrive.nocomplete as non-blocking.

Return op name nvgpu.mbarrier.create as a bitstring.

nvgpu.mbarrier.create - Creates a nvgpu.mbarrier object.

Return op name nvgpu.mbarrier.get as a bitstring.

nvgpu.mbarrier.get - Return a pointer to an nvgpu.mbarrier.

Return op name nvgpu.mbarrier.init as a bitstring.

nvgpu.mbarrier.init - Initialize the nvgpu.mbarrier.

Return op name nvgpu.mbarrier.test.wait as a bitstring.

nvgpu.mbarrier.test.wait - Checks if the nvgpu.mbarrier has completed its current phase.

Return op name nvgpu.mbarrier.try_wait.parity as a bitstring.

nvgpu.mbarrier.try_wait.parity - Waits for the nvgpu.mbarrier to complete its current phase.

Return op name nvgpu.mma.sp.sync as a bitstring.

nvgpu.mma.sp.sync

Return op name nvgpu.mma.sync as a bitstring.

nvgpu.mma.sync

Return op name nvgpu.rcp as a bitstring.

nvgpu.rcp - The reciprocal calculation for vector types

Return op name nvgpu.tma.async.load as a bitstring.

nvgpu.tma.async.load - TMA asynchronous load

Return op name nvgpu.tma.async.store as a bitstring.

nvgpu.tma.async.store - TMA asynchronous store

Return op name nvgpu.tma.create.descriptor as a bitstring.

nvgpu.tma.create.descriptor - TMA create descriptor

Return op name nvgpu.tma.fence.descriptor as a bitstring.

nvgpu.tma.fence.descriptor - Insert fence given nvgpu.tensormap.descriptor

Return op name nvgpu.tma.prefetch.descriptor as a bitstring.

nvgpu.tma.prefetch.descriptor - Prefetch given nvgpu.tensormap.descriptor

Return op name nvgpu.truncf as a bitstring.

nvgpu.truncf - Truncate floating-point to narrower floating-point

Return op name nvgpu.warpgroup.generate.descriptor as a bitstring.

nvgpu.warpgroup.generate.descriptor - Generate a warpgroup matrix descriptor

Return op name nvgpu.warpgroup.mma as a bitstring.

nvgpu.warpgroup.mma

Return op name nvgpu.warpgroup.mma.init.accumulator as a bitstring.

nvgpu.warpgroup.mma.init.accumulator - Initializes the accumulator matrix

Return op name nvgpu.warpgroup.mma.store as a bitstring.

nvgpu.warpgroup.mma.store

Functions

device_async_copy()

Return op name nvgpu.device_async_copy as a bitstring.

device_async_copy(ssa)

nvgpu.device_async_copy - device-side asynchronous copy

This op has support for result type inference.

Attributes

  • dstElements - Single, IndexAttr, index attribute
  • bypassL1 - Optional, UnitAttr, unit attribute

Operands

  • dst - Single, AnyMemRef, memref of any non-token type values
  • dstIndices - Variadic, Index, variadic of index
  • src - Single, AnyMemRef, memref of any non-token type values
  • srcIndices - Variadic, Index, variadic of index
  • srcElements - Optional, Index, index

Results

  • asyncToken - Single, NVGPU_DeviceAsyncToken, device async token type

Description

The nvgpu.device_async_copy op initiates an asynchronous copy operation of elements from source (global memory) to the destination (shared memory) without blocking the thread. The async copy is added to a group.

This op is meant to be used with nvgpu.device_async_create_group and nvgpu.device_async_wait to synchronize copies as explained in those ops descriptions.

bypassL1 attribute is hint to the hardware to bypass the L1 cache during async copy, this hint may be ignored by the hardware.

dstElements attribute is the total number of elements written to destination (shared memory).

srcElements argument is the total number of elements read from source (global memory).

srcElements is an optional argument and when present the op only reads srcElements number of elements from the source (global memory) and zero fills the rest of the elements in the destination (shared memory).

In order to do a copy and wait for the result we need the following combination:

// copy 1.
%cp1 = nvgpu.device_async_copy %A[%c0], %B[%c0], 4 :memref<16xf32> to memref<16xf32, 3>
// copy 2.
%cp2 = nvgpu.device_async_copy %C[%c0], %D[%c0], 4 : memref<16xf32> to memref<16xf32, 3>
// group 1 contains copy 1 and copy 2.
%token1 = nvgpu.device_async_create_group %cp1, %cp2
// copy 3.
%cp3 = nvgpu.device_async_copy %E[%c0], %F[%c0], 4 : memref<16xf32> to memref<16xf32, 3>
// group 2 contains copy 3.
%token2 = nvgpu.device_async_create_group %cp3
// after the wait copy 1 and copy 2 are complete.
nvgpu.device_async_wait %token1
// after the wait copy 3 is complete.
nvgpu.device_async_wait %token2

Example:

%0 = nvgpu.device_async_copy %src[%c0, %c0], %dst[%c0, %c0, %c0], 4 :
  memref<4x5xf32> to memref<2x7x5xf32, 3>

device_async_create_group()

Return op name nvgpu.device_async_create_group as a bitstring.

device_async_create_group(ssa)

nvgpu.device_async_create_group - device side asynchronous create group operation

This op has support for result type inference.

Operands

  • inputTokens - Variadic, NVGPU_DeviceAsyncToken, variadic of device async token type

Results

  • asyncToken - Single, NVGPU_DeviceAsyncToken, device async token type

Description

The nvgpu.device_async_create_group op creates a group of memory accesses containing all the pending device_async_copy operations associated with argument tokens. Each token can only be part of one group.

It returns a token that can be use to wait until the group fully completes.

This is meant to be used with nvgpu.device_async_wait to synchronize copies as explained in those ops descriptions.

Groups are executed in the order they are created.

Example:

  %0 = nvgpu.device_async_create_group

device_async_wait()

Return op name nvgpu.device_async_wait as a bitstring.

device_async_wait(ssa)

nvgpu.device_async_wait - Wait for async gpu ops to complete.

Attributes

  • numGroups - Optional, I32Attr, 32-bit signless integer attribute

Operands

  • asyncDependencies - Single, NVGPU_DeviceAsyncToken, device async token type

Description

The nvgpu.device_async_wait op will block the execution thread until the group associated with the source token is fully completed.

The optional $numGroups attribute gives an upper bound of the number of groups uncompleted when the wait can unblock the thread. For example, if 16 async groups are pushe and $numGroups is set to 12, then the thread will unblock when 12 groups or fewer are in flight (4 groups have completed).

Example:

nvgpu.device_async_wait %0

extf()

Return op name nvgpu.extf as a bitstring.

extf(ssa)

nvgpu.extf - Extend floating-point to wider floating-point

Attributes

  • rnd - Single, FPRoundingModeAttr, NVVM FPRoundingMode kind
  • relu - Single, BoolAttr, bool attribute

Operands

  • in - Single, anonymous/composite constraint, floating-point or vector of floating-point values

Results

  • out - Single, anonymous/composite constraint, floating-point or vector of floating-point values

Description

Extend a floating-point value to a wider floating-point type. Destination must be strictly wider than source.

Supported paths: f4 (f4E2M1FN) -> f16, bf16 f6 (f6E2M3FN, f6E3M2FN) -> f16, bf16 f8 (f8E4M3FN, f8E5M2, f8E8M0FNU) -> f16, bf16 f16 -> f32 bf16 -> f32

Example:

%r = nvgpu.extf %in : vector<8xf8E4M3FN> to vector<8xf16>
%r = nvgpu.extf %in : vector<8xf6E3M2FN> to vector<8xf16>
%r = nvgpu.extf %in : vector<4xf8E5M2> to vector<4xf32>
%r = nvgpu.extf %in : f8E4M3FN to f32

ldmatrix()

Return op name nvgpu.ldmatrix as a bitstring.

ldmatrix(ssa)

nvgpu.ldmatrix

Attributes

  • transpose - Single, BoolAttr, bool attribute
  • numTiles - Single, I32Attr, 32-bit signless integer attribute

Operands

  • srcMemref - Single, AnyMemRef, memref of any non-token type values
  • indices - Variadic, Index, variadic of index

Results

  • res - Single, AnyVectorOfNonZeroRank, vector of any non-token type values

Description

The nvgpu.ldmatrix op represents loading a matrix fragment from memory to registers. The source and result type must be compatible with lowering to the nvvm.ldmatrix instruction. This op represents the distributed version of a vector.transfer_read as an intermediate step between lowering from vector.transfer_read to nvvm.ldmatrix.

This operation is meant to follow the semantic of described here: https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-ldmatrix

Example:

%0 = nvgpu.ldmatrix %sm[%c0, %c0] {numTiles = 4 : i32, transpose = false} :
  memref<?x?xf16, 3> -> vector<4x2xf16>

mbarrier_arrive()

Return op name nvgpu.mbarrier.arrive as a bitstring.

mbarrier_arrive(ssa)

nvgpu.mbarrier.arrive - Performs arrive operation on the nvgpu.mbarrier.arrive.

This op has support for result type inference.

Operands

  • barriers - Single, NVGPU_MBarrierGroup, mbarrier barrier type
  • mbarId - Single, Index, index

Results

  • token - Single, NVGPU_MBarrierToken,

Description

The Op performs arrive-on operation on the mbarrier object and returns a nvgpu.mbarrier.token.

For more information, see https://docs.nvidia.com/cuda/parallel-thread-execution/#arrive-on-operation-on-mbarrier-object

Example:

  %token = nvgpu.mbarrier.arrive %barrier : !nvgpu.mbarrier.barrier<memorySpace = #gpu.address_space<workgroup>> -> !nvgpu.mbarrier.token

mbarrier_arrive_expect_tx()

Return op name nvgpu.mbarrier.arrive.expect_tx as a bitstring.

mbarrier_arrive_expect_tx(ssa)

nvgpu.mbarrier.arrive.expect_tx - Performs expect_tx operation on the nvgpu.mbarrier.arrive

Operands

  • barriers - Single, NVGPU_MBarrierGroup, mbarrier barrier type
  • txcount - Single, Index, index
  • mbarId - Single, Index, index
  • predicate - Optional, I1, 1-bit signless integer

Description

A thread executing the Op performs an expect-tx operation on the mbarrier object at the location specified by the address operand $barrier. The expect-tx operation, with an $txcount argument, increases the tx-count of an mbarrier object by the value specified by $txcount. This makes the current phase of the mbarrier object to expect and track the completion of additional asynchronous transactions.

The $txCount specifies the number of element to the expect-tx operation.

Example:

  nvgpu.mbarrier.arrive.expect_tx %barrier, %ic0 : !nvgpu.mbarrier.barrier<memorySpace = #gpu.address_space<workgroup>>

mbarrier_arrive_nocomplete()

Return op name nvgpu.mbarrier.arrive.nocomplete as a bitstring.

mbarrier_arrive_nocomplete(ssa)

nvgpu.mbarrier.arrive.nocomplete - Performs arrive operation on the nvgpu.mbarrier.arrive.nocomplete as non-blocking.

This op has support for result type inference.

Operands

  • barriers - Single, NVGPU_MBarrierGroup, mbarrier barrier type
  • mbarId - Single, Index, index
  • count - Single, Index, index

Results

  • token - Single, NVGPU_MBarrierToken,

Description

The Op performs arrive-on operation on the mbarrier object and returns a nvgpu.mbarrier.token.

The Op does not cause the nvgpu.mbarrier to complete its current phase.

Example:

  %token = nvgpu.mbarrier.arrive.noComplete %barrier, %count : !nvgpu.mbarrier.barrier<memorySpace = #gpu.address_space<workgroup>> -> !nvgpu.mbarrier.token

mbarrier_create()

Return op name nvgpu.mbarrier.create as a bitstring.

mbarrier_create(ssa)

nvgpu.mbarrier.create - Creates a nvgpu.mbarrier object.

Results

  • barriers - Single, NVGPU_MBarrierGroup, mbarrier barrier type

Description

The Op generates one or more mbarrier object, which is a barrier created in shared memory and supports various synchronization behaviors for threads.

The mbarrier object has the following type and alignment requirements: Type: .b64, Alignment: 8, Memory space: .shared

Example:

  %barrier = nvgpu.mbarrier.create -> !nvgpu.mbarrier.barrier<memorySpace = #gpu.address_space<workgroup>>

mbarrier_get()

Return op name nvgpu.mbarrier.get as a bitstring.

mbarrier_get(ssa)

nvgpu.mbarrier.get - Return a pointer to an nvgpu.mbarrier.

Operands

  • barriers - Single, NVGPU_MBarrierGroup, mbarrier barrier type
  • mbarId - Single, Index, index

Results

  • mbarrierPointer - Single, anonymous/composite constraint, 32-bit signless integer or 64-bit signless integer

Description

The nvgpu.mbarrier.get operation retrieves a pointer to a specific mbarrier object from a group of barriers created by the nvgpu.mbarrier.create operation.

Example:

  %mbars = nvgpu.mbarrier.create -> !nvgpu.mbarrier.group<memorySpace = #gpu.address_space<workgroup>, num_barriers = 10>
  %mbar_pointer = nvgpu.mbarrier.get %mbars[%c2] : !nvgpu.mbarrier.group<memorySpace = #gpu.address_space<workgroup>>

mbarrier_init()

Return op name nvgpu.mbarrier.init as a bitstring.

mbarrier_init(ssa)

nvgpu.mbarrier.init - Initialize the nvgpu.mbarrier.

Operands

  • barriers - Single, NVGPU_MBarrierGroup, mbarrier barrier type
  • count - Single, Index, index
  • mbarId - Single, Index, index
  • predicate - Optional, I1, 1-bit signless integer

Description

The Op initializes the mbarrier object with the given number of threads.

Example:

  %num_threads = gpu.block_dim x
  %barrier = nvgpu.mbarrier.create -> !nvgpu.mbarrier.barrier<memorySpace = #gpu.address_space<workgroup>>
  nvgpu.mbarrier.init %barrier, %num_threads : !nvgpu.mbarrier.barrier<memorySpace = #gpu.address_space<workgroup>>

mbarrier_test_wait()

Return op name nvgpu.mbarrier.test.wait as a bitstring.

mbarrier_test_wait(ssa)

nvgpu.mbarrier.test.wait - Checks if the nvgpu.mbarrier has completed its current phase.

This op has support for result type inference.

Operands

  • barriers - Single, NVGPU_MBarrierGroup, mbarrier barrier type
  • token - Single, NVGPU_MBarrierToken,
  • mbarId - Single, Index, index

Results

  • waitComplete - Single, I1, 1-bit signless integer

Description

Checks whether the mbarrier object has completed the phase. It is a non-blocking instruction which tests for the completion of the phase.

Example:

  %isComplete = nvgpu.mbarrier.test.wait %barrier, %token : !nvgpu.mbarrier.barrier<memorySpace = #gpu.address_space<workgroup>>, !nvgpu.mbarrier.token

mbarrier_try_wait_parity()

Return op name nvgpu.mbarrier.try_wait.parity as a bitstring.

mbarrier_try_wait_parity(ssa)

nvgpu.mbarrier.try_wait.parity - Waits for the nvgpu.mbarrier to complete its current phase.

Operands

  • barriers - Single, NVGPU_MBarrierGroup, mbarrier barrier type
  • phaseParity - Single, I1, 1-bit signless integer
  • ticks - Single, Index, index
  • mbarId - Single, Index, index

Description

Checks whether the mbarrier object has completed the phase. It is a potentially blocking instruction which tests for the completion of the phase. Suspended thread resumes execution when the specified phase completes OR before the phase completes following a system-dependent time limit.

The $phaseParity specifies either even phase (0) or odd phase (1) to wait.

Example:

  nvgpu.mbarrier.try_wait.parity %barrier, %phaseParity, %ticks : !nvgpu.mbarrier.barrier<memorySpace = #gpu.address_space<workgroup>>

mma_sp_sync()

Return op name nvgpu.mma.sp.sync as a bitstring.

mma_sp_sync(ssa)

nvgpu.mma.sp.sync

Attributes

  • mmaShape - Single, I64ArrayAttr, 64-bit integer array attribute
  • sparsitySelector - Single, I32Attr, 32-bit signless integer attribute
  • tf32Enabled - Optional, UnitAttr, unit attribute

Operands

  • matrixA - Single, AnyVectorOfNonZeroRank, vector of any non-token type values
  • matrixB - Single, AnyVectorOfNonZeroRank, vector of any non-token type values
  • matrixC - Single, AnyVectorOfNonZeroRank, vector of any non-token type values
  • sparseMetadata - Single, NVGPU_MmaSparseSyncMetadataType, fixed-length vector of 16-bit signless integer values of length 2

Results

  • res - Single, AnyVectorOfNonZeroRank, vector of any non-token type values

Description

The nvgu.mma.sp.sync operation performs a warp-distributed MMA operation where operand A is "structured sparse". In this case, the matrixA operand represents the (warp-distributed) non-zero values of operand A, and the sparse_metadata operand provides the indices.

The full description of the sparsity storage format and distribution scheme is described in the PTX docs. This operation is meant to follow the semantic described in the PTX documentation here: https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-for-sparse-mma

The way the indices are distributed among the threads in a warp is controlled by the optional sparsity_selector operand, which is 0 by default. For more information, please consult the PTX documentation linked above.

Example (targetingthe f16 16x8x32 mma.sp PTX instruction):

nvgpu.mma.sp.sync (%a, %b, %c) metadata (%meta) {mmaShape = [16, 8, 32]} :
  (vector<4x2xf16>, vector<2x2xf16>, vector<2x2xf16>) -> vector<2x2xf16>

mma_sync()

Return op name nvgpu.mma.sync as a bitstring.

mma_sync(ssa)

nvgpu.mma.sync

Attributes

  • mmaShape - Single, I64ArrayAttr, 64-bit integer array attribute
  • tf32Enabled - Optional, UnitAttr, unit attribute

Operands

  • matrixA - Single, AnyVectorOfNonZeroRank, vector of any non-token type values
  • matrixB - Single, AnyVectorOfNonZeroRank, vector of any non-token type values
  • matrixC - Single, AnyVectorOfNonZeroRank, vector of any non-token type values

Results

  • res - Single, AnyVectorOfNonZeroRank, vector of any non-token type values

Description

The nvgpu.mma.sync op represents the warp-level matrix-multiply-and- accumulate (mma) operation that is compatible with nvvm.mma.sync. The operands and results vector sizes are thread-level onwership to the warp-level mma operation shape. mmaShape attribute holds the warp-level matrix-multiply shape.

The nvgpu.mma.sync op serves as an intermediate point between lowering from vector.contract to nvvm.mma.sync.

This operation is meant to follow the semantic of described here: https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-mma

Example:

%res = nvgpu.mma.sync (%matrixA, %matrixB, %matrixC) {mmaShape = [16, 8, 16]} :
    (vector<4x2xf16>, vector<2x2xf16>, vector<2x2xf32>) -> vector<2x2xf32>

rcp()

Return op name nvgpu.rcp as a bitstring.

rcp(ssa)

nvgpu.rcp - The reciprocal calculation for vector types

This op has support for result type inference.

Attributes

  • rounding - Single, FPRoundingModeAttr, NVVM FPRoundingMode kind
  • approx - Single, BoolAttr, bool attribute
  • ftz - Single, BoolAttr, bool attribute

Operands

  • in - Single, anonymous/composite constraint, vector of 32-bit float values

Results

  • out - Single, anonymous/composite constraint, vector of 32-bit float values

Description

Reciprocal calculation for vector types using nvvm.rcp OPs.

Currently, only the approx rounding mode and ftz are supported, and only for the f32 type.

The input and output must be of the same vector type and shape.

tma_async_load()

Return op name nvgpu.tma.async.load as a bitstring.

tma_async_load(ssa)

nvgpu.tma.async.load - TMA asynchronous load

Operands

  • dst - Single, AnyMemRef, memref of any non-token type values
  • barriers - Single, NVGPU_MBarrierGroup, mbarrier barrier type
  • tensorMapDescriptor - Single, NVGPU_TensorMapDescriptor, TensorMap descriptor
  • coordinates - Variadic, Index, variadic of index
  • mbarId - Single, Index, index
  • multicastMask - Optional, I16, 16-bit signless integer
  • predicate - Optional, I1, 1-bit signless integer

Description

The Op loads a tile memory region from global memory to shared memory by Tensor Memory Access (TMA).

$tensorMapDescriptor is tensor map descriptor which has information about tile shape. The descriptor is created by nvgpu.tma.create.descriptor

The Op uses $barrier mbarrier based completion mechanism.

tma_async_store()

Return op name nvgpu.tma.async.store as a bitstring.

tma_async_store(ssa)

nvgpu.tma.async.store - TMA asynchronous store

Operands

  • src - Single, AnyMemRef, memref of any non-token type values
  • tensorMapDescriptor - Single, NVGPU_TensorMapDescriptor, TensorMap descriptor
  • coordinates - Variadic, Index, variadic of index
  • predicate - Optional, I1, 1-bit signless integer

Description

The Op store a tile memory region from global memory to shared memory by Tensor Memory Access (TMA).

$tensorMapDescriptor is tensor map descriptor which has information about tile shape. The descriptor is created by nvgpu.tma.create.descriptor

tma_create_descriptor()

Return op name nvgpu.tma.create.descriptor as a bitstring.

tma_create_descriptor(ssa)

nvgpu.tma.create.descriptor - TMA create descriptor

Operands

  • tensor - Single, AnyUnrankedMemRef, unranked.memref of any non-token type values
  • boxDimensions - Variadic, Index, variadic of index

Results

  • tensorMap - Single, NVGPU_TensorMapDescriptor, TensorMap descriptor

Description

The Op creates a tensor map descriptor object representing tiled memory region. To do that it calls CUDA Driver's cuTensorMapEncodeTiled. The descriptor is used by Tensor Memory Access (TMA).

The tensor is the source tensor to be tiled.

The boxDimensions is the size of the tiled memory region in each dimension.

For more information see below: https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__TENSOR__MEMORY.html

tma_fence_descriptor()

Return op name nvgpu.tma.fence.descriptor as a bitstring.

tma_fence_descriptor(ssa)

nvgpu.tma.fence.descriptor - Insert fence given nvgpu.tensormap.descriptor

Operands

  • tensorMapDescriptor - Single, NVGPU_TensorMapDescriptor, TensorMap descriptor

Description

The Op fences the given $tmaDescriptor. This is necessary if the tensor map descriptor was modified from the host using cudaMemcpy. In this case, the kernel needs a fence after which it is safe to use tensor.map.

tma_prefetch_descriptor()

Return op name nvgpu.tma.prefetch.descriptor as a bitstring.

tma_prefetch_descriptor(ssa)

nvgpu.tma.prefetch.descriptor - Prefetch given nvgpu.tensormap.descriptor

Operands

  • tensorMapDescriptor - Single, NVGPU_TensorMapDescriptor, TensorMap descriptor
  • predicate - Optional, I1, 1-bit signless integer

Description

The Op brings the cache line containing the given $tmaDescriptor for subsequent use by the tma.async.load instruction.

truncf()

Return op name nvgpu.truncf as a bitstring.

truncf(ssa)

nvgpu.truncf - Truncate floating-point to narrower floating-point

Attributes

  • rnd - Single, FPRoundingModeAttr, NVVM FPRoundingMode kind
  • sat - Single, SaturationModeAttr, Describes the saturation mode whose value is one of {none, satfinite}
  • relu - Single, BoolAttr, bool attribute

Operands

  • in - Single, anonymous/composite constraint, floating-point or vector of floating-point values
  • random_bits - Optional, I32, 32-bit signless integer

Results

  • out - Single, anonymous/composite constraint, floating-point or vector of floating-point values

Description

Truncate a floating-point value to a smaller floating-point type. Destination must be strictly narrower than source.

Supported paths: f16 -> f8 (f8E4M3FN, f8E5M2, f8E8M0FNU), f6 (f6E2M3FN, f6E3M2FN),

      f4 (f4E2M1FN)

bf16 -> f8 (f8E4M3FN, f8E5M2, f8E8M0FNU), f6 (f6E2M3FN, f6E3M2FN),

      f4 (f4E2M1FN)

f32 -> f16, bf16, f8 (f8E4M3FN, f8E5M2, f8E8M0FNU),

      f6 (f6E2M3FN, f6E3M2FN), f4 (f4E2M1FN)

f64 -> f32, f16, bf16

The random_bits operand enables stochastic rounding (RS mode) for f32->f16/bf16 conversions; when provided, rnd must be RS.

Example:

%r = nvgpu.truncf %in : vector<8xf32> to vector<8xf8E4M3FN>
%r = nvgpu.truncf %in : vector<8xf32> to vector<8xf6E2M3FN>
%r = nvgpu.truncf %in : f32 to f16
%r = nvgpu.truncf %in : vector<2x4xf16> to vector<2x4xf8E5M2>

warpgroup_generate_descriptor()

Return op name nvgpu.warpgroup.generate.descriptor as a bitstring.

warpgroup_generate_descriptor(ssa)

nvgpu.warpgroup.generate.descriptor - Generate a warpgroup matrix descriptor

Operands

  • tensor - Single, AnyMemRef, memref of any non-token type values
  • tensorMap - Single, NVGPU_TensorMapDescriptor, TensorMap descriptor

Results

  • descriptor - Single, NVGPU_WarpgroupMatrixDescriptor, Warpgroup matrix descriptor type

Description

This Op builds a nvgpu.warpgroup.descriptor that is used by nvgpu.warpgroup.mma to perform warpgroup-level matrix multiply and accumulate.

The descriptor specifies the properties of the matrix in shared memory that is a multiplicand in the matrix multiply and accumulate operation.

warpgroup_mma()

Return op name nvgpu.warpgroup.mma as a bitstring.

warpgroup_mma(ssa)

nvgpu.warpgroup.mma

Attributes

  • waitGroup - Optional, I64Attr, 64-bit signless integer attribute
  • transposeA - Optional, UnitAttr, unit attribute
  • transposeB - Optional, UnitAttr, unit attribute

Operands

  • descriptorA - Single, NVGPU_WarpgroupMatrixDescriptor, Warpgroup matrix descriptor type
  • descriptorB - Single, NVGPU_WarpgroupMatrixDescriptor, Warpgroup matrix descriptor type
  • matrixC - Single, NVGPU_WarpgroupAccumulator,

Results

  • matrixD - Single, NVGPU_WarpgroupAccumulator,

Description

The nvgpu.warpgroup.mma op performs the warpgroup-level (4 warps) matrix-multiply-and-accumulate (mma) operation that results in nvvm.wgmma.mma_async.

The operands are descriptorA and descriptorB that are wgmma matrix descriptors that shows the properties of the matrix in shared memory. The results are thread-level ownership to the warpgroup-level mma operation shape. The shape is deduced from the descriptor types and output vector.

The Op encapsulates multiple nvvm.wgmma.mma_async operations to complete the given shape. As nvvm.wgmma.async Op, or its corresponding PTX instruction, is asynchronous, this Op groups the nvvm.wgmma.async and surrounds them between wgmma.fence.aligned and wgmma.commit.group.sync.aligned, wgmma.wait.group.sync.aligned Ops.

Example:

  %r1,%r2 = nvgpu.warpgroup.mma %descA, %descB, %acc1, %acc2: 
             !nvgpu.warpgroup.descriptor<tensor = memref<128x64xf16, 3>>, 
             !nvgpu.warpgroup.descriptor<tensor = memref<64x128xf16, 3>>, 
             !nvgpu.warpgroup.accumulator<fragmented = vector<64x128xf32>>,
             !nvgpu.warpgroup.accumulator<fragmented = vector<64x128xf32>>
             -> 
             !nvgpu.warpgroup.accumulator<fragmented = vector<64x128xf32>>,
             !nvgpu.warpgroup.accumulator<fragmented = vector<64x128xf32>>

warpgroup_mma_init_accumulator()

Return op name nvgpu.warpgroup.mma.init.accumulator as a bitstring.

warpgroup_mma_init_accumulator(ssa)

nvgpu.warpgroup.mma.init.accumulator - Initializes the accumulator matrix

Results

  • matrixC - Single, NVGPU_WarpgroupAccumulator,

Description

This Op generates and initializes the accumulator matrix for nvgpu.warpgroup.mma op to perform matrix-multiply-and-accumulate.

warpgroup_mma_store()

Return op name nvgpu.warpgroup.mma.store as a bitstring.

warpgroup_mma_store(ssa)

nvgpu.warpgroup.mma.store

Operands

  • matrixD - Single, NVGPU_WarpgroupAccumulator,
  • dstMemref - Single, AnyMemRef, memref of any non-token type values

Description

The nvgpu.warpgroup.mma.store op performs the store of fragmented result in $matrixD to given memref.

[See the details of register fragment layout for accumulator matrix D] (https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#wgmma-64n16-d)

Note that, the op must be run with warp group.