Mosaic GPU Dialect
Reactant.MLIR.Dialects.mosaic_gpu.alloc_smem Method
alloc_smem
Allocates a statically shaped SMEM memref that is valid for the duration of the enclosing alloca_scope operation.
An optional alignment attribute specifies the minimum byte alignment for the allocation and must be a positive power of 2.
Reactant.MLIR.Dialects.mosaic_gpu.alloca_scope Method
alloca_scope
Defines a single-block allocation scope for scoped allocation operations (such as mosaic_gpu.alloc_smem). Any scoped allocations inside this region are valid only for the duration of the region and may not escape.
alloca_scope can capture outer values directly and optionally return non-memref results via mosaic_gpu.alloca_scope.return.
Reactant.MLIR.Dialects.mosaic_gpu.alloca_scope_return Method
alloca_scope_return
The alloca_scope.return op is a terminator that indicates the end of execution within an AllocaScopeOp's region. It can optionally return some values, which become the results of the parent AllocaScopeOp.
The declared results of the parent AllocaScopeOp must match the operand types of this op.
Reactant.MLIR.Dialects.mosaic_gpu.arrive_expect_tx Method
arrive_expect_tx
Executes an arrive.expect_tx operation on the given barrier.
`expect_tx` must be positive or 0. If `expect_tx` is negative, behavior is undefined.Reactant.MLIR.Dialects.mosaic_gpu.assume_multiple Method
assume_multiple
This operation is a hint to the compiler that the input value is guaranteed to be a multiple of multiple. This can be used to satisfy divisibility checks in some compiler passes.
The result is the same as the input value.
Reactant.MLIR.Dialects.mosaic_gpu.async_load Function
async_load
Schedules an async copy of the contents of the source MemRef in GMEM to the destination MemRef in SMEM. The destination MemRef in SMEM must be contiguous.
Upon completion of the copy, the complete-tx(complete-count) operation will be executed on the provided barrier, if one is given.
The barrier is required on Hopper and newer GPUs, which use the TMA implementation, and must be omitted on pre-Hopper GPUs, which use the cp.async implementation. In that case, completion is tracked via cp.async commit groups and has to be awaited with nvvm.cp_async_wait_group rather than the barrier. When the barrier is omitted, collective and leader_tracked must not be set, and oob_fill_mode must be promise_in_bounds.
The indices and slice_lengths inputs define what slice of the GMEM source corresponds to the SMEM destination. Both indices and slice_lengths must have a length equal to the rank of the source. The values in indices are the starting indices of each dimension and the values in slice_lengths are the lengths. Providing -1 in slice_lengths indicates that the slice length is 1 and that the corresponding dimension should be collapsed and does not appear in the destination MemRef. If an index is a vector of ints, its elements serve as GMEM indices from which the data should be gathered from GMEM. In this case, the slice_lengths must have a value equal to the size of the vector. Only one index may be a vector.
The data is written in row-major order to the contiguous SMEM destination. The source data does not need to be contiguous, except for the last (and minor-most) dimension.
The collective attribute can be provided to use TMA multicast to more efficiently load the GMEM data in cases where multiple thread blocks are grouped together in a cluster and need to load the same data. Each block in a cluster will first load a slice from GMEM to SMEM and then the slices will be multicast to all other blocks in the cluster. In this way TMA multicast guarantees L2 cache hits. The collective attribute is the list of cluster dimensions along which to partition the input data loads. Only supported with TMA (Hopper and newer).
The leader_tracked attribute can be provided to only track the completion of the copy in the leader block in the cluster. If CopyPartitioned(axis), performs a partitioned collective copy along the given axis. If CopyReplicated, all blocks load the same data. Only supported with TMA (Hopper and newer).
The oob_fill_mode attribute controls the behavior of out-of-bounds accesses. zeros and undefined are only supported with TMA (Hopper and newer); the cp.async implementation only supports promise_in_bounds.
The predicate allows scheduling the transfer conditionally.
The gmem_peer_id parameter specifies a particular device id to load from in a multi-GPU setting.
Reactant.MLIR.Dialects.mosaic_gpu.async_load_im2col Function
async_load_im2col
Schedules an async im2col copy of the contents of the source MemRef in GMEM to the destination MemRef in SMEM. Each row of destination receives the channels of one pixel of one window in source. The destination MemRef in SMEM must be contiguous.
Upon completion of the copy, the complete-tx(complete-count) operation will be executed on the provided barrier.
sourceis a batch of channels-last images of shape(N, W, C),(N, H, W, C)or(N, D, H, W, C)(rank 3, 4 or 5) with an element bitwidth of at least 8.destinationhas shape(*rows, cols)of rank at least 2. Its leading dimensions are flattened intoprod(rows)windows (in[1, 1024]) ofcolschannels each (in[1, 256]), wherecolshas a bytewidth divisible by 16.window_lower_boundsandwindow_upper_boundsgive, for each spatial dimensioniof sizeS, the(lo, hi)bounds such that windows start at positionslo,lo + s,lo + 2 * s, ... belowS + hi, wheresiswindow_strides[i], with the image implicitly zero-padded outside[0, S). Each bound must fit in a signed16 / (rank - 2)-bit integer.indiceshas length equal to the rank ofsourceand specifies(n, *start, c): the batch index and start position of the window loaded into the first row, and the first channel to load. Each subsequent row takes the next window in row-major order over(n, *start), so a copy can span several images. Rows past the last image are zero-filled.filter_offsetshas lengthrank - 2and selects the pixel to load within each window as an offset in[0, 2**(16 / (rank - 2)) - 1]from the window's start position.window_strideshas lengthrank - 2and gives the step in[1, 8]between consecutive window start positions along each spatial dimension.
The predicate allows scheduling the transfer conditionally.
Reactant.MLIR.Dialects.mosaic_gpu.async_prefetch Function
async_prefetch
Schedules an async prefetch of the contents of the source MemRef in GMEM to the L2 cache, making subsequent loads of the same data from GMEM faster.
The indices and slice_lengths inputs define what slice of the GMEM source is going to be prefetched. Both indices and slice_lengths must have a length equal to the rank of the source. The values in indices are the starting indices of each dimension and the values in slice_lengths are the lengths. Providing -1 in slice_lengths indicates that the slice length is 1. If an index is a vector of ints, its elements serve as GMEM indices from which the data should be gathered from GMEM. In this case, the slice_lengths must have a value equal to the size of the vector. Only one index may be a vector.
The collective attribute can be provided to partition the prefetch over multiple blocks in a cluster.
The predicate allows scheduling the prefetch conditionally.
Reactant.MLIR.Dialects.mosaic_gpu.async_store Function
async_store
Schedules an async store of the contents of the source MemRef in SMEM to the destination MemRef in GMEM. The source MemRef in SMEM must be contiguous.
The indices and slice_lengths inputs define what slice of the GMEM destination corresponds to the SMEM source. Both indices and slice_lengths must have a length equal to the rank of the destination. The values in indices are the starting indices of each dimension and the values in slice_lengths are the lengths. Providing -1 in slice_lengths indicates that this dimension is collapsed in the source and needs to be expanded to a slice of size 1 in the destination. If an index is a vector of ints, its elements serve as GMEM indices to which the data should be scattered to GMEM. In this case, the slice_lengths must have a value equal to the size of the vector. Only one index may be a vector.
The data is written in row-major order to the GMEM destination. The source data in SMEM needs to be contiguous, but the destination GMEM does not.
The predicate allows scheduling the transfer conditionally.
The reduction_op attribute can be provided to perform a reduction when storing to GMEM. For example, using add will add the SMEM values to existing values in GMEM.
The gmem_peer_id and is_global_broadcast parameters define the destination of the store in a multi-GPU setting. gmem_peer_id specifies a particular device id to store to, while is_global_broadcast indicates that the value should be broadcast to all peers. These two parameters are mutually exclusive (i.e. both cannot be set).
Reactant.MLIR.Dialects.mosaic_gpu.async_store_scales_smem_to_tmem Method
async_store_scales_smem_to_tmem
Schedules an async copy of the scaling factors from the source MemRef in SMEM to the destination MemRef in TMEM.
The source and destination MemRefs must have the same element type (F8E8M0FNU or F8E4M3FN).
Given a destination shape [M, N], the source shape must be [M/128, N/4, 32, 16], and M % 128 == 0 and N % 4 == 0.
If collective is true, the copy is performed by a cluster of 2 CTAs, and we expect that this op is issued by a single CTA. If collective is false, the copy is performed by a single CTA.
Reactant.MLIR.Dialects.mosaic_gpu.async_store_smem Method
async_store_smem
cluster_dim and cluster_idx specify the target block.
The target block can wait on the provided barrier (its local copy) to observe the completion of this copy.
When atomic is specified, the store performs an atomic reduction on target SMEM instead of overwriting it.
Reactant.MLIR.Dialects.mosaic_gpu.async_store_smem_to_tmem Method
async_store_smem_to_tmem
Schedules an async copy of the contents of the source MemRef in SMEM to the destination MemRef in TMEM.
The source and destination MemRefs must have the same element type, and shape.
If collective is true, the copy is performed by a cluster of 2 CTAs, and we expect that this op is issued by a single CTA. If collective is false, the copy is performed by a single CTA.
Reactant.MLIR.Dialects.mosaic_gpu.async_store_sparse_metadata_smem_to_tmem Method
async_store_sparse_metadata_smem_to_tmem
Schedules an async copy of the sparse metadata from the source MemRef in SMEM to the destination MemRef in TMEM.
The source and destination MemRefs must have the same element type (i2). Given a destination shape of [M, N], the source shape must be [M/128, N/64, 128, 64], and M % 128 == 0 and N % 64 == 0.
If collective is true, the copy is performed by a cluster of 2 CTAs, and we expect that this op is issued by a single CTA. If collective is false, the copy is performed by a single CTA.
Reactant.MLIR.Dialects.mosaic_gpu.broadcast_in_dim Method
broadcast_in_dim
The operand is broadcast to the shape described by the result type. broadcast_dimensions maps the dimensions of operand to the dimensions of the target shape, i.e. the i'th dimension of the operand is mapped to the broadcast_dimensions[i]'th dimension of the output shape. The dimensions of the operand must have size 1 or be the same size as the dimension in the output shape they are mapped to.
This op matches the semantics of lax.broadcast_in_dim.
Reactant.MLIR.Dialects.mosaic_gpu.broadcasted_iota Method
broadcasted_iota
Creates an array that has the specified shape and holds values starting at zero and incrementing by one along the specified dimension.
sourceReactant.MLIR.Dialects.mosaic_gpu.custom_primitive Method
custom_primitive
Allows defining a custom Mosaic GPU primitive.
Custom primitives should carry input and output layouts for each of their vector operands and outputs, and input transforms for each of their memref operands that live in SMEM.
Custom primitives can only return vectors.
sourceReactant.MLIR.Dialects.mosaic_gpu.get_cluster_ref Function
get_cluster_ref
Maps a memref to a peer block in the cluster along specified dimensions.
sourceReactant.MLIR.Dialects.mosaic_gpu.initialize_barrier Method
initialize_barrier
Initializes num_barriers barriers each meant to synchronize exactly arrival_count threads.
base_pointer must be a pointer to a shared memory location.
Reactant.MLIR.Dialects.mosaic_gpu.layout_cast Method
layout_cast
Casts a vector value to a new strided or tiled layout.
If strict is true, only allows relayouts that are free at runtime and forces layout equality otherwise.
Reactant.MLIR.Dialects.mosaic_gpu.mma Method
mma
Computes a @ b + accumulator synchronously using Ampere MMA instructions.
The output has an identical shape and type as the input accumulator.
sourceReactant.MLIR.Dialects.mosaic_gpu.multimem_load_reduce Method
multimem_load_reduce
Loads from a multicast memref specified by the source and reduces the loaded values into a single output using hardware-accelerated multimem instructions.
The reduction is defined by reduction_op and the devices which participate in the operation are backed by a multicast pointer source.
And, Or and Xor reduction types are compatible only with int inputs. Umin, Umax, Smin and Smax may only be specified for int types.
sourceReactant.MLIR.Dialects.mosaic_gpu.query_cluster_cancel Method
query_cluster_cancel
It interprets the 16-byte opaque response written to shared memory by a completed try_cluster_cancel call to determine if a new work unit was successfully claimed.
If success is true, then x, y, and z are set to the CTA ID of the first CTA that is part of the cluster that was claimed.
Reactant.MLIR.Dialects.mosaic_gpu.reinterpret_cast Method
reinterpret_cast
Casts a memref to a new memref type with a potentially different shape but the same total number of elements and the same element type. This is a pure reinterpretation of the shape and layout; no data is moved or copied.
Unlike memref.reinterpret_cast, this op does not take explicit offsets, sizes, or strides. The result type fully determines the reinterpretation.
Reactant.MLIR.Dialects.mosaic_gpu.return_ Method
return_
The return op is a terminator that indicates the end of execution within a CustomPrimitiveOp or RunScopedOp region. It can optionally return some values, which become the results of the parent op.
The declared results of the parent op must match the operand types of this op.
sourceReactant.MLIR.Dialects.mosaic_gpu.run_scoped Method
run_scoped
Evaluates a single-block region where block arguments represent scoped allocations whose lifetimes are confined to the region.
Each block argument must be a statically shaped memref in SMEM, with a corresponding #mosaic_gpu.smem_alloc entry in buffer_attrs.
The region may capture outer values and return non-memref results via mosaic_gpu.return.
Reactant.MLIR.Dialects.mosaic_gpu.slice_smem Method
slice_smem
Constructs a SMEM ref starting at the specified offset.
slice_smem may optionally carry an alias_id attribute that is used to differentiate between potentially aliasing allocations.
Reactant.MLIR.Dialects.mosaic_gpu.slice_tmem Method
slice_tmem
The principal use case for this op is to do a single TMEM allocation and slice it into multiple smaller TMEM references. source is the large TMEM allocation and offset is the number of columns to start slicing from.
slice_tmem may optionally carry an alias_id attribute that is used to differentiate between potentially aliasing allocations.
Reactant.MLIR.Dialects.mosaic_gpu.tcgen05_commit_arrive Function
tcgen05_commit_arrive
Makes the barrier object track the completion of all prior async tcgen05 operations.
if collective is true, allow signaling on the barrier object of multiple CTAs within the cluster.
Reactant.MLIR.Dialects.mosaic_gpu.tcgen05_mma Function
tcgen05_mma
Schedules tcgen05.mma instructions that perform the following matrix multiply and accumulate:
accumulator += a * b
This operation supports larger inputs than the PTX-level MMA instruction and will schedule as many PTX-level MMA instructions as needed to accomplish the calculation.
The inputs should have the following shapes:
a: [groups_m * m, groups_k * s]
b: [groups_k * s, groups_n * s]
accumulator: [groups_m * m, groups_n * s]
where s == swizzle / element_bytewidth and m is specified according to https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-matrix-shape.
The accumulator, a and b matrices need to be provided as 2-dimensional memrefs. The accumulator is always in TMEM and b is always in SMEM. a can be in TMEM or SMEM. a and b must have the same element type and when a is in TMEM only F16 or BF16 are supported.
a_scale and b_scale are optional scaling matrices that reside in TMEM. When set the operation is defined as:
accumulator += (a * a_scale) * (b * b_scale)
a_sparse_metadata is an optional sparse metadata matrix that resides in TMEM. When set, a is treated as a sparse matrix and the metadata encodes the sparsity pattern. With sparse MMA, a's contracting dimension is halved relative to b's contracting dimension. a_sparse_metadata must have the same shape as a.
accumulate is a boolean that indicates whether to perform the accumulate step.
Reactant.MLIR.Dialects.mosaic_gpu.tmem_alloc Method
tmem_alloc
This op allocates a chunk of TMEM and stores the pointer to the memory in the provided SMEM memref.
The smem_ptr is a pointer in SMEM where a pointer to the allocated TMEM will be stored. The op returns a memref to the allocated TMEM. The result must have a shape with dimensions [rows, logical_columns]. If packing is 1, then the number of logical (unpacked) columns is equal to the number of allocated columns in TMEM. Otherwise, these constraints must hold:
packing = 32 / bitwidth(element type of result)
unpacked_columns = allocated_columns * packingThe number of allocated columns in TMEM can be any power of two in the range [32, 512]. If the calculated number of allocated columns is less than 32 or not a power of two, then it will be rounded up to the nearest power of two larger or equal to 32.
If collective is true 2 CTAs will perform the allocation collectively, otherwise, only one CTA will perform the allocation.
Reactant.MLIR.Dialects.mosaic_gpu.tmem_relinquish_alloc_permit Method
tmem_relinquish_alloc_permit
The instruction specifies that the CTA of the executing thread is relinquishing the right to allocate Tensor Memory. So, it is illegal for a CTA to perform tmem_alloc after any of its constituent threads execute tmem_relinquish_alloc_permit.
If collective is true, applies to collective TMEM allocations.
Reactant.MLIR.Dialects.mosaic_gpu.try_cluster_cancel Function
try_cluster_cancel
It allows an SM to dynamically acquire work by atomically canceling the launch of a pending cluster from the grid and claiming its CTA ID as the next unit of work.
sourceReactant.MLIR.Dialects.mosaic_gpu.vector_concat Method
vector_concat
Concatenates a variadic list of operands along dimension.
All operands must have the same rank and element type. This op matches the semantics of lax.concatenate.
Reactant.MLIR.Dialects.mosaic_gpu.vector_load Method
vector_load
Similar to vector.load (vector dialect) but supports loading from non-contiguous memory.
If optimized is true, raises an error if we cannot generate an optimised transfer. optimized defaults to true when unset.
Reactant.MLIR.Dialects.mosaic_gpu.vector_store Method
vector_store
Similar to vector.store (vector dialect) but supports storing to non-contiguous memory.
If optimized is true, raises an error if we cannot generate an optimised transfer. optimized defaults to true when unset.
If atomic_type is set, performs an atomic store-(add|min|...) of the value.
If multimem is true, then the store is performed on all the devices in the underlying mesh using multimem instructions.
Reactant.MLIR.Dialects.mosaic_gpu.wgmma Method
wgmma
Schedules WGMMA operations that perform the following matrix multiply and accumulate:
accumulator = a * b + accumulator
This operation supports larger inputs than the PTX-level WGMMA operation and will schedule as many PTX-level WGMMA operations as needed to accomplish the calculation. The b matrix, and optionally a, need to be provided as a 2-dimensional memref.
The inputs should have the following shapes:
a: [groups_m * 64, groups_k * s]
b: [groups_k * s, groups_n * s]
accumulator: [groups_m * 64, groups_n * s]
where s == swizzle / element_bytewidth.
The output has an identical shape and type as the input accumulator.
The accumulator is always in registers and b is always in shared memory. a and b must have the same element type and when a is in registers only F16 or BF16 are supported.
The accumulator must be a vector with a FragmentedLayout. The WGMMA operation will be executed in the async proxy and any inputs in registers need to be synchronized with a memory fence. Note that the caller is responsible for fencing the accumulator prior to issuing the op, whereas a is fenced automatically in the lowering if it is passed in registers.
Usually a is read from shared memory if it is used directly in the WGMMA operation. If a needs to be transformed before it is used in the WGMMA operation, it may be more convenient to read it directly form registers. This avoids the need to store the data and wait for a fence.
Reactant.MLIR.Dialects.mosaic_gpu.with_transforms Method
with_transforms
This op enforces the provided transforms on the parameter memref.
source