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

This module defines functions for Ops in GPU dialect.

Summary

Functions

all_reduce()

Return op name gpu.all_reduce as a bitstring.

all_reduce(ssa)

gpu.all_reduce - Reduce values among workgroup.

alloc()

Return op name gpu.alloc as a bitstring.

alloc(ssa)

gpu.alloc - GPU memory allocation operation.

ballot()

Return op name gpu.ballot as a bitstring.

ballot(ssa)

gpu.ballot - Collects predicate values from all threads in a subgroup.

barrier()

Return op name gpu.barrier as a bitstring.

barrier(ssa)

gpu.barrier - Synchronizes work items within an execution scope.

binary()

Return op name gpu.binary as a bitstring.

binary(ssa)

gpu.binary - An Op for storing serialized GPU binary objects.

block_dim()

Return op name gpu.block_dim as a bitstring.

block_dim(ssa)

gpu.block_dim

block_id()

Return op name gpu.block_id as a bitstring.

block_id(ssa)

gpu.block_id

cluster_block_id()

Return op name gpu.cluster_block_id as a bitstring.

cluster_block_id(ssa)

gpu.cluster_block_id

cluster_dim()

Return op name gpu.cluster_dim as a bitstring.

cluster_dim(ssa)

gpu.cluster_dim

cluster_dim_blocks()

Return op name gpu.cluster_dim_blocks as a bitstring.

cluster_dim_blocks(ssa)

gpu.cluster_dim_blocks

cluster_id()

Return op name gpu.cluster_id as a bitstring.

cluster_id(ssa)

gpu.cluster_id

container_module_attribute_name()

Get the name of the attribute used to annotate the modules that contain kernel modules.

create_2to4_spmat()

Return op name gpu.create_2to4_spmat as a bitstring.

create_2to4_spmat(ssa)

gpu.create_2to4_spmat - Create sparse matrix with 2:4 sparsity operation

create_bsr()

Return op name gpu.create_bsr as a bitstring.

create_bsr(ssa)

gpu.create_bsr - Create sparse matrix in BSR format operation

create_coo()

Return op name gpu.create_coo as a bitstring.

create_coo(ssa)

gpu.create_coo - Create sparse matrix in COO format operation

create_coo_aos()

Return op name gpu.create_coo_aos as a bitstring.

create_coo_aos(ssa)

gpu.create_coo_aos - Create sparse matrix in COO format operation (AoS)

create_csc()

Return op name gpu.create_csc as a bitstring.

create_csc(ssa)

gpu.create_csc - Create sparse matrix in CSC format operation

create_csr()

Return op name gpu.create_csr as a bitstring.

create_csr(ssa)

gpu.create_csr - Create sparse matrix in CSR format operation

create_dn_tensor()

Return op name gpu.create_dn_tensor as a bitstring.

create_dn_tensor(ssa)

gpu.create_dn_tensor - Create dense tensor operation

dealloc()

Return op name gpu.dealloc as a bitstring.

dealloc(ssa)

gpu.dealloc - GPU memory deallocation operation

destroy_dn_tensor()

Return op name gpu.destroy_dn_tensor as a bitstring.

destroy_dn_tensor(ssa)

gpu.destroy_dn_tensor - Destroy dense tensor operation

destroy_sp_mat()

Return op name gpu.destroy_sp_mat as a bitstring.

destroy_sp_mat(ssa)

gpu.destroy_sp_mat - Destroy sparse matrix operation

dynamic_shared_memory()

Return op name gpu.dynamic_shared_memory as a bitstring.

dynamic_shared_memory(ssa)

gpu.dynamic_shared_memory - Get the memref for dynamic shared memory

func()

Return op name gpu.func as a bitstring.

func(ssa)

gpu.func - Function executable on a GPU

global_id()

Return op name gpu.global_id as a bitstring.

global_id(ssa)

gpu.global_id

grid_dim()

Return op name gpu.grid_dim as a bitstring.

grid_dim(ssa)

gpu.grid_dim

host_register()

Return op name gpu.host_register as a bitstring.

host_register(ssa)

gpu.host_register - Registers a memref for access from device.

host_unregister()

Return op name gpu.host_unregister as a bitstring.

host_unregister(ssa)

gpu.host_unregister - Unregisters a memref for access from device.

initialize_named_barrier()

Return op name gpu.initialize_named_barrier as a bitstring.

initialize_named_barrier(ssa)

gpu.initialize_named_barrier - Initialize a named barrier with a member count.

lane_id()

Return op name gpu.lane_id as a bitstring.

lane_id(ssa)

gpu.lane_id

launch()

Return op name gpu.launch as a bitstring.

launch(ssa)

gpu.launch - GPU kernel launch operation

launch_func()

Return op name gpu.launch_func as a bitstring.

launch_func(ssa)

gpu.launch_func - Launches a function as a GPU kernel

memcpy()

Return op name gpu.memcpy as a bitstring.

memcpy(ssa)

gpu.memcpy - GPU memcpy operation

memset()

Return op name gpu.memset as a bitstring.

memset(ssa)

gpu.memset - GPU memset operation

module()

Return op name gpu.module as a bitstring.

module(ssa)

gpu.module - A top level compilation unit containing code to be run on a GPU.

num_subgroups()

Return op name gpu.num_subgroups as a bitstring.

num_subgroups(ssa)

gpu.num_subgroups

number_of_buffers_in_workgroup_attributions_attribute_name()

Returns the name of the attribute containing the number of buffers located in the workgroup memory.

printf()

Return op name gpu.printf as a bitstring.

printf(ssa)

gpu.printf - Device-side printf, as in CUDA or OpenCL, for debugging

return()

Return op name gpu.return as a bitstring.

return(ssa)

gpu.return - Terminator for GPU functions.

rotate()

Return op name gpu.rotate as a bitstring.

rotate(ssa)

gpu.rotate - Rotate values within a subgroup.

sddmm()

Return op name gpu.sddmm as a bitstring.

sddmm(ssa)

gpu.sddmm - SDDMM operation

sddmm_buffer_size()

Return op name gpu.sddmm_buffer_size as a bitstring.

sddmm_buffer_size(ssa)

gpu.sddmm_buffer_size - Precompute buffersize for SDDMM operation

set_csr_pointers()

Return op name gpu.set_csr_pointers as a bitstring.

set_csr_pointers(ssa)

gpu.set_csr_pointers - SpGEMM get size operation

set_default_device()

Return op name gpu.set_default_device as a bitstring.

set_default_device(ssa)

gpu.set_default_device - Set default GPU for operations after this by index

shuffle()

Return op name gpu.shuffle as a bitstring.

shuffle(ssa)

gpu.shuffle - Shuffles values within a subgroup.

spgemm_copy()

Return op name gpu.spgemm_copy as a bitstring.

spgemm_copy(ssa)

gpu.spgemm_copy - SpGEMM copy operation

spgemm_create_descr()

Return op name gpu.spgemm_create_descr as a bitstring.

spgemm_create_descr(ssa)

gpu.spgemm_create_descr - SpGEMM Create Descr operation

spgemm_destroy_descr()

Return op name gpu.spgemm_destroy_descr as a bitstring.

spgemm_destroy_descr(ssa)

gpu.spgemm_destroy_descr - SpGEMM Destroy Descr operation

spgemm_work_estimation_or_compute()

Return op name gpu.spgemm_work_estimation_or_compute as a bitstring.

spgemm_work_estimation_or_compute(ssa)

gpu.spgemm_work_estimation_or_compute - SpGEMM work estimation operation

spmat_get_size()

Return op name gpu.spmat_get_size as a bitstring.

spmat_get_size(ssa)

gpu.spmat_get_size - SpMat get size operation

spmm()

Return op name gpu.spmm as a bitstring.

spmm(ssa)

gpu.spmm - SpMM operation

spmm_buffer_size()

Return op name gpu.spmm_buffer_size as a bitstring.

spmm_buffer_size(ssa)

gpu.spmm_buffer_size - Precompute buffersize for SpMM operation

spmv()

Return op name gpu.spmv as a bitstring.

spmv(ssa)

gpu.spmv - SpMV operation

spmv_buffer_size()

Return op name gpu.spmv_buffer_size as a bitstring.

spmv_buffer_size(ssa)

gpu.spmv_buffer_size - Precompute buffersize for SpMV operation

subgroup_broadcast()

Return op name gpu.subgroup_broadcast as a bitstring.

subgroup_broadcast(ssa)

gpu.subgroup_broadcast - Broadcasts a value from the specific lane across subgroup

subgroup_id()

Return op name gpu.subgroup_id as a bitstring.

subgroup_id(ssa)

gpu.subgroup_id

subgroup_mma_compute()

Return op name gpu.subgroup_mma_compute as a bitstring.

subgroup_mma_compute(ssa)

gpu.subgroup_mma_compute - GPU warp synchronous matrix multiply accumulate

subgroup_mma_constant_matrix()

Return op name gpu.subgroup_mma_constant_matrix as a bitstring.

subgroup_mma_constant_matrix(ssa)

gpu.subgroup_mma_constant_matrix - GPU warp synchronous constant matrix

subgroup_mma_elementwise()

Return op name gpu.subgroup_mma_elementwise as a bitstring.

subgroup_mma_elementwise(ssa)

gpu.subgroup_mma_elementwise - GPU warp elementwise operation on a matrix

subgroup_mma_extract_thread_local()

Return op name gpu.subgroup_mma_extract_thread_local as a bitstring.

subgroup_mma_extract_thread_local(ssa)

gpu.subgroup_mma_extract_thread_local - Extract a value from GPU warp by invocation and indices

subgroup_mma_insert_thread_local()

Return op name gpu.subgroup_mma_insert_thread_local as a bitstring.

subgroup_mma_insert_thread_local(ssa)

gpu.subgroup_mma_insert_thread_local - Insert a value into GPU warp by invocation and indices

subgroup_mma_load_matrix()

Return op name gpu.subgroup_mma_load_matrix as a bitstring.

subgroup_mma_load_matrix(ssa)

gpu.subgroup_mma_load_matrix - GPU warp synchronous matrix load

subgroup_mma_store_matrix()

Return op name gpu.subgroup_mma_store_matrix as a bitstring.

subgroup_mma_store_matrix(ssa)

gpu.subgroup_mma_store_matrix - GPU warp synchronous matrix store

subgroup_reduce()

Return op name gpu.subgroup_reduce as a bitstring.

subgroup_reduce(ssa)

gpu.subgroup_reduce - Reduce values among subgroup.

subgroup_size()

Return op name gpu.subgroup_size as a bitstring.

subgroup_size(ssa)

gpu.subgroup_size

terminator()

Return op name gpu.terminator as a bitstring.

terminator(ssa)

gpu.terminator - Terminator for GPU launch regions.

thread_id()

Return op name gpu.thread_id as a bitstring.

thread_id(ssa)

gpu.thread_id

wait()

Return op name gpu.wait as a bitstring.

wait(ssa)

gpu.wait - Wait for async gpu ops to complete.

warp_execute_on_lane_0()

Return op name gpu.warp_execute_on_lane_0 as a bitstring.

warp_execute_on_lane_0(ssa)

gpu.warp_execute_on_lane_0 - Executes operations in the associated region on thread #0 of aSPMD program

yield()

Return op name gpu.yield as a bitstring.

yield(ssa)

gpu.yield - GPU yield operation

Functions

all_reduce()

Return op name gpu.all_reduce as a bitstring.

all_reduce(ssa)

gpu.all_reduce - Reduce values among workgroup.

This op has support for result type inference.

Attributes

op - Optional, GPU_AllReduceOperationAttr, built-in reduction operations supported by gpu.allreduce.
uniform - Optional, UnitAttr, unit attribute

Operands

value - Single, AnyIntegerOrFloat, Integer or Float

Results

result - Single, AnyIntegerOrFloat, Integer or Float

Description

The all_reduce op reduces the value of every work item across a local workgroup. The result is equal for all work items of a workgroup.

For example, both

%1 = gpu.all_reduce add %0 {} : (f32) -> (f32)
%2 = gpu.all_reduce %0 {
^bb(%lhs : f32, %rhs : f32):
  %sum = arith.addf %lhs, %rhs : f32
  "gpu.yield"(%sum) : (f32) -> ()
} : (f32) -> (f32)

compute the sum of each work item's %0 value. The first version specifies the accumulation as operation, whereas the second version specifies the accumulation as code region. The reduction operation must be one of:

Integer types: add, mul, minui, minsi, maxui, maxsi, and, or, xor
Floating point types: add, mul, minnumf, maxnumf, minimumf, maximumf

If uniform flag is set either none or all work items of a workgroup need to execute this op in convergence.

alloc()

Return op name gpu.alloc as a bitstring.

alloc(ssa)

gpu.alloc - GPU memory allocation operation.

Attributes

hostShared - Optional, UnitAttr, unit attribute

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
dynamicSizes - Variadic, Index, variadic of index
symbolOperands - Variadic, Index, variadic of index

Results

memref - Single, AnyMemRef, memref of any non-token type values
asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.alloc operation allocates a region of memory on the GPU. It is similar to the memref.alloc op, but supports asynchronous GPU execution.

The op does not execute before all async dependencies have finished executing.

If the async keyword is present, the op is executed asynchronously (i.e. it does not block until the execution has finished on the device). In that case, it also returns a !gpu.async.token.

If the host_shared keyword is present, the memory will be allocated in a memory accessible both on host and on device.

Example:

%memref, %token = gpu.alloc async [%dep] host_shared (%width) : memref<64x?xf32, 1>

ballot()

Return op name gpu.ballot as a bitstring.

ballot(ssa)

gpu.ballot - Collects predicate values from all threads in a subgroup.

Operands

predicate - Single, I1, 1-bit signless integer

Results

result - Single, AnySignlessInteger, signless integer

Description

The ballot op performs a ballot operation across all threads in a subgroup. Each thread contributes its predicate value as a single bit. The result is an integer where the Nth bit is set iff the Nth thread's predicate is true.

The result type must be a signless integer type. The exact supported widths are target-dependent. Common GPU targets support i32 and i64.

Example:

%0 = gpu.ballot %pred : i32
%1 = gpu.ballot %pred : i64

barrier()

Return op name gpu.barrier as a bitstring.

barrier(ssa)

gpu.barrier - Synchronizes work items within an execution scope.

Attributes

address_spaces - Optional, GPU_AddressSpaceAttrArray, GPU Address Space array
scope - Single, GPU_BarrierScopeAttr, barrier execution scope

Operands

named_barrier - Optional, GPU_NamedBarrier, named barrier type

Description

The barrier op synchronizes work items within the specified execution scope. By default, the scope is workgroup, synchronizing all work items in a workgroup.

// Synchronize all work items in the workgroup, making all prior
// memory accesses visible.
gpu.barrier

waits until all work items in the workgroup have reached the operation and all memory accesses made by these work items prior to the op are visible to all work items in the workgroup. Data hazards between work items accessing the same memory can be avoided by synchronizing work items in-between these accesses.

The scope attribute controls the execution scope of the barrier:

// Synchronize within a subgroup (warp/wavefront).
gpu.barrier scope <subgroup>
// Synchronize within a cluster.
gpu.barrier scope <cluster>

A named barrier allows synchronizing a specific subset of subgroups that have been associated with a named barrier handle. Named barriers require workgroup scope.

// Initialize a named barrier for 4 participating members.
%nb = gpu.initialize_named_barrier %c4 : i32 -> !gpu.named_barrier
// Wait on the named barrier.
gpu.barrier named(%nb : !gpu.named_barrier)

If the memfence attribute is specified, the set of memory accesses that must be completed after the barrier resolves is limited to only those accesses that read from or write to the specified address spaces. In particular, specifying memfence [] creates a barrier that is not required to affect the visibility of any memory operations and is purely used for synchronizing work items.

// Only workgroup address space accesses required to be visible.
gpu.barrier memfence [#gpu.address_space<workgroup>]
// No memory accesses required to be visible.
gpu.barrier memfence []
// All memory accesses required to be visible.
gpu.barrier

The three clauses can be combined in any order, but not all combinations may be supported on a given target:

// Named barrier with a workgroup-only memory fence.
gpu.barrier named(%nb : !gpu.named_barrier) memfence [#gpu.address_space<workgroup>]
// Subgroup barrier with a global fence.
gpu.barrier memfence [#gpu.address_space<global>] scope <subgroup>

Once one thread of execution in a given scope (say, thread in a workgroup) has executed a particular dynamic instance of gpu.barrier, all other threads in that scope are required to execute the same dynamic instance of gpu.barrier before any thread executes any other instance of it. That is, you cannot, for example, have the two subgroups of a workgroup arrive at gpu.barrier ops in different branches of an if statement and have this work.

binary()

Return op name gpu.binary as a bitstring.

binary(ssa)

gpu.binary - An Op for storing serialized GPU binary objects.

Attributes

sym_name - Single, SymbolNameAttr, string attribute
offloadingHandler - Optional, OffloadingTranslationAttr, any attribute with the OffloadingTranslationAttrTrait trait.
objects - Single, GPUObjectArrayAttr, an array of GPU object attributes with at least 1 elements

Description

GPU binaries provide a semantic mechanism for storing GPU objects, e.g. the result of compiling a GPU module to an object file.

This operation has 3 arguments:

The name of the binary.
An optional attribute implementing the offloading LLVM translation interface.
An array of GPU object attributes.

During translation, the offloading attribute will be called for translating GPU binary and launch_func operations. The default offloading handler is: #gpu.select_object, this handler selects the first object from the array and embeds it as a string.

Examples:

  // Selects the first object.
  gpu.binary @myobject [#gpu.object<...>, #gpu.object<...>]
  // Uses the `#foo.my_handler` for handling the binary during translation.
  gpu.binary @myobject <#foo.my_handler> [#gpu.object<...>, #gpu.object<...>]
  // Selects the object with the `#rocdl.target` target attribute.
  gpu.binary @myobject <#gpu.select_object<#rocdl.target>> [#gpu.object<...>, #gpu.object<#rocdl.target, ...>]

block_dim()

Return op name gpu.block_dim as a bitstring.

block_dim(ssa)

gpu.block_dim

This op has support for result type inference.

Attributes

dimension - Single, GPU_DimensionAttr, a dimension, either 'x', 'y', or 'z'
upper_bound - Optional, IndexAttr, index attribute

Results

anonymous - Single, Index, index

Description

Returns the number of threads in the thread block (aka the block size) along the x, y, or z dimension.

Example:

%bDimX = gpu.block_dim x

If known_block_size is set on an this operation's enclosing gpu.func, or gpu.known_block_size is set on an enclosing FunctionOpInterface implementor, or if the enclosing gpu.launch specifies a constant size for dimension's blocks, these contextual facts may be used to infer that this operation has a constant value, though such a transformation will not be performed by canonicalization or the default constant folder. Executions which cause that constant-value assumption to be false incur undefined behavior.

If upper_bound is set, executions where the bblock size along dimension exceeds upper_bound cause undefined behavior.

There is an implicit upper bound of kMaxDim (currently uint32_t::max).

block_id()

Return op name gpu.block_id as a bitstring.

block_id(ssa)

gpu.block_id

This op has support for result type inference.

Attributes

dimension - Single, GPU_DimensionAttr, a dimension, either 'x', 'y', or 'z'
upper_bound - Optional, IndexAttr, index attribute

Results

anonymous - Single, Index, index

Description

Returns the block id, i.e. the index of the current block within the grid along the x, y, or z dimension.

Example:

%bIdY = gpu.block_id y

If upper_bound is set, or if one can be inferred from known_grid_size-type annotations in context, executions where the block index in dimension would be greater than or equal to that bound cause undefined behavior. upper_bound takes priority over bounds inferrable from context.

There is an implicit upper bound of kMaxDim (currently uint32_t::max).

cluster_block_id()

Return op name gpu.cluster_block_id as a bitstring.

cluster_block_id(ssa)

gpu.cluster_block_id

This op has support for result type inference.

Attributes

dimension - Single, GPU_DimensionAttr, a dimension, either 'x', 'y', or 'z'
upper_bound - Optional, IndexAttr, index attribute

Results

anonymous - Single, Index, index

Description

Returns the block id within the cluster along the x, y, or z dimension.

Example:

%cBlockIdY = gpu.cluster_block_id y

If upper_bound is set, then executing (a lowering of) this operation in an environment where the number of thread blocks per cluster along dimension is greater than upper_bound causes undefined behavior.

There is an implicit upper bound of kMaxClusterDim (currently 8).

cluster_dim()

Return op name gpu.cluster_dim as a bitstring.

cluster_dim(ssa)

gpu.cluster_dim

This op has support for result type inference.

Attributes

dimension - Single, GPU_DimensionAttr, a dimension, either 'x', 'y', or 'z'
upper_bound - Optional, IndexAttr, index attribute

Results

anonymous - Single, Index, index

Description

Returns the number of cluster identifiers per grid along the x, y, or z dimension.

Example:

%cDimX = gpu.cluster_dim x

If upper_bound is set, then executing (a lowering of) this operation in an environment where the clusters per grid is greater than upper_bound causes undefined behavior.

There is an implicit upper bound of kMaxDim (currently uint32_t::max).

cluster_dim_blocks()

Return op name gpu.cluster_dim_blocks as a bitstring.

cluster_dim_blocks(ssa)

gpu.cluster_dim_blocks

This op has support for result type inference.

Attributes

dimension - Single, GPU_DimensionAttr, a dimension, either 'x', 'y', or 'z'
upper_bound - Optional, IndexAttr, index attribute

Results

anonymous - Single, Index, index

Description

Returns the number of thread blocks in the cluster along the x, y, or z dimension.

Example:

%cDimBlocksX = gpu.cluster_dim_blocks x

If upper_bound is set, then executing (a lowering of) this operation in an environment where the thread blocks per cluster is greater than upper_bound causes undefined behavior.

There is an implicit upper bound of kMaxClusterDim (currently 8).

cluster_id()

Return op name gpu.cluster_id as a bitstring.

cluster_id(ssa)

gpu.cluster_id

This op has support for result type inference.

Attributes

dimension - Single, GPU_DimensionAttr, a dimension, either 'x', 'y', or 'z'
upper_bound - Optional, IndexAttr, index attribute

Results

anonymous - Single, Index, index

Description

Returns the cluster id, i.e. the index of the current cluster within the grid along the x, y, or z dimension.

Example:

%cIdY = gpu.cluster_id y

If upper_bound is set, then executing (a lowering of) this operation in an environment where the number of clusters in the grid along dimension is greater than upper_bound causes undefined behavior.

There is an implicit upper bound of kMaxDim (currently uint32_t::max).

container_module_attribute_name()

Get the name of the attribute used to annotate the modules that contain kernel modules.

Examples

iex> Beaver.MLIR.Dialect.GPU.container_module_attribute_name() :"gpu.container_module"

create_2to4_spmat()

Return op name gpu.create_2to4_spmat as a bitstring.

create_2to4_spmat(ssa)

gpu.create_2to4_spmat - Create sparse matrix with 2:4 sparsity operation

Attributes

pruneFlag - Single, GPU_Prune2To4SpMatFlagAttr, pruning strategy for 2:4 sparse matrix

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
rows - Single, Index, index
cols - Single, Index, index
memref - Single, AnyMemRef, memref of any non-token type values

Results

spMat - Single, GPU_SparseSpMatHandle, sparse matrix handle type
asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.create_2to4_spmat operation initializes a sparse matrix in dense format with 2:4 sparsity. The buffers must already be copied from the host to the device prior to using this operation. The operation returns a handle to the sparse matrix descriptor.

If the async keyword is present, the op is executed asynchronously (i.e. it does not block until the execution has finished on the device). In that case, it returns a !gpu.async.token in addition to the environment.

Example:

%spmat, %token = gpu.create_2to4_spmat async [%dep] {PRUNE_AND_CHECK} %rows, %cols, %mem: memref<?xf64>

create_bsr()

Return op name gpu.create_bsr as a bitstring.

create_bsr(ssa)

gpu.create_bsr - Create sparse matrix in BSR format operation

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
brows - Single, Index, index
bcols - Single, Index, index
bnnz - Single, Index, index
rBlockSize - Single, Index, index
cBlockSize - Single, Index, index
bRowPos - Single, AnyMemRef, memref of any non-token type values
bColIdxs - Single, AnyMemRef, memref of any non-token type values
values - Single, AnyMemRef, memref of any non-token type values

Results

spmat - Single, GPU_SparseSpMatHandle, sparse matrix handle type
asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.create_bsr operation initializes a sparse matrix in BSR format with the given sizes for the matrix and blocks from the given position, index, and values buffers. The buffers must already be copied from the host to the device prior to using this operation. The operation returns a handle to the sparse matrix descriptor.

The BSR format is similar to CSR, where the column indices represent two-dimensional blocks instead of a single matrix entry. Note that this operation (currently) only supports storage with square blocks, i.e., rBlockSize == cBlockSize.

Example:

%spmat, %token = gpu.create_bsr async [%dep]
   %brows, %bcols, %bnnz, %rBlockSize, %cBlockSize,
   %bRowPos, %bColIdxs, %values : memref<?xindex>, memref<?xindex>, memref<?xf64>

create_coo()

Return op name gpu.create_coo as a bitstring.

create_coo(ssa)

gpu.create_coo - Create sparse matrix in COO format operation

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
rows - Single, Index, index
cols - Single, Index, index
nnz - Single, Index, index
rowIdxs - Single, AnyMemRef, memref of any non-token type values
colIdxs - Single, AnyMemRef, memref of any non-token type values
values - Single, AnyMemRef, memref of any non-token type values

Results

spmat - Single, GPU_SparseSpMatHandle, sparse matrix handle type
asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.create_coo operation initializes a sparse matrix in COO format with the given sizes from the given index and values buffers. The buffers must already be copied from the host to the device prior to using this operation. The operation returns a handle to the sparse matrix descriptor. Note that this operation builds the COO in SoA format.

Example:

%spmat, %token = gpu.create_coo async [%dep] %rows, %cols, %nnz, %rowIdx,
    %colIdx, %values : memref<?xindex>, memref<?xindex>, memref<?xf64>

create_coo_aos()

Return op name gpu.create_coo_aos as a bitstring.

create_coo_aos(ssa)

gpu.create_coo_aos - Create sparse matrix in COO format operation (AoS)

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
rows - Single, Index, index
cols - Single, Index, index
nnz - Single, Index, index
idxs - Single, AnyMemRef, memref of any non-token type values
values - Single, AnyMemRef, memref of any non-token type values

Results

spmat - Single, GPU_SparseSpMatHandle, sparse matrix handle type
asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.create_coo_aos operation initializes a sparse matrix in COO format with the given sizes from the given index and values buffers. The buffers must already be copied from the host to the device prior to using this operation. The operation returns a handle to the sparse matrix descriptor. Unlike the default gpu.create_coo operation, this operation builds the COO format from a single index buffer in AoS format (note that this feature has been deprecated in cuSparse 11.2).

Example:

%spmat, %token = gpu.create_coo_aos async [%dep] %rows, %cols, %nnz, %idxs,
    %values : memref<?xindex>, memref<?xf64>

create_csc()

Return op name gpu.create_csc as a bitstring.

create_csc(ssa)

gpu.create_csc - Create sparse matrix in CSC format operation

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
rows - Single, Index, index
cols - Single, Index, index
nnz - Single, Index, index
colPos - Single, AnyMemRef, memref of any non-token type values
rowIdxs - Single, AnyMemRef, memref of any non-token type values
values - Single, AnyMemRef, memref of any non-token type values

Results

spmat - Single, GPU_SparseSpMatHandle, sparse matrix handle type
asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.create_csc operation initializes a sparse matrix in CSC format with the given sizes from the given position, index, and values buffers. The buffers must already be copied from the host to the device prior to using this operation. The operation returns a handle to the sparse matrix descriptor.

The CSC format has exactly the same memory layout as its transpose in CSR format (and vice versa).

Example:

%spmat, %token = gpu.create_csc async [%dep] %rows, %cols, %nnz, %colPos,
    %rowIdx, %values : memref<?xindex>, memref<?xindex>, memref<?xf64>

create_csr()

Return op name gpu.create_csr as a bitstring.

create_csr(ssa)

gpu.create_csr - Create sparse matrix in CSR format operation

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
rows - Single, Index, index
cols - Single, Index, index
nnz - Single, Index, index
rowPos - Single, AnyMemRef, memref of any non-token type values
colIdxs - Single, AnyMemRef, memref of any non-token type values
values - Single, AnyMemRef, memref of any non-token type values

Results

spmat - Single, GPU_SparseSpMatHandle, sparse matrix handle type
asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.create_csr operation initializes a sparse matrix in CSR format with the given sizes from the given position, index, and values buffers. The buffers must already be copied from the host to the device prior to using this operation. The operation returns a handle to the sparse matrix descriptor.

The CSR format has exactly the same memory layout as its transpose in CSC format (and vice versa).

Example:

%spmat, %token = gpu.create_csr async [%dep] %rows, %cols, %nnz, %rowPos,
    %colIdx, %values : memref<?xindex>, memref<?xindex>, memref<?xf64>

create_dn_tensor()

Return op name gpu.create_dn_tensor as a bitstring.

create_dn_tensor(ssa)

gpu.create_dn_tensor - Create dense tensor operation

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
memref - Single, AnyMemRef, memref of any non-token type values
dims - Variadic, Index, variadic of index

Results

dnTensor - Single, GPU_SparseDnTensorHandle, dense tensor handle type
asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.create_dn_tensor operation initializes a dense tensor from the given values buffer and sizes. The buffer must already be copied from the host to the device prior to using this operation. The operation returns a handle to the dense tensor descriptor.

Example:

%dmat, %token = gpu.create_dn_tensor async [%dep] %mem, %dims : index, index into memref<?xf64>

dealloc()

Return op name gpu.dealloc as a bitstring.

dealloc(ssa)

gpu.dealloc - GPU memory deallocation operation

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
memref - Single, AnyMemRef, memref of any non-token type values

Results

asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.dealloc operation frees the region of memory referenced by a memref which was originally created by the gpu.alloc operation. It is similar to the memref.dealloc op, but supports asynchronous GPU execution.

The op does not execute before all async dependencies have finished executing.

If the async keyword is present, the op is executed asynchronously (i.e. it does not block until the execution has finished on the device). In that case, it returns a !gpu.async.token.

Example:

%token = gpu.dealloc async [%dep] %memref : memref<8x64xf32, 1>

destroy_dn_tensor()

Return op name gpu.destroy_dn_tensor as a bitstring.

destroy_dn_tensor(ssa)

gpu.destroy_dn_tensor - Destroy dense tensor operation

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
dnTensor - Single, GPU_SparseDnTensorHandle, dense tensor handle type

Results

asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.destroy_dn_tensor operation releases all resources of a dense tensor represented by a handle that was previously created by a gpu.create_dn_tensor operation.

Example:

%token = gpu.destroy_dn_tensor async [%dep] %dnTensor

destroy_sp_mat()

Return op name gpu.destroy_sp_mat as a bitstring.

destroy_sp_mat(ssa)

gpu.destroy_sp_mat - Destroy sparse matrix operation

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
spmat - Single, GPU_SparseSpMatHandle, sparse matrix handle type

Results

asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.destroy_sp_mat operation releases all resources of a sparse matrix represented by a handle that was previously created by a one of the sparse matrix creation operations.

Example:

%token = gpu.destroy_sp_mat async [%dep] %spmat

dynamic_shared_memory()

Return op name gpu.dynamic_shared_memory as a bitstring.

dynamic_shared_memory(ssa)

gpu.dynamic_shared_memory - Get the memref for dynamic shared memory

Results

resultMemref - Single, anonymous/composite constraint, 1D memref of 8-bit signless integer values

Description

This operation provides a memref pointer to the start of dynamic shared memory, often referred to as workgroup memory. It's important to note that this dynamic shared memory needs to be allocated at kernel launch. One can conveniently utilize the dynamic_shared_memory_size parameter of gpu.launch for this purpose.

Examples:

%0 = gpu.dynamic.shared.memory : memref<?xi8, #gpu.address_space<workgroup>>
%1 = memref.view %0[%c8192][] : memref<?xi8, #gpu.address_space<workgroup>>
                        to memref<32x64xf32, #gpu.address_space<workgroup>>
%2 = memref.view %0[%c16384][] : memref<?xi8, #gpu.address_space<workgroup>>
                        to memref<32x64xf32, #gpu.address_space<workgroup>>

func()

Return op name gpu.func as a bitstring.

func(ssa)

gpu.func - Function executable on a GPU

Attributes

function_type - Single, anonymous/composite constraint, type attribute of function type
arg_attrs - Optional, DictArrayAttr, Array of dictionary attributes
res_attrs - Optional, DictArrayAttr, Array of dictionary attributes
workgroup_attrib_attrs - Optional, DictArrayAttr, Array of dictionary attributes
private_attrib_attrs - Optional, DictArrayAttr, Array of dictionary attributes
known_block_size - Optional, GPU_OptionalDimSizeHintAttr, i32 dense array attribute with 3 elements (if present) and all elements >= 1
known_grid_size - Optional, GPU_OptionalDimSizeHintAttr, i32 dense array attribute with 3 elements (if present) and all elements >= 1
known_cluster_size - Optional, GPU_OptionalDimSizeHintAttr, i32 dense array attribute with 3 elements (if present) and all elements >= 1
workgroup_attributions - Optional, I64Attr, 64-bit signless integer attribute whose value is non-negative
kernel - Optional, UnitAttr, unit attribute

Description

Defines a function that can be executed on a GPU. This supports memory attribution and its body has a particular execution model.

GPU functions are either kernels (as indicated by the kernel attribute) or regular functions. The former can be launched from the host side, while the latter are device side only.

The memory attribution defines SSA values that correspond to memory buffers allocated in the memory hierarchy of the GPU (see below).

The operation has one attached region that corresponds to the body of the function. The region arguments consist of the function arguments without modification, followed by buffers defined in memory annotations. The body of a GPU function, when launched, is executed by multiple work items. There are no guarantees on the order in which work items execute, or on the connection between them. In particular, work items are not necessarily executed in lock-step. Synchronization ops such as "gpu.barrier" should be used to coordinate work items. Declarations of GPU functions, i.e. not having the body region, are not supported.

A function may optionally be annotated with the block and/or grid sizes that will be used when it is launched using the known_block_size and known_grid_size attributes, respectively. If set, these attributes must be arrays of three 32-bit integers giving the x, y, and z launch dimensions. Launching a kernel that has these annotations, or that calls a function with these annotations, using a block size or grid size other than what is specified is undefined behavior. These attributes may be set on non-gpu.func functions by using gpu.known_block_size or gpu.known_grid_size, but this carries the risk that they will de discarded.

Syntax:

op ::= `gpu.func` symbol-ref-id `(` argument-list `)` (`->`
function-result-list)?
       memory-attribution `kernel`? function-attributes? region

memory-attribution ::= (`workgroup` `(` ssa-id-and-type-list `)`)?
                       (`private` `(` ssa-id-and-type-list `)`)?

Example:

gpu.func @foo(%arg0: index)
    workgroup(%workgroup: memref<32xf32, 3>)
    private(%private: memref<1xf32, 5>)
    kernel
    attributes {qux: "quux"} {
  gpu.return
}

The generic form illustrates the concept

"gpu.func"(%arg: index) {sym_name: "foo", kernel, qux: "quux"} ({
^bb0(%arg0: index, %workgroup: memref<32xf32, 3>,
     %private: memref<1xf32, 5>):
  "gpu.return"() : () -> ()
}) : (index) -> ()

Note the non-default memory spaces used in memref types in memory attribution.

global_id()

Return op name gpu.global_id as a bitstring.

global_id(ssa)

gpu.global_id

This op has support for result type inference.

Attributes

dimension - Single, GPU_DimensionAttr, a dimension, either 'x', 'y', or 'z'
upper_bound - Optional, IndexAttr, index attribute

Results

anonymous - Single, Index, index

Description

Returns the unique global workitem/thread id, i.e., the unique index of the current workitem/thread within all workgroups / grid along the x, y, or z dimension.

Example:

%gidX = gpu.global_id x
%gidX = gpu.global_id x upper_bound 65536

The upper_bound attribute defines an upper bound analogously to the ones on thread_id and block_id. If one is not set, the bound may be inferred from a combination of known_block_size and known_grid_size-type annotations.

grid_dim()

Return op name gpu.grid_dim as a bitstring.

grid_dim(ssa)

gpu.grid_dim

This op has support for result type inference.

Attributes

dimension - Single, GPU_DimensionAttr, a dimension, either 'x', 'y', or 'z'
upper_bound - Optional, IndexAttr, index attribute

Results

anonymous - Single, Index, index

Description

Returns the number of thread blocks in the grid along the x, y, or z dimension.

Example:

%gDimZ = gpu.grid_dim z

If known_grid_size is set on an this operation's enclosing gpu.func, or gpu.known_grid_size is set on an enclosing FunctionOpInterface implementor, or if the enclosing gpu.launch specifies a constant size for dimension's grid length, these contextual facts may be used to infer that this operation has a constant value, though such a transformation will not be performed by canonicalization or the default constant folder. Executions which cause that constant-value assumption to be false incur undefined behavior.

If upper_bound is set, executions where the grid size in dimension would exceed upper_bound cause undefined behavior.

There is an implicit upper bound of kMaxDim (currently uint32_t::max).

host_register()

Return op name gpu.host_register as a bitstring.

host_register(ssa)

gpu.host_register - Registers a memref for access from device.

Operands

value - Single, AnyUnrankedMemRef, unranked.memref of any non-token type values

Description

This op maps the provided host buffer into the device address space.

This operation may not be supported in every environment, there is not yet a way to check at runtime whether this feature is supported.

Writes from the host are guaranteed to be visible to device kernels that are launched afterwards. Writes from the device are guaranteed to be visible on the host after synchronizing with the device kernel completion.

host_unregister()

Return op name gpu.host_unregister as a bitstring.

host_unregister(ssa)

gpu.host_unregister - Unregisters a memref for access from device.

Operands

value - Single, AnyUnrankedMemRef, unranked.memref of any non-token type values

Description

This op unmaps the provided host buffer from the device address space.

This operation may not be supported in every environment, there is not yet a

way to check at runtime whether this feature is supported.

initialize_named_barrier()

Return op name gpu.initialize_named_barrier as a bitstring.

initialize_named_barrier(ssa)

gpu.initialize_named_barrier - Initialize a named barrier with a member count.

This op has support for result type inference.

Operands

member_count - Single, I32, 32-bit signless integer

Results

result - Single, GPU_NamedBarrier, named barrier type

Description

Initializes a named barrier object with the given number of participating members (subgroups) and returns a handle to it. All members that will synchronize on this barrier must be accounted for in the count.

%nb = gpu.initialize_named_barrier %num_members : i32 -> !gpu.named_barrier

lane_id()

Return op name gpu.lane_id as a bitstring.

lane_id(ssa)

gpu.lane_id

This op has support for result type inference.

Attributes

upper_bound - Optional, IndexAttr, index attribute

Results

result - Single, Index, index

Description

Returns the lane id within the subgroup (warp/wave).

Example:

%laneId = gpu.lane_id

If upper_bound is set, executions with more than upper_bound lanes per subgroup cause undefined behavior. In the abscence of upper_bound, the lane id is still assumed to be non-negative and less than the target-independent kMaxSubgroupSize (currently 128).

launch()

Return op name gpu.launch as a bitstring.

launch(ssa)

gpu.launch - GPU kernel launch operation

Attributes

cooperative - Optional, UnitAttr, unit attribute
module - Optional, FlatSymbolRefAttr, flat symbol reference attribute
function - Optional, FlatSymbolRefAttr, flat symbol reference attribute
workgroup_attributions - Optional, I64Attr, 64-bit signless integer attribute whose value is non-negative

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
gridSizeX - Single, Index, index
gridSizeY - Single, Index, index
gridSizeZ - Single, Index, index
blockSizeX - Single, Index, index
blockSizeY - Single, Index, index
blockSizeZ - Single, Index, index
clusterSizeX - Optional, Index, index
clusterSizeY - Optional, Index, index
clusterSizeZ - Optional, Index, index
dynamicSharedMemorySize - Optional, I32, 32-bit signless integer
asyncObject - Optional, AnyType, any non-token type

Results

asyncToken - Optional, GPU_AsyncToken, async token type

Description

Launch a kernel on the specified grid of thread blocks. The body of the kernel is defined by the single region that this operation contains.

The async execution model is equivalent to the gpu.launch_func op, refer to its description.

The first three operands (following any async dependencies) are grid sizes along the x,y,z dimensions and the following three are block sizes along the x,y,z dimensions. When a lower-dimensional kernel is required, unused sizes must be explicitly set to 1. The last operand is optional and corresponds to the amount of dynamic shared memory a kernel's workgroup should be allocated; when this operand is not present, a zero size is assumed.

The body region has at least twelve arguments, or eighteen if cluster dimensions are present, grouped as follows:

three optional arguments that contain cluster identifiers along x,y,z dimensions;
three arguments that contain block identifiers along x,y,z dimensions;
three arguments that contain thread identifiers along x,y,z dimensions;
operands of the gpu.launch operation as is (i.e. the operands for grid and block sizes).
a variadic number of Workgroup memory attributions.
a variadic number of Private memory attributions.

The function and module attributes are optional and specifies the kernel name and a module in which the kernel should be outlined.

Syntax:

operation ::= `gpu.launch` (`async` (`[` ssa-id-list `]`)? )?
                         ( `clusters` `(` ssa-id-list `)` `in` ssa-reassignment )?
                         `blocks` `(` ssa-id-list `)` `in` ssa-reassignment
                         `threads` `(` ssa-id-list `)` `in` ssa-reassignment
                         (dynamic_shared_memory_size ssa-use)?
                         (`module(` symbol-ref-id `)`)?
                         (`function(` symbol-ref-id `)`)?
                         memory-attribution
                         region attr-dict?
ssa-reassignment ::= `(` ssa-id `=` ssa-use (`,` ssa-id `=` ssa-use)* `)`
memory-attribution ::= (`workgroup` `(` ssa-id-and-type-list `)`)?
                       (`private` `(` ssa-id-and-type-list `)`)?

Example:

gpu.launch blocks(%bx, %by, %bz) in (%sz_bx = %0, %sz_by = %1, %sz_bz = %2)
           threads(%tx, %ty, %tz) in (%sz_tx = %3, %sz_ty = %4, %sz_tz = %5) {
  // Block and thread identifiers, as well as block/grid sizes are
  // immediately usable inside body region.
  "some_op"(%bx, %tx) : (index, index) -> ()
  // Assuming %val1 is defined outside the gpu.launch region.
  %42 = load %val1[%bx] : memref<?xf32, 1>
}

// Generic syntax explains how the pretty syntax maps to the IR structure.
"gpu.launch"(%cst, %cst, %c1,  // Grid sizes.
             %cst, %c1, %c1)   // Block sizes.

    {/*attributes*/}
    // All sizes and identifiers have "index" size.
    : (index, index, index, index, index, index) -> () {
// The operation passes block and thread identifiers, followed by grid and
// block sizes.
^bb0(%bx : index, %by : index, %bz : index,
     %tx : index, %ty : index, %tz : index,
     %num_bx : index, %num_by : index, %num_bz : index,
     %num_tx : index, %num_ty : index, %num_tz : index)
  "some_op"(%bx, %tx) : (index, index) -> ()
  %3 = "memref.load"(%val1, %bx) : (memref<?xf32, 1>, index) -> f32
}

// Launch with memory attributions.
gpu.launch blocks(%bx, %by, %bz) in (%sz_bx = %0, %sz_by = %1, %sz_bz = %2)
           threads(%tx, %ty, %tz) in (%sz_tx = %3, %sz_ty = %4, %sz_tz = %5)
           workgroup(%workgroup: memref<32xf32, 3>)
           private(%private: memref<1xf32, 5>) {
  // Block and thread identifiers, as well as block/grid sizes are
  // immediately usable inside body region.
  "some_op"(%bx, %tx) : (index, index) -> ()
  // Assuming %val1 is defined outside the gpu.launch region.
  %42 = load %workgroup[%bx] : memref<32xf32, 3>
}

// Launch with clusters.
gpu.launch clusters(%cx, %cy, %cz) in (%sz_cx = %0, %sz_cy = %1, %sz_cz = %2)
           blocks(%bx, %by, %bz) in (%sz_bx = %3, %sz_by = %4, %sz_bz = %5)
           threads(%tx, %ty, %tz) in (%sz_tx = %6, %sz_ty = %7, %sz_tz = %8)
{
  // Cluster, block and thread identifiers, as well as cluster/block/grid
  // sizes are immediately usable inside body region.
  "some_op"(%cx, %bx, %tx) : (index, index, index) -> ()
}

// Launch with module and function attributes.
gpu.launch blocks(%bx, %by, %bz) in (%sz_bx = %0, %sz_by = %1, %sz_bz = %2)
           threads(%tx, %ty, %tz) in (%sz_tx = %3, %sz_ty = %4, %sz_tz = %5)
           module(@kernel_module) function(@kernel_func) {
  "some_op"(%bx, %tx) : (index, index) -> ()
  %42 = load %val1[%bx] : memref<?xf32, 1>
}

Rationale: using operation/block arguments gives analyses a clear way of understanding that a value has additional semantics (e.g., we will need to know what value corresponds to threadIdx.x for coalescing). We can recover these properties by analyzing the operations producing values, but it is easier just to have that information by construction.

launch_func()

Return op name gpu.launch_func as a bitstring.

launch_func(ssa)

gpu.launch_func - Launches a function as a GPU kernel

Attributes

kernel - Single, SymbolRefAttr, symbol reference attribute
cooperative - Optional, UnitAttr, unit attribute

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
gridSizeX - Single, LaunchIndx, index or 32-bit signless integer or 64-bit signless integer
gridSizeY - Single, LaunchIndx, index or 32-bit signless integer or 64-bit signless integer
gridSizeZ - Single, LaunchIndx, index or 32-bit signless integer or 64-bit signless integer
blockSizeX - Single, LaunchIndx, index or 32-bit signless integer or 64-bit signless integer
blockSizeY - Single, LaunchIndx, index or 32-bit signless integer or 64-bit signless integer
blockSizeZ - Single, LaunchIndx, index or 32-bit signless integer or 64-bit signless integer
clusterSizeX - Optional, LaunchIndx, index or 32-bit signless integer or 64-bit signless integer
clusterSizeY - Optional, LaunchIndx, index or 32-bit signless integer or 64-bit signless integer
clusterSizeZ - Optional, LaunchIndx, index or 32-bit signless integer or 64-bit signless integer
dynamicSharedMemorySize - Optional, I32, 32-bit signless integer
kernelOperands - Variadic, AnyType, variadic of any non-token type
asyncObject - Optional, AnyType, any non-token type

Results

asyncToken - Optional, GPU_AsyncToken, async token type

Description

Launch a kernel function on the specified grid of thread blocks. gpu.launch operations are lowered to gpu.launch_func operations by outlining the kernel body into a function in a dedicated module, which reflects the separate compilation process. The kernel function is required to have the gpu.kernel attribute. The module containing the kernel function is required to be a gpu.module. And finally, the module containing the kernel module (which thus cannot be the top-level module) is required to have the gpu.container_module attribute. The gpu.launch_func operation has a symbol attribute named kernel to identify the fully specified kernel function to launch (both the gpu.module and func).

By default, the host implicitly blocks until kernel execution has completed.

Otherwise, the operation supports two async models.

The first one is dependency-based and is enabled when the async keyword is present in text form, and corresponds to when the operation produces the optional token result of type !gpu.async.token. Other async GPU ops can take this token as dependency. In this case, the gpu.launch_func does not block, and supports specifying async dependencies: the kernel does not start executing until the ops producing those async dependencies have completed.

The second async model is stream-based. When asyncObject is present, the launch operation is queued to execute on the queue represented by it.

The two async models are mutually exclusive.

The operation requires at least the grid and block sizes along the x,y,z dimensions as arguments. When a lower-dimensional kernel is required, unused sizes must be explicitly set to 1.

The remaining operands are optional. The first optional operand corresponds to the amount of dynamic shared memory a kernel's workgroup should be allocated; when this operand is not present, a zero size is assumed.

The remaining operands if present are passed as arguments to the kernel function.

The gpu.launch_func also supports kernel launching with clusters if supported by the target architecture. The cluster size can be set by clusterSizeX, clusterSizeY, and clusterSizeZ arguments. When these arguments are present, the Op launches a kernel that clusters the given thread blocks. This feature is exclusive to certain architectures.

The cooperative attribute indicates that the kernel should be launched cooperatively, guaranteeing that all thread blocks in the grid are co-resident on the GPU simultaneously. This enables grid-wide synchronization patterns.

Example:

module attributes {gpu.container_module} {

  // This module creates a separate compilation unit for the GPU compiler.
  gpu.module @kernels {
    func.func @kernel_1(%arg0 : f32, %arg1 : memref<?xf32, 1>)
        attributes { nvvm.kernel = true } {

      // Operations that produce block/thread IDs and dimensions are
      // injected when outlining the `gpu.launch` body to a function called
      // by `gpu.launch_func`.
      %tIdX = gpu.thread_id x
      %tIdY = gpu.thread_id y
      %tIdZ = gpu.thread_id z

      %bDimX = gpu.block_dim x
      %bDimY = gpu.block_dim y
      %bDimZ = gpu.block_dim z

      %bIdX = gpu.block_id x
      %bIdY = gpu.block_id y
      %bIdZ = gpu.block_id z

      %gDimX = gpu.grid_dim x
      %gDimY = gpu.grid_dim y
      %gDimZ = gpu.grid_dim z

      // (Optional)  Cluster size only for support architectures
      %cIdX = gpu.cluster_id x
      %cIdY = gpu.cluster_id y
      %cIdZ = gpu.cluster_id z

      %cDimX = gpu.cluster_dim x
      %cDimY = gpu.cluster_dim y
      %cDimZ = gpu.cluster_dim z

      "some_op"(%bx, %tx) : (index, index) -> ()
      %42 = load %arg1[%bx] : memref<?xf32, 1>
    }
  }

  %t0 = gpu.wait async
  gpu.launch_func
      async                           // (Optional) Don't block host, return token.
      [%t0]                           // (Optional) Execute only after %t0 has completed.
      @kernels::@kernel_1             // Kernel function.
      clusters in (%cst, %cst, %cst)  // (Optional) Cluster size only for support architectures.
      blocks in (%cst, %cst, %cst)    // Grid size.
      threads in (%cst, %cst, %cst)   // Block size.
      dynamic_shared_memory_size %s   // (Optional) Amount of dynamic shared
                                      // memory to allocate for a workgroup.
      args(%arg0 : f32,               // (Optional) Kernel arguments.
           %arg1 : memref<?xf32, 1>)
}

memcpy()

Return op name gpu.memcpy as a bitstring.

memcpy(ssa)

gpu.memcpy - GPU memcpy operation

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
dst - Single, AnyMemRef, memref of any non-token type values
src - Single, AnyMemRef, memref of any non-token type values

Results

asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.memcpy operation copies the content of one memref to another.

The op does not execute before all async dependencies have finished executing.

If the async keyword is present, the op is executed asynchronously (i.e. it does not block until the execution has finished on the device). In that case, it returns a !gpu.async.token.

Example:

%token = gpu.memcpy async [%dep] %dst, %src : memref<?xf32, 1>, memref<?xf32>

memset()

Return op name gpu.memset as a bitstring.

memset(ssa)

gpu.memset - GPU memset operation

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
dst - Single, AnyMemRef, memref of any non-token type values
value - Single, AnyType, any non-token type

Results

asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.memset operation sets the content of memref to a scalar value.

The op does not execute before all async dependencies have finished executing.

If the async keyword is present, the op is executed asynchronously (i.e. it does not block until the execution has finished on the device). In that case, it returns a !gpu.async.token.

Example:

%token = gpu.memset async [%dep] %dst, %value : memref<?xf32, 1>, f32

module()

Return op name gpu.module as a bitstring.

module(ssa)

gpu.module - A top level compilation unit containing code to be run on a GPU.

Attributes

sym_name - Single, SymbolNameAttr, string attribute
targets - Optional, GPUNonEmptyTargetArrayAttr, array of GPU target attributes with at least 1 elements
offloadingHandler - Optional, OffloadingTranslationAttr, any attribute with the OffloadingTranslationAttrTrait trait.

Description

GPU module contains code that is intended to be run on a GPU. A host device can launch this code through a gpu.launc_func that creates a fully qualified symbol through the gpu.module's symbol and a gpu.func symbol contained in the gpu.module.

The module's top-level scope is modeled by a single region with a single block. GPU modules are required to have a name that is used for symbol resolution by the gpu.launch_func operation.

Using an op with a region to define a GPU module enables "embedding" GPU modules with SIMT execution models in other dialects in a clean manner and allows filtering of code regions to execute passes on only code intended to or not intended to be run on the separate device.

Modules can contain zero or more target attributes. These attributes encode how to transform modules into binary strings and are used by the gpu-module-to-binary pass to transform modules into GPU binaries.

Modules can contain an optional OffloadingTranslationAttr attribute. This attribute will be used during the gpu-module-to-binary pass to specify the OffloadingTranslationAttr used when creating the gpu.binary operation.

gpu.module @symbol_name {
  gpu.func {}
    ...
}
// Module with offloading handler and target attributes.
gpu.module @symbol_name2 <#gpu.select_object<1>> [
    #nvvm.target,
    #rocdl.target<chip = "gfx90a">] {
  gpu.func {}
    ...
}

num_subgroups()

Return op name gpu.num_subgroups as a bitstring.

num_subgroups(ssa)

gpu.num_subgroups

This op has support for result type inference.

Attributes

upper_bound - Optional, IndexAttr, index attribute

Results

result - Single, Index, index

Description

Returns the number of subgroups within a workgroup.

Example:

%numSg = gpu.num_subgroups : index

If upper_bound is set, executions with more than upper_bound subgroups per workgroup cause undefined behavior. There is a default upper bound of kMaxDim (currently uint32_t::max).

number_of_buffers_in_workgroup_attributions_attribute_name()

Returns the name of the attribute containing the number of buffers located in the workgroup memory.

Examples

iex> Beaver.MLIR.Dialect.GPU.number_of_buffers_in_workgroup_attributions_attribute_name() :workgroup_attributions

printf()

Return op name gpu.printf as a bitstring.

printf(ssa)

gpu.printf - Device-side printf, as in CUDA or OpenCL, for debugging

Attributes

format - Single, StrAttr, string attribute

Operands

args - Variadic, anonymous/composite constraint, variadic of integer or index or floating-point

Description

gpu.printf takes a literal format string format and an arbitrary number of scalar arguments that should be printed.

The format string is a C-style printf string, subject to any restrictions imposed by one's target platform.

return()

Return op name gpu.return as a bitstring.

return(ssa)

gpu.return - Terminator for GPU functions.

Operands

operands - Variadic, AnyType, variadic of any non-token type

Description

A terminator operation for regions that appear in the body of gpu.func functions. The operands to the gpu.return are the result values returned by an invocation of the gpu.func.

rotate()

Return op name gpu.rotate as a bitstring.

rotate(ssa)

gpu.rotate - Rotate values within a subgroup.

This op has support for result type inference.

Attributes

offset - Single, I32Attr, 32-bit signless integer attribute whose minimum value is 0
width - Single, I32Attr, 32-bit signless integer attribute whose value is a power of two > 0

Operands

value - Single, AnyIntegerOrFloatOr1DVector, Integer or Float or fixed-length vector of Integer or Float values of ranks 1

Results

rotateResult - Single, AnyIntegerOrFloatOr1DVector, Integer or Float or fixed-length vector of Integer or Float values of ranks 1
valid - Single, I1, 1-bit signless integer

Description

The "rotate" op moves values across lanes in a subgroup (a.k.a., local invocations) within the same subgroup. The width attribute specifies the number of lanes that participate in the rotation, and must be uniform across all participating lanes. Further, the first width lanes of the subgroup must be active.

width must be a power of two, and offset must be in the range [0, width).

Return the rotateResult of the invocation whose id within the group is calculated as follows:

Invocation ID = ((LaneId + offset) & (width - 1)) + (LaneId & ~(width - 1))

Returns the rotateResult and true if the current lane id is smaller than width, and poison value and false otherwise.

example:

%1, %2 = gpu.rotate %0, 1, 16 : f32

For lane k, returns the value from lane (k + cst1) % width.

sddmm()

Return op name gpu.sddmm as a bitstring.

sddmm(ssa)

gpu.sddmm - SDDMM operation

Attributes

modeA - Single, GPU_TransposeModeAttr, transpose mode of sparse matrix supported by sparse tensor ops
modeB - Single, GPU_TransposeModeAttr, transpose mode of sparse matrix supported by sparse tensor ops
computeType - Single, TypeAttr, any type attribute

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
dnmatA - Single, GPU_SparseDnTensorHandle, dense tensor handle type
dnmatB - Single, GPU_SparseDnTensorHandle, dense tensor handle type
spmatC - Single, GPU_SparseSpMatHandle, sparse matrix handle type
buffer - Single, AnyMemRef, memref of any non-token type values

Results

asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.sddmm operation performs the SDDMM operation on the given sparse and dense matrices, and buffer. The operation expects handles returned by previous sparse operations to construct an environment and the operands for SDDMM. The buffer must have been allocated on the device.

Example:

%token = gpu.sddmm async [%dep] %dnmatA{TRANSPOSE}, %dnmatB{TRANSPOSE}, %spmatC, %buffer into f32

The matrix arguments can also be associated with one of the following operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value is NON_TRANSPOSE.

sddmm_buffer_size()

Return op name gpu.sddmm_buffer_size as a bitstring.

sddmm_buffer_size(ssa)

gpu.sddmm_buffer_size - Precompute buffersize for SDDMM operation

Attributes

modeA - Single, GPU_TransposeModeAttr, transpose mode of sparse matrix supported by sparse tensor ops
modeB - Single, GPU_TransposeModeAttr, transpose mode of sparse matrix supported by sparse tensor ops
computeType - Single, TypeAttr, any type attribute

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
dnmatA - Single, GPU_SparseDnTensorHandle, dense tensor handle type
dnmatB - Single, GPU_SparseDnTensorHandle, dense tensor handle type
spmatC - Single, GPU_SparseSpMatHandle, sparse matrix handle type

Results

bufferSz - Single, Index, index
asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.sddmm_buffer_size operation returns the buffer size required to perform the SDDMM operation on the given sparse and dense matrices. The operation expects handles returned by previous sparse operations to construct an environment and the operands for SDDMM.

Example:

%buffersz, %token = gpu.sddmm_buffer_size async [%dep] %dnmatA{TRANSPOSE}, %dnmatB{TRANSPOSE}, %spmatC into f32

The matrix arguments can also be associated with one of the following operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value is NON_TRANSPOSE.

set_csr_pointers()

Return op name gpu.set_csr_pointers as a bitstring.

set_csr_pointers(ssa)

gpu.set_csr_pointers - SpGEMM get size operation

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
spmat - Single, GPU_SparseSpMatHandle, sparse matrix handle type
positions - Single, AnyMemRef, memref of any non-token type values
coordinates - Single, AnyMemRef, memref of any non-token type values
values - Single, AnyMemRef, memref of any non-token type values

Results

asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.set_csr_pointers assigns the given positions, coordinates, and values buffer that reside on the device directly to the given sparse matrix descriptor in csr format.

If the async keyword is present, the op is executed asynchronously (i.e. it does not block until the execution has finished on the device). In that case, it returns a !gpu.async.token in addition to the environment.

Example:

%token = gpu.set_csr_pointers async [%dep] %positions, %coordinates, %values
      : memref<?xf32>, memref<?xindex>, memref<?xindex>

set_default_device()

Return op name gpu.set_default_device as a bitstring.

set_default_device(ssa)

gpu.set_default_device - Set default GPU for operations after this by index

Operands

devIndex - Single, I32, 32-bit signless integer

Description

Operation that sets the current default GPU, using a zero-based index into the set of GPUs on the system. The default GPU setting may be thread-local.

shuffle()

Return op name gpu.shuffle as a bitstring.

shuffle(ssa)

gpu.shuffle - Shuffles values within a subgroup.

This op has support for result type inference.

Attributes

mode - Single, GPU_ShuffleModeAttr, Indexing modes supported by gpu.shuffle.

Operands

value - Single, AnyIntegerOrFloatOr1DVector, Integer or Float or fixed-length vector of Integer or Float values of ranks 1
offset - Single, I32, 32-bit signless integer
width - Single, I32, 32-bit signless integer

Results

shuffleResult - Single, AnyIntegerOrFloatOr1DVector, Integer or Float or fixed-length vector of Integer or Float values of ranks 1
valid - Single, I1, 1-bit signless integer

Description

The "shuffle" op moves values across lanes in a subgroup (a.k.a., local invocation) within the same subgroup. The width argument specifies the number of lanes that participate in the shuffle, and must be uniform across all lanes. Further, the first width lanes of the subgroup must be active.

The intepretation of the offset arguments depends on the selected mode.

Returns the shuffleResult and true if the current lane id is smaller than width, and an unspecified value and false otherwise.

xor example:

%1, %2 = gpu.shuffle xor %0, %offset, %width : f32

For lane k, returns the value %0 from lane k ^ offset. Every lane trades value with exactly one other lane.

down example:

%cst1 = arith.constant 1 : i32
%3, %4 = gpu.shuffle down %0, %cst1, %width : f32

For lane k, returns the value from lane (k + cst1). If (k + cst1) is bigger than or equal to width, the value is poison and valid is false.

up example:

%cst1 = arith.constant 1 : i32
%5, %6 = gpu.shuffle up %0, %cst1, %width : f32

For lane k, returns the value from lane (k - cst1). If (k - cst1) is smaller than 0, the value is poison and valid is false.

idx example:

%cst0 = arith.constant 0 : i32
%7, %8 = gpu.shuffle idx %0, %cst0, %width : f32

Broadcasts the value from lane 0 to all lanes.

spgemm_copy()

Return op name gpu.spgemm_copy as a bitstring.

spgemm_copy(ssa)

gpu.spgemm_copy - SpGEMM copy operation

Attributes

modeA - Single, GPU_TransposeModeAttr, transpose mode of sparse matrix supported by sparse tensor ops
modeB - Single, GPU_TransposeModeAttr, transpose mode of sparse matrix supported by sparse tensor ops
computeType - Single, TypeAttr, any type attribute

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
desc - Single, GPU_SparseSpGEMMOpHandle, SpGEMM operation handle type
spmatA - Single, GPU_SparseSpMatHandle, sparse matrix handle type
spmatB - Single, GPU_SparseSpMatHandle, sparse matrix handle type
spmatC - Single, GPU_SparseSpMatHandle, sparse matrix handle type

Results

asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.spgemm_copy operation copies the sparse matrix result of a SpGEMM computation.

Example:

gpu.spgemm_copy %spmatA, %spmatB, %spmatC, %spgemmDesc: f32

The matrix arguments can also be associated with one of the following operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value is NON_TRANSPOSE.

spgemm_create_descr()

Return op name gpu.spgemm_create_descr as a bitstring.

spgemm_create_descr(ssa)

gpu.spgemm_create_descr - SpGEMM Create Descr operation

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type

Results

desc - Single, GPU_SparseSpGEMMOpHandle, SpGEMM operation handle type
asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.spgemm_create_descr creates a descriptor for the SpGEMM operation. The descriptor describes the SpGEMM operation and stores the internal data throughout the computation. It needs to be passed as an argument to spgemm_* operations.

Example:

%desc, %token = gpu.spgemm_create_descr async [%dep]

spgemm_destroy_descr()

Return op name gpu.spgemm_destroy_descr as a bitstring.

spgemm_destroy_descr(ssa)

gpu.spgemm_destroy_descr - SpGEMM Destroy Descr operation

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
desc - Single, GPU_SparseSpGEMMOpHandle, SpGEMM operation handle type

Results

asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.spgemm_destroy_descr destroys the SpGEMM operation descriptor.

Example:

%token = gpu.spgemm_destroy_descr async [%dep] %desc

spgemm_work_estimation_or_compute()

Return op name gpu.spgemm_work_estimation_or_compute as a bitstring.

spgemm_work_estimation_or_compute(ssa)

gpu.spgemm_work_estimation_or_compute - SpGEMM work estimation operation

Attributes

modeA - Single, GPU_TransposeModeAttr, transpose mode of sparse matrix supported by sparse tensor ops
modeB - Single, GPU_TransposeModeAttr, transpose mode of sparse matrix supported by sparse tensor ops
computeType - Single, TypeAttr, any type attribute
kind - Single, GPU_SpGEMMWorkEstimationOrComputeKindAttr, choose whether spgemm_work_estimation_or_compute does work estimation or compute

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
desc - Single, GPU_SparseSpGEMMOpHandle, SpGEMM operation handle type
spmatA - Single, GPU_SparseSpMatHandle, sparse matrix handle type
spmatB - Single, GPU_SparseSpMatHandle, sparse matrix handle type
spmatC - Single, GPU_SparseSpMatHandle, sparse matrix handle type
bufferSz - Single, Index, index
buffer - Single, AnyMemRef, memref of any non-token type values

Results

bufferSzNew - Single, Index, index
asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.spgemm_work_estimation_or_compute is used to call cusparseSpGEMM_workEstimation or cusparseSpGEMM_compute. Both of them are for both determining the buffer size and performing the actual computation. The operation expects handles returned by previous sparse operations to construct an environment and the operands for SpGEMM. The buffer must have been allocated on the device.

C' = alpha op(A) op(B) + beta * C

Example:

%bufferSz, %token = gpu.spgemm_work_estimation_or_compute async [%dep] {COMPUTE}
                      %desc, %spmatA{NON_TRANSPOSE}, %spmatB{NON_TRANSPOSE},
                      %spmatC, %spgemmDesc, %c0, %alloc: f32 into
                      memref<0xi8>

The matrix arguments can also be associated with one of the following operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value is NON_TRANSPOSE.

spmat_get_size()

Return op name gpu.spmat_get_size as a bitstring.

spmat_get_size(ssa)

gpu.spmat_get_size - SpMat get size operation

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
spmat - Single, GPU_SparseSpMatHandle, sparse matrix handle type

Results

rows - Single, Index, index
cols - Single, Index, index
nnz - Single, Index, index
asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.spmat_get_size operation retrieves the number of rows, number of columns, and number of non-zero elements of a sparse matrix.

Example:

%rows, %cols, %nnz, %token = gpu.spmat_get_size async [%dep] %spmatC

spmm()

Return op name gpu.spmm as a bitstring.

spmm(ssa)

gpu.spmm - SpMM operation

Attributes

modeA - Single, GPU_TransposeModeAttr, transpose mode of sparse matrix supported by sparse tensor ops
modeB - Single, GPU_TransposeModeAttr, transpose mode of sparse matrix supported by sparse tensor ops
computeType - Single, TypeAttr, any type attribute

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
spmatA - Single, GPU_SparseSpMatHandle, sparse matrix handle type
dnmatB - Single, GPU_SparseDnTensorHandle, dense tensor handle type
dnmatC - Single, GPU_SparseDnTensorHandle, dense tensor handle type
buffers - Variadic, AnyMemRef, variadic of memref of any non-token type values

Results

asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.spmm operation performs the SpMM operation on the given sparse and dense matrix, and buffer. The operation expects handles returned by previous sparse operations to construct an environment and the operands for SpMM. The buffer must have been allocated on the device.

The matrix arguments can also be associated with one of the following operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value is NON_TRANSPOSE.

Example:

%token = gpu.spmm async [%dep] %spmatA{TRANSPOSE}, %dnmatB{TRANSPOSE}, %dnmatC, %buffers : type($buffers) into f32

spmm_buffer_size()

Return op name gpu.spmm_buffer_size as a bitstring.

spmm_buffer_size(ssa)

gpu.spmm_buffer_size - Precompute buffersize for SpMM operation

Attributes

modeA - Single, GPU_TransposeModeAttr, transpose mode of sparse matrix supported by sparse tensor ops
modeB - Single, GPU_TransposeModeAttr, transpose mode of sparse matrix supported by sparse tensor ops
computeType - Single, TypeAttr, any type attribute

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
spmatA - Single, GPU_SparseSpMatHandle, sparse matrix handle type
dnmatB - Single, GPU_SparseDnTensorHandle, dense tensor handle type
dnmatC - Single, GPU_SparseDnTensorHandle, dense tensor handle type

Results

bufferSzs - Variadic, Index, variadic of index
asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.spmm_buffer_size operation returns the buffer size required to perform the SpMM operation on the given sparse and dense matrix. The operation expects handles returned by previous sparse operations to construct an environment and the operands for SpMM.

The matrix arguments can also be associated with one of the following operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value is NON_TRANSPOSE.

Example:

%bufferszs, %token = gpu.spmm_buffer_size async [%dep] %spmatA{TRANSPOSE}, %dnmatB{TRANSPOSE}, %dnmatC : i64 into f32

spmv()

Return op name gpu.spmv as a bitstring.

spmv(ssa)

gpu.spmv - SpMV operation

Attributes

modeA - Single, GPU_TransposeModeAttr, transpose mode of sparse matrix supported by sparse tensor ops
computeType - Single, TypeAttr, any type attribute

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
spmatA - Single, GPU_SparseSpMatHandle, sparse matrix handle type
dnX - Single, GPU_SparseDnTensorHandle, dense tensor handle type
dnY - Single, GPU_SparseDnTensorHandle, dense tensor handle type
buffer - Single, AnyMemRef, memref of any non-token type values

Results

asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.spmv operation performs the SpMV operation on the given sparse matrix, dense vectors, and buffer. The operation expects handles returned by previous sparse operations to construct an environment and the operands for SpMV. The buffer must have been allocated on the device.

The matrix arguments can also be associated with one of the following operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value is NON_TRANSPOSE.

Example:

%token = gpu.spmv async [%dep] %spmatA{TRANSPOSE}, %dnX, %dnY : memref<?xf64> into bf16

spmv_buffer_size()

Return op name gpu.spmv_buffer_size as a bitstring.

spmv_buffer_size(ssa)

gpu.spmv_buffer_size - Precompute buffersize for SpMV operation

Attributes

modeA - Single, GPU_TransposeModeAttr, transpose mode of sparse matrix supported by sparse tensor ops
computeType - Single, TypeAttr, any type attribute

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type
spmatA - Single, GPU_SparseSpMatHandle, sparse matrix handle type
dnX - Single, GPU_SparseDnTensorHandle, dense tensor handle type
dnY - Single, GPU_SparseDnTensorHandle, dense tensor handle type

Results

bufferSz - Single, Index, index
asyncToken - Optional, GPU_AsyncToken, async token type

Description

The gpu.spmv_buffer_size operation returns the buffer size required to perform the SpMV operation on the given sparse matrix and dense vectors. The operation expects handles returned by previous sparse operations to construct an environment and the operands for SpMV.

The matrix arguments can also be associated with one of the following operators: NON_TRANSPOSE, TRANSPOSE, CONJUGATE_TRANSPOSE. The default value is NON_TRANSPOSE.

Example:

%buffersz, %token = gpu.spmv_buffer_size async [%dep] %spmatA{TRANSPOSE}, %dnX, %dnY into f32

subgroup_broadcast()

Return op name gpu.subgroup_broadcast as a bitstring.

subgroup_broadcast(ssa)

gpu.subgroup_broadcast - Broadcasts a value from the specific lane across subgroup

This op has support for result type inference.

Attributes

broadcast_type - Single, GPU_BroadcastTypeAttr, a lane to broadcast from

Operands

src - Single, AnyType, any non-token type
lane - Optional, I32, 32-bit signless integer

Results

result - Single, AnyType, any non-token type

Description

Broadcasts a value from one lane to all active lanes in a subgroup. The result is guaranteed to be uniform across the active lanes in subgroup.

The possible broadcast types are:

first_active_lane - broadcasts the value from the first active lane in the subgroup.
specific_lane - broadcasts from the specified lane. The lane index must be uniform and within the subgroup size. The result is poison if the lane index is invalid, non subgroup-uniform, or if the source lane is not active.

subgroup_id()

Return op name gpu.subgroup_id as a bitstring.

subgroup_id(ssa)

gpu.subgroup_id

This op has support for result type inference.

Attributes

upper_bound - Optional, IndexAttr, index attribute

Results

result - Single, Index, index

Description

Returns the subgroup id, i.e., the index of the current subgroup within the workgroup.

Example:

%sgId = gpu.subgroup_id : index

Executions where there are more than upper_bound subgroups per workgroup cause undefined behavior. There is an implicit upper bound of kMaxDim (currently uint32_t::max).

subgroup_mma_compute()

Return op name gpu.subgroup_mma_compute as a bitstring.

subgroup_mma_compute(ssa)

gpu.subgroup_mma_compute - GPU warp synchronous matrix multiply accumulate

This op has support for result type inference.

Attributes

a_transpose - Optional, UnitAttr, unit attribute
b_transpose - Optional, UnitAttr, unit attribute

Operands

opA - Single, anonymous/composite constraint, gpu.mma_matrix of 8-bit signed integer or 8-bit unsigned integer or 16-bit float or 32-bit float or 64-bit float values
opB - Single, anonymous/composite constraint, gpu.mma_matrix of 8-bit signed integer or 8-bit unsigned integer or 16-bit float or 32-bit float or 64-bit float values
opC - Single, anonymous/composite constraint, gpu.mma_matrix of 32-bit signless integer or 16-bit float or 32-bit float or 64-bit float values

Results

res - Single, GPU_MMAMatrix, MMAMatrix type

Description

The gpu.subgroup_mma_compute operation performs a matrix-multiply accumulate (mma) operation using all the threads in a subgroup.

This operation takes three !gpu.mma_matrixs as arguments: these hold A, B and Coperands for the mma operation. The operation performed is represented as C += A * B. The op returns a !gpu.mma_matrix which contains the result of the operation held by all threads in a subgroup. a_transpose or b_transpose if present, signify that the respective operand was loaded in a transposed manner. The transpose operands are required to map to correct underlying intrisics but they currently do not seem to affect correctness even if they are absent given that the operands were loaded correctly using the transpose attribute in gpu.subgroup_mma_load_matrix op.

For integer types, the A and B matrices carry their signedness with their types. The accumulator type is expected to be signless and imply a signed integer with a greater width than the other two operands.

This op is meant to be used along with gpu.subgroup_mma_store_matrix and gpu.subgroup_mma_load_matrix ops.

Example:

%D = gpu.subgroup_mma_compute_matrix %A, %B, %C :
  !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp">>
  -> !gpu.mma_matrix<16x16xf16, "COp">

subgroup_mma_constant_matrix()

Return op name gpu.subgroup_mma_constant_matrix as a bitstring.

subgroup_mma_constant_matrix(ssa)

gpu.subgroup_mma_constant_matrix - GPU warp synchronous constant matrix

Operands

value - Single, anonymous/composite constraint, 8-bit signed integer or 8-bit unsigned integer or 32-bit signless integer or 16-bit float or 32-bit float

Results

res - Single, GPU_MMAMatrix, MMAMatrix type

Description

The gpu.subgroup_mma_constant_matrix creates a !gpu.mma_matrix with constant elements.

The operation takes a scalar input and return a !gpu.mma_matrix where each element of is equal to the operand constant. The destination mma_matrix type must have elememt type equal to the constant type. Since the layout of !gpu.mma_matrix is opaque this only support setting all the elements to the same value.

This op is meant to be used along with gpu.subgroup_mma_compute.

Example:

 %0 = gpu.subgroup_mma_constant_matrix %a :
   !gpu.mma_matrix<16x16xf16, "AOp">
 %1 = gpu.subgroup_mma_constant_matrix %b :
   !gpu.mma_matrix<16x16xf32, "COp">

subgroup_mma_elementwise()

Return op name gpu.subgroup_mma_elementwise as a bitstring.

subgroup_mma_elementwise(ssa)

gpu.subgroup_mma_elementwise - GPU warp elementwise operation on a matrix

Attributes

opType - Single, MMAElementWiseAttr, elementwise operation to apply to mma matrix

Operands

args - Variadic, GPU_MMAMatrix, variadic of MMAMatrix type

Results

res - Single, GPU_MMAMatrix, MMAMatrix type

Description

The gpu.subgroup_mma_elementwise takes !gpu.mma_matrix inputs and compute a new !gpu.mma_matrix by applying an elementwise operation to each element.

Since the operation is elementwise and the matrix type must match, the matrix elements are processed independently of the matrix layout.

This op is meant to be used along with gpu.subgroup_mma_compute.

Example:

 %0 =  %A, %B { opType = "ADD" } :
  (!gpu.mma_matrix<16x16xf16, "COp">, !gpu.mma_matrix<16x16xf16, "COp">)
  -> !gpu.mma_matrix<16x16xf16, "COp">

subgroup_mma_extract_thread_local()

Return op name gpu.subgroup_mma_extract_thread_local as a bitstring.

subgroup_mma_extract_thread_local(ssa)

gpu.subgroup_mma_extract_thread_local - Extract a value from GPU warp by invocation and indices

This op has support for result type inference.

Operands

matrix - Single, GPU_MMAMatrix, MMAMatrix type
indices - Variadic, Index, variadic of index

Results

res - Single, AnyIntegerOrFloat, Integer or Float

Description

The gpu.subgroup_mma_extract_thread_local operation extracts a value from !gpu.mma_matrix that is stored at subgroup level.

This operation takes !gpu.mma_matrix as its first operand. It is the source matrix across a subgroup. The op returns a scalar value stored in the invocation in the subgroup.

Since matrix is packed into the the threads within a subgroup, indices are the indices into the values stored by each thread. That is, an index of 0 (or [0, 0]) does not necessarily refer to the first element of the matrix, but the first element that a particular thread holds.

The mapping of matrix elements to threads is not defined by this operation and may not be defined by some lowerings (such as the lowering to SPIR-V). However, if the size of the subgroup is S, then subgroup_mma_extract_thread_local at each index in [0, (M * N) / S) will have the entire matrix extracted across the subgroup.

Example:

%c0 = arith.constant 0 : index
%val = gpu.subgroup_mma_extract_thread_local %m[%c0] : !gpu.mma_matrix<16x16xf32, "AOp"> -> f32

subgroup_mma_insert_thread_local()

Return op name gpu.subgroup_mma_insert_thread_local as a bitstring.

subgroup_mma_insert_thread_local(ssa)

gpu.subgroup_mma_insert_thread_local - Insert a value into GPU warp by invocation and indices

Operands

value - Single, AnyIntegerOrFloat, Integer or Float
matrix - Single, GPU_MMAMatrix, MMAMatrix type
indices - Variadic, Index, variadic of index

Results

res - Single, GPU_MMAMatrix, MMAMatrix type

Description

The gpu.subgroup_mma_insert_thread_local operation inserts a value to !gpu.mma_matrix that is stored at subgroup level.

This operation takes scalar value as its first operand and !gpu.mma_matrix as its second operand. The op inserts the scalar value to the matrix.

The mapping of matrix elements to threads is not defined by this operation and may not be defined by some lowerings (such as the lowering to SPIR-V). However, if the size of the subgroup is S, then subgroup_mma_insert_thread_local at each index in [0, (M * N) / S) will have the entire matrix inserted across the subgroup.

The op returns !gpu.mma_matrix with the updated value.

Example:

%c0 = arith.constant 0 : index
%s0 = gpu.subgroup_mma_insert_thread_local %val, %m[%c0] : f16, !gpu.mma_matrix<16x16xf16, "COp">
        -> !gpu.mma_matrix<16x16xf16, "COp">

subgroup_mma_load_matrix()

Return op name gpu.subgroup_mma_load_matrix as a bitstring.

subgroup_mma_load_matrix(ssa)

gpu.subgroup_mma_load_matrix - GPU warp synchronous matrix load

Attributes

leadDimension - Single, IndexAttr, index attribute
transpose - Optional, UnitAttr, unit attribute

Operands

srcMemref - Single, GPU_MMAMemRef, memref of 8-bit signless integer or 32-bit signless integer or 16-bit float or 32-bit float or 64-bit float or vector of 8-bit signless integer or 32-bit signless integer or 16-bit float or 32-bit float or 64-bit float values of ranks 1 values
indices - Variadic, Index, variadic of index

Results

res - Single, GPU_MMAMatrix, MMAMatrix type

Description

The gpu.subgroup_mma_load_matrix operation loads a matrix collectively using all the threads in a subgroup.

This operation takes a memref as its first operand: it is the source matrix from which data is to be loaded. The op returns a !gpu.mma_matrix. The source memref can be in global memory or shared memory. The load address is determined using indices. The matrix being loaded into is the result. The leadDimension attribute specifies the leading dimension size of the source matrix which eventually allows the lowering to determine the size of each row. If the transpose attribute is present then the op does a transposed load.

The memory indices along each dimension must be in-bounds for that dimension as with an ordinary memref.load.

For integer types, the resulting !gpu.mma_matrix type needs to specify the signedness of the data if the matrix type is an A or B operand for gpu.subgroup_mma_compute.

This op is often meant to be used along with gpu.subgroup_mma_store_matrix and gpu.subgroup_mma_compute.

Example:

 %0 = gpu.subgroup_mma_load_matrix src[%i,%j] : {leadDimension = 32 : i32}
      : memref<32x32xf16, 3>, !gpu.mma_matrix<16x16xf16, "AOp">

subgroup_mma_store_matrix()

Return op name gpu.subgroup_mma_store_matrix as a bitstring.

subgroup_mma_store_matrix(ssa)

gpu.subgroup_mma_store_matrix - GPU warp synchronous matrix store

Attributes

leadDimension - Single, IndexAttr, index attribute
transpose - Optional, UnitAttr, unit attribute

Operands

src - Single, anonymous/composite constraint, gpu.mma_matrix of 8-bit signed integer or 8-bit unsigned integer or 32-bit signless integer or 16-bit float or 32-bit float or 64-bit float values
dstMemref - Single, GPU_MMAMemRef, memref of 8-bit signless integer or 32-bit signless integer or 16-bit float or 32-bit float or 64-bit float or vector of 8-bit signless integer or 32-bit signless integer or 16-bit float or 32-bit float or 64-bit float values of ranks 1 values
indices - Variadic, Index, variadic of index

Description

The gpu.subgroup_mma_store_matrix operation stores a matrix collectively using all the threads in a subgroup.

This operation takes a !gpu.mma_matrix and a memref as operands. !gpu.mma_matrix is the source value containing the data to be stored into the destination memref which can be in global or shared memory. The store address is determined using the indices provided. The leadDimension attribute specifies the leading dimension of the destination matrix. If the transpose attribute is present then the op does a transposed store.

This op is often meant to be used along with gpu.subgroup_mma_load_matrix and gpu.subgroup_mma_compute.

The memory indices along each dimension must be in-bounds for that dimension as with an ordinary memref.load.

Example:

gpu.subgroup_mma_store_matrix %D, %sg[%i,%j] : { leadDimension = 32 : i32}
                : !gpu.mma_matrix<16x16xf16, "COp">, memref<32x32xf16, 3>

subgroup_reduce()

Return op name gpu.subgroup_reduce as a bitstring.

subgroup_reduce(ssa)

gpu.subgroup_reduce - Reduce values among subgroup.

This op has support for result type inference.

Attributes

op - Single, GPU_AllReduceOperationAttr, built-in reduction operations supported by gpu.allreduce.
uniform - Optional, UnitAttr, unit attribute
cluster_size - Optional, I32Attr, 32-bit signless integer attribute
cluster_stride - Single, I32Attr, 32-bit signless integer attribute

Operands

value - Single, AnyIntegerOrFloatOr1DVector, Integer or Float or fixed-length vector of Integer or Float values of ranks 1

Results

result - Single, AnyIntegerOrFloatOr1DVector, Integer or Float or fixed-length vector of Integer or Float values of ranks 1

Description

The subgroup_reduce op reduces the values of lanes (work items) across a subgroup.

The subgroup is divided into clusters starting at lane index 0. Within each cluster, there are size lanes, and the lane index advances by stride. A reduction is done for each cluster in parallel: every lane in the cluster is reduced, and the result is equal for all lanes in the cluster. If size is omitted, there is a single cluster covering the entire subgroup. If stride is omitted, the stride is 1 (the cluster's lanes are contiguous).

When the reduced value is of a vector type, each vector element is reduced independently. Only 1-d vector types are allowed.

Example:

%1 = gpu.subgroup_reduce add %a : (f32) -> f32
%2 = gpu.subgroup_reduce add %b : (vector<4xf16>) -> vector<4xf16>
%3 = gpu.subgroup_reduce add %c cluster(size = 4) : (f32) -> f32
%3 = gpu.subgroup_reduce add %c cluster(size = 4, stride = 2) : (f32) -> f32

If uniform flag is set either none or all lanes of a subgroup need to execute this op in convergence.

The reduction operation must be one of:

Integer types: add, mul, minui, minsi, maxui, maxsi, and, or, xor
Floating point types: add, mul, minnumf, maxnumf, minimumf, maximumf

subgroup_size()

Return op name gpu.subgroup_size as a bitstring.

subgroup_size(ssa)

gpu.subgroup_size

This op has support for result type inference.

Attributes

upper_bound - Optional, IndexAttr, index attribute

Results

result - Single, Index, index

Description

Returns the number of threads within a subgroup.

Example:

%sgSz = gpu.subgroup_size : index

Executions where the number of threads per subgroup exceed upper_bound cause undefined behavior. When no upper_bound is specified, range analyses and similar machinery assume the default bound of kMaxSubgroupSize, currently 128.

terminator()

Return op name gpu.terminator as a bitstring.

terminator(ssa)

gpu.terminator - Terminator for GPU launch regions.

Description

A terminator operation for regions that appear in the body of gpu.launch operation. These regions are not expected to return any value so the terminator takes no operands.

thread_id()

Return op name gpu.thread_id as a bitstring.

thread_id(ssa)

gpu.thread_id

This op has support for result type inference.

Attributes

dimension - Single, GPU_DimensionAttr, a dimension, either 'x', 'y', or 'z'
upper_bound - Optional, IndexAttr, index attribute

Results

anonymous - Single, Index, index

Description

Returns the thread id, i.e. the index of the current thread within the block along the x, y, or z dimension.

Example:

%tIdX = gpu.thread_id x

If upper_bound is set, or if one can be inferred from known_block_size-type annotations in context, executions where the thread index would be greater than or equal to that bound cause undefined behavior.

There is an implicit upper bound of kMaxDim (currently uint32_t::max).

wait()

Return op name gpu.wait as a bitstring.

wait(ssa)

gpu.wait - Wait for async gpu ops to complete.

Operands

asyncDependencies - Variadic, GPU_AsyncToken, variadic of async token type

Results

asyncToken - Optional, GPU_AsyncToken, async token type

Description

This op synchronizes the host or the device with a list of dependent ops.

If the op contains the async keyword, it returns a new async token which is synchronized with the op arguments. This new token is merely a shortcut to the argument list, and one could replace the uses of the result with the arguments for the same effect. The async version of this op is primarily used to make each async token have a single use during lowering and thereby make forks in async execution explicit. Example usage:

%t0 = gpu.foo async : !gpu.async.token
%t1 = gpu.bar async : !gpu.async.token
%t2 = gpu.wait async [%t0, %t1]
// gpu.baz doesn't run until gpu.foo and gpu.bar have both completed, just
// as if the async dependencies were [%t0, %t1].
%t3 = gpu.baz async [%t2]

If the op does not contain the async keyword, it does not return a new async token but blocks until all ops producing the async dependency tokens finished execution. All dependent memory operations are visible to the host once this op completes. Example usage:

%t0 = gpu.foo async : !gpu.async.token
%t1 = gpu.bar async : !gpu.async.token
// The gpu.wait op blocks until gpu.foo and gpu.bar have completed.
gpu.wait [%t0, %t1]

warp_execute_on_lane_0()

Return op name gpu.warp_execute_on_lane_0 as a bitstring.

warp_execute_on_lane_0(ssa)

gpu.warp_execute_on_lane_0 - Executes operations in the associated region on thread #0 of aSPMD program

Attributes

warp_size - Single, I64Attr, 64-bit signless integer attribute

Operands

laneid - Single, Index, index
args - Variadic, AnyType, variadic of any non-token type

Results

results - Variadic, AnyType, variadic of any non-token type

Description

warp_execute_on_lane_0 is an operation used to bridge the gap between vector programming and SPMD programming model like GPU SIMT. It allows to trivially convert a region of vector code meant to run on a multiple threads into a valid SPMD region and then allows incremental transformation to distribute vector operations on the threads.

Any code present in the region would only be executed on first thread/lane based on the laneid operand. The laneid operand is an integer ID between [0, warp_size). The warp_size attribute indicates the number of lanes in a warp.

Operands are vector values distributed on all lanes that may be used by the single lane execution. The matching region argument is a vector of all the values of those lanes available to the single active lane. The distributed dimension is implicit based on the shape of the operand and argument. the properties of the distribution may be described by extra attributes (e.g. affine map).

Return values are distributed on all lanes using laneId as index. The vector is distributed based on the shape ratio between the vector type of the yield and the result type. If the shapes are the same this means the value is broadcasted to all lanes. In the future the distribution can be made more explicit using affine_maps and will support having multiple Ids.

Therefore the warp_execute_on_lane_0 operations allow to implicitly copy between lane0 and the lanes of the warp. When distributing a vector from lane0 to all the lanes, the data are distributed in a block cyclic way. For example vector<64xf32> gets distributed on 32 threads and map to vector<2xf32> where thread 0 contains vector[0] and vector[1].

During lowering values passed as operands and return value need to be visible to different lanes within the warp. This would usually be done by going through memory.

The region is not isolated from above. For values coming from the parent region not going through operands only the lane 0 value will be accesible so it generally only make sense for uniform values.

Example:

// Execute in parallel on all threads/lanes.
gpu.warp_execute_on_lane_0 (%laneid)[32] {
  // Serial code running only on thread/lane 0.
  ...
}
// Execute in parallel on all threads/lanes.

This may be lowered to an scf.if region as below:

  // Execute in parallel on all threads/lanes.
  %cnd = arith.cmpi eq, %laneid, %c0 : index
  scf.if %cnd {
    // Serial code running only on thread/lane 0.
    ...
  }
  // Execute in parallel on all threads/lanes.

When the region has operands and/or return values:

// Execute in parallel on all threads/lanes.
%0 = gpu.warp_execute_on_lane_0(%laneid)[32]
args(%v0 : vector<4xi32>) -> (vector<1xf32>) {
^bb0(%arg0 : vector<128xi32>) :
  // Serial code running only on thread/lane 0.
  ...
  gpu.yield %1 : vector<32xf32>
}
// Execute in parallel on all threads/lanes.

values at the region boundary would go through memory:

// Execute in parallel on all threads/lanes.
...
// Store the data from each thread into memory and Synchronization.
%tmp0 = memreg.alloc() : memref<128xf32>
%tmp1 = memreg.alloc() : memref<32xf32>
%cnd = arith.cmpi eq, %laneid, %c0 : index
vector.store %v0, %tmp0[%laneid] : memref<128xf32>, vector<4xf32>
some_synchronization_primitive
scf.if %cnd {
  // Serialized code running only on thread 0.
  // Load the data from all the threads into a register from thread 0. This
  // allow threads 0 to access data from all the threads.
  %arg0 = vector.load %tmp0[%c0] : memref<128xf32>, vector<128xf32>
  ...
  // Store the data from thread 0 into memory.
  vector.store %1, %tmp1[%c0] : memref<32xf32>, vector<32xf32>
}
// Synchronization and load the data in a block cyclic way so that the
// vector is distributed on all threads.
some_synchronization_primitive
%0 = vector.load %tmp1[%laneid] : memref<32xf32>, vector<32xf32>
// Execute in parallel on all threads/lanes.

yield()

Return op name gpu.yield as a bitstring.

yield(ssa)

gpu.yield - GPU yield operation

Operands

values - Variadic, AnyType, variadic of any non-token type

Description

gpu.yield is a special terminator operation for blocks inside regions in gpu ops. It returns values to the immediately enclosing gpu op.

Example:

gpu.yield %f0, %f1 : f32, f32