API reference

Layout(shape, strides)

A hierarchical coordinate-to-element-offset map. Shape and stride trees must match, with positive extents and nonnegative strides. Zero strides represent broadcasting; a layout is not necessarily injective. Use static(n) at leaves that must specialize; ordinary integers remain runtime values.

l((row, col)) evaluates zero-based coordinates. An integer coordinate is decomposed first-mode-fastest, recursively. Evaluation assumes in-range coordinates; memory views provide checked windows.

LocalOwnership{N,Axis}()
StripedOwnership{N,Axis}()

Two explicit warp ownership patterns. Each lane holds N local values along logical Axis (1 or 2). LocalOwnership assigns one independent sequence to each lane; StripedOwnership interleaves the 32 lanes within one sequence. Neither describes memory strides. Coordinates take a zero-based lane (0:31), local to this warp, and a zero-based value slot. Use Ownership for an explicit mapping.

Ownership(Val((rows, cols)), thread_value_layout)

Map (lane, logical_value) to a matrix coordinate through a hierarchical layout. The layout's output is a column-major logical index, independent of physical memory layout and register packing.

Swizzle{Bits,Base,Shift}()

XOR two disjoint bit fields in an element offset. Base low bits remain unchanged; positive Shift moves the source field right. It is its own inverse. Units come from the inner layout, not implicitly from bytes.

Drop hierarchy and merge adjacent modes when their strides are contiguous.

Function composition: compose(f,l)(c) == f(l(c)); retains the inner domain.

Number of storage elements up to the largest offset, including holes.

tile(l::Layout, Val((m, n, ...)))

Factor each flat mode into (within_tile, tile_index), first component fastest. For example a 16×32 matrix tiled by 8×16 has shape ((8,2),(16,2)). This algebraic operation changes coordinates, not storage or ownership.

window(layout, origin, Val(shape))

A static-size logical window with a potentially runtime origin. Offsets stay relative to the parent's allocation base, preserving swizzle phase. Bounds are checked unless the caller uses @inbounds after establishing validity.

ownership_table(ownership) -> Matrix{Tuple{Int,Int}}

Logical coordinate of every (thread, slot) pair, zero-based, as a matrix indexed by (thread+1, slot+1).

The address operand of a copy atom: the 16-byte row each lane supplies.

CopyAtom{:load|:store,Trans}()
CopyAtom(op, trans=false)

A warp-collective copy of one 8×8 matrix of 16-bit units between shared memory and registers: ldmatrix for :load, stmatrix for :store, with the .trans modifier when Trans is true. Like an MMAAtom, the atom contributes only ownerships. They use memory coordinates: axis 1 indexes the eight 16-byte rows, axis 2 the eight units of a row. operand_layout with Addresses gives the row each lane addresses, and with Registers the two units each lane holds in one register word.

Copies of larger ownerships derive from these tables. matrix_copy_plan covers a target ownership with 8×8 blocks along the axis that memory stores contiguously, chooses the transposed variant when the register pattern requires it, groups blocks into .x1, .x2 and .x4 instructions and permutes the result words into the target's slot order. Adjacent pairs of 8-bit elements form one unit, so 8-bit operands qualify as well. load_fragment and store! use the derived copy for shared tiles.

CopyPlan{Shape,Threads,Axis}()

Distribute a matrix's 16-byte vectors over Threads threads, with vectors running along logical Axis (1 or 2). Shape is static; source and destination may have different memory layouts. This plan issues cp.async copies only. Call commit_copies() and wait_copies(Val(groups_remaining)) explicitly, then synchronize the consumers before reading shared memory.

Call validate_copy(plan, T, destination_layout, source_layout) on the host to check vector contiguity/alignment and unique destination ownership. Raw pointers must additionally be 16-byte aligned. Reusing a shared allocation requires synchronization with its readers, separately from copy completion.

Float8E4M3

OCP FP8 E4M3 (no infinities, NaN at 0x7f/0xff), converting from FP32 with round-to-nearest-even and saturation like cvt.rn.satfinite.e4m3x2.f32.

Float8E5M2

FP8 E5M2 with IEEE-style infinities and NaNs, converting from FP32 with round-to-nearest-even and saturation like cvt.rn.satfinite.e5m2x2.f32.

Fragment(values::NTuple, ownership)

This thread's immutable share of a logical tile. The ownership layout maps (thread, value_slot) to logical coordinates; it does not describe memory strides, and the values need not form a row or be contiguous. Arithmetic preserves ownership. Collective support depends on that ownership layout.

MMAAtom{(m,n,k),TA,TB,TC}()
MMAAtom((m,n,k), TA, TC=Float32)

A warp-collective mma.sync instruction described by its shape and element types. The atom contributes exactly one instruction-specific fact: the thread/value ownership of its A, B and accumulator operands, from operand_layout. Fragments, loads, stores, reductions, tiling and conversions derive from those ownerships. Float32 inputs select the TF32 instruction of that shape. The PTX extension supplies the instruction for each atom it implements; an atom without one is still a valid description.

MatrixCopyPlan

The derived decomposition of a 32-lane ownership into CopyAtom blocks. trans selects the transposed instruction, axis is the logical axis that memory stores contiguously, per_unit the elements per 16-bit unit, grid the block counts along memory rows and columns, and words[b] the fragment word (1-based) that block b transfers, with blocks numbered column-major over the grid. Consecutive blocks group into instructions of up to four matrices.

GlobalTile(pointer, layout, Val(align)=Val(sizeof(T)))
SharedTile(pointer, layout, Val(align)=Val(sizeof(T)))

Borrow typed global (address space 1) or shared (address space 3) memory. The layout maps logical coordinates to element offsets from the allocation base. Pointer addition occurs in bytes only at pointer(tile, coordinate). The caller owns allocation, lifetime, alignment and synchronization.

align declares, in bytes, that the pointer and every stride of the layout other than the unit stride are multiples of align, so that any element whose coordinate along the unit-stride axis is a multiple of align ÷ sizeof(T) has an align-aligned address. Loads and stores of fragments vectorize up to that width. A window's origin along the unit-stride axis must keep the declaration. Bounds checks verify the pointer, the strides and window origins unless elided with @inbounds.

PackedFragment(T, words::NTuple{W,UInt32}, ownership)

Register representation of logical 8- or 16-bit elements of type T. Each word holds 32 ÷ bits elements, lowest element first. Ownership counts logical elements, not words. This constructor interprets bits; pack(T, f) numerically converts values and unpack(p) exposes ordinary typed values.

Bindings of one matrix, under their earlier names.

ReductionPlan

A thread-uniform recipe for reducing one logical axis of an ownership. groups lists the slots (1-based) that share each kept coordinate within a thread; bits lists thread-index bits whose flip preserves the kept coordinate, so an xor shuffle over each of them completes the reduction; result is the replicated result ownership.

The register operand of a copy atom: two units per lane in one word.

SoftmaxState(fragment; dims=2)

FP32 running maxima and unnormalized exponential sums in the fragment's reduced ownership. The empty state is (-Inf, 0). Scores must be finite or -Inf (masked). State arithmetic does not allocate storage or synchronize. The state type depends on the axis; pass dims=Val(axis) when the axis is not a literal.

TMALoad(T, Val(S), Val(A)): the 128-byte-row TMATile, under its earlier name.

TMATile(T, Val((m,k)), Val(2))            # logical A(M,K), 128-byte rows
TMATile(T, Val((k,n)), Val(1), Val(64))   # logical B(K,N), 64-byte rows

A 2D tile that TMA moves between global and shared memory. Elements are one, two or four bytes wide. Axis A is the inner axis and spans one swizzle row of W bytes (32, 64 or 128, default 128); the other extent is a positive multiple of eight, at most 256. Shared storage is shared_layout(plan), allocated at a multiple of 8W bytes. transfer_bytes counts every element of the box, including zero-filled out-of-bounds elements of a load. The plan owns no memory, descriptor or barrier; prepare_tma binds it to a device array.

Tcgen05MMA{(m,n,k),TA,TB,TC}()
Tcgen05MMA((m,n,k), TA, TC=Float32)

A single-thread tcgen05.mma.cta_group::1 instruction described by its shape and element types: .kind::f16 for BF16/FP16, .kind::tf32 for Float32, .kind::f8f6f4 for FP8 and .kind::i8 for INT8 inputs. k is the instruction's fixed K (256 bits of each operand row), m is 128 and n a multiple of 16 up to 256.

Its operands are not thread-owned. A and B are shared-memory encodings, a canonical swizzled core-matrix layout (Swizzle{B,M,3} composed with a layout of packed 32-, 64- or 128-byte rows, which is Tylo's TMA storage; see swizzled_structure), or a TmemTile for A; the accumulator is a TmemTile whose layout is operand_layout with Accumulator. Descriptors and the instruction descriptor derive from those encodings (tcgen05_operand); mma steps the operands along K through their layouts, so K-major and MN-major storage share one path. Completion is asynchronous: commit_mma arrives on an mbarrier once every prior instruction of this thread has finished.

Tcgen05Operand

A shared-memory A or B operand of a Tcgen05MMA: the descriptor of the tile's base, the tile's layout and the logical origin. mma derives each K step's descriptor from the layout, so the same operand serves K-major and MN-major storage. Construct with tcgen05_operand.

TiledMMA(atom, Val((warps_m,warps_n)), Val((repeat_m,repeat_n)), Val(k))

A CTA's MMA decomposition. Each warp computes repeatm×repeatn atoms; K is a positive multiple of the atom's K. Warp numbering is M-fastest. This plan describes ownership and instruction repetition, with no allocation or hidden barriers. Its accumulator is one flat Fragment in TiledMMAOwnership.

TmemTile(T, address::UInt32, layout)

Borrow TMEM with a logical-to-storage layout. Layout offsets count typed slots in a virtual column-major grid with 128 hardware lanes: offset i selects lane i % 128 and typed position i ÷ 128 along its storage. FP32 has one position per hardware word; BF16/FP16 have two. Logical axes may be permuted or partitioned independently of this hardware addressing convention.

The caller owns the allocation, lifetime and synchronization. Construction neither allocates storage nor proves that the view fits the allocation.

TmemTransfer{Shape,Axis}()

A warp-collective TMEM/register transfer using the .32x32b instruction. Shape is the logical partition shape. Each thread holds values along logical Axis; the other axis has extent 32. This plan also describes the resulting register ownership. permutedims(plan) exchanges its logical axes.

Bind it to a storage view with partition(plan, tile). Supported payloads are powers of two from 1 through 128 words per thread: FP32 and packed BF16/FP16 loads and stores. Instruction shape, ownership and storage compatibility are checked explicitly.

WGMMA64(T, Val(n), Val(k)=Val(64), Val(partials)=Val(1))

Hopper shared/shared warpgroup MMA: 128 threads compute a 64×n×k product. BF16/FP16 inputs, FP32 accumulation; n=8:8:256, k in (16,32,64). Each K=16 instruction updates one of partials independent accumulators cyclically. Partials must divide k/16; total registers per thread must not exceed 128. partials=4 at n=8,k=64 preserves Cohere's four independent K chains. Use finish_mma to sum partials after waiting. No allocation or CTA barrier.

permutedims(fragment, (2, 1))

Exchange the two logical axes without moving values between registers or threads. Applying the permutation twice restores the original ownership. For example, maximum(permutedims(f); dims=1) performs the same collective as maximum(f; dims=2), with the result's logical axes exchanged as well. This does not convert values into another MMA operand distribution.

sum(fragment; dims)
maximum(fragment; dims)
minimum(fragment; dims)

Reduce FP32 values over a logical axis, retaining a singleton dimension and explicit result ownership. Broadcast the result back with ordinary dotted arithmetic, for example fragment .- maximum(fragment; dims=2).

The recipe is derived from the ownership: a local tree over the slots that share a kept coordinate, then xor shuffles over the lane bits that replicate it. An axis whose values span warps has no recipe and is rejected rather than silently reducing only this thread's values. Reductions with shuffles require all 32 lanes of each warp. The result type depends on the axis, so dims must be a constant in a kernel; dims=Val(2) states that explicitly.

Consecutive block ranges, each one instruction of up to four matrices.

_reduce_values(op, data, ownership, Val(axis))

Reduce this thread's values along a logical axis using the recipe derived from the ownership: a local tree over the slots sharing each kept coordinate, then xor shuffles over the lane bits that replicate it. One result per kept coordinate held by this thread.

The replicated result ownership of reducing axis, or an error when no recipe exists.

One xor-shuffle exchange at a lane offset.

The ownership instance of a singleton ownership type, or nothing.

Butterfly reduction over lane offsets W/2 down to 1.

accumulator(atom, address::UInt32)

The atom's accumulator as a TmemTile at a TMEM address: lane m holds row m, column n holds column n. The caller owns the allocation; the first mma with accumulate=false initializes it.

broadcast_slots(x, anchor) -> Vector{Int} or nothing

For each slot of anchor, the slot of x supplying its value under array broadcasting: x has the anchor's logical shape, or extent one on the axes it broadcasts along. The mapping must be identical on every thread.

Close this thread's uncommitted cp.async operations as one group.

commit_mma(barrier)

Arrive on a shared mbarrier once all prior tcgen05.mma instructions issued by this thread have completed. Single thread; no wait is implied.

commit_tma_stores()

Close the issuing thread's bulk group of TMA stores.

copy_async!(plan, destination, source, origin::Tuple, thread)

Copy a fixed-capacity tile from logical origin in a bounded global source. Valid aligned contiguous 16-byte vectors use cp.async. Boundary or unaligned vectors use scalar loads/stores, zero-filling every invalid element. Source coordinates are checked before a pointer is formed. Destination shape and vector alignment must satisfy the same contract as the unmasked form.

Both paths require explicit commit/wait and consumer synchronization. A copy wait alone does not publish scalar shared stores to other threads. This form uses scalar zero-fill because the pinned PTX wrapper exposes full-vector copies.

The four matrix copy atoms: loads and stores, plain and transposed.

Order subsequent TMEM operations after a preceding thread synchronization.

Order preceding TMEM operations before a following thread synchronization.

fit_ownership(table, shape) -> Ownership or nothing

An explicit Layouts.Ownership reproducing table over a logical shape, when the table is affine in the bits of the thread and slot indices.

Atoms with an instruction binding.

instruction_descriptor(atom; a_major=:K, b_major=:K) -> UInt32

The tcgen05.mma instruction descriptor of an atom for the given operand majorness: a pure function of the atom's shape and element types, as mma derives it from its operands' encodings.

Every logical coordinate of the ownership's shape is held by some thread.

No logical coordinate is held twice.

load_a(atom, tile, lane)
load_b(atom, tile, lane)

Load an atom's A (m×k) or B (k×n) operand from a shared tile through load_fragment with the operand's ownership.

load_async(partition)

Issue a warp-collective TMEM load using an explicit transfer partition. All 32 lanes must execute with the same partition and the correct warp band. The result is pending; call wait_load before using its values.

load_a(atom, tile, lane)
load_b(atom, tile, lane)

Load an atom's A (m×k) or B (k×n) operand from a shared tile through load_fragment with the operand's ownership.

load_fragment(ownership, tile, thread)

Load this thread's values of a memory tile at the coordinates the ownership assigns to it. Scalar loads are correct for any static ownership. A shared tile of 8- or 16-bit elements whose ownership decomposes into CopyAtom blocks along the layout's static unit-stride axis loads through ldmatrix instead, with the same result. Packed element types return a PackedFragment.

matrix_copy_plan(ownership, T, axis) -> MatrixCopyPlan or nothing

Cover a 32-lane ownership of T elements with 8×8 ldmatrix/stmatrix blocks, given the logical axis that memory stores contiguously. Every lane must hold complete words of adjacent elements along that axis, every block must be held by the same two units of every lane, and those units must follow the atom's register pattern, plain or transposed. Otherwise there is no plan and scalar accesses remain the correct route.

Collectively multiply operands and return the updated immutable accumulator.

mma(atom::Tcgen05MMA, d::TmemTile, a, b, Val(K), accumulate::Bool)

Issue K ÷ k instructions accumulating a*b over logical K into the TMEM accumulator d, from one thread. a is a Tcgen05Operand or a TmemTile of the atom's A type in the accumulator convention; b is a Tcgen05Operand. With accumulate=false the first instruction overwrites d. Nothing waits: order TMEM and shared-memory readiness with fences and barriers, and observe completion through commit_mma.

mma_async(plan, a, b, accumulator)

128-thread collective. Fence accumulator registers, issue the plan's WGMMA instructions and commit ONE group. Returns pending registers; only wait_mma accepts them. Operands must already be visible to the async proxy and remain unchanged until the wait completes. All four warps execute in convergence.

operand_layout(atom, role)

The thread/value ownership of an atom's A, B or accumulator operand: an explicit Layouts.Ownership following the PTX ISA figures. Evaluated while generating, so it is a constant inside kernels.

pack(T, fragment)
pack(fragment)

Convert logical values to T and pack adjacent local slots into register words. The one-argument form preserves the element type and all bits. Neither form moves data between threads or changes logical ownership. Supported element types are 8- and 16-bit: BFloat16, Float16, Float8E4M3, Float8E5M2, Int8 and UInt8. A fragment must hold complete words.

pack_operand_a(atom, left, right)
pack_operand_a(atom, accumulator, Val(m), Val(k))

Conversion from two horizontally adjacent FP32 16×8 accumulator fragments to one 16×16 A operand, without lane communication. Round to the atom's BF16/FP16 type, nearest even, and pack low element first. Signed zero and infinities survive; NaNs remain NaNs but their payload/sign are unspecified. No finite saturation or flush-to-zero modifier is applied. Finite overflow follows the destination format. This is numerical conversion, not an FP32 bit reinterpretation.

The tiled overload selects zero-based M repetition m and 16-column pair k. It requires one N warp and complete pairs of N atoms. Logical axis-1 ownership must match that of the consuming MMA; the surrounding kernel owns this correspondence.

partition(transfer::TmemTransfer, tile::TmemTile)

Bind a transfer to a logical storage view with matching shape and ownership. Use window to select a warp's region before partitioning. The transfer must cover one aligned 32-lane hardware band and contiguous storage positions. The executing warp must match that band's index within its warpgroup, and all 32 threads must use the same partition. Types do not prove participation.

Bounds checks may be omitted with @inbounds after establishing the storage contract; unsupported instruction widths/layouts are always rejected.

prepare_tma(plan, array; bounds=size(array))

Encode and upload a descriptor on the host. array is a dense device matrix stored as (inner, outer), irrespective of the plan's logical axis order, or a three-dimensional array of such matrices along its third axis, in which case origins take a third coordinate selecting the matrix and every matrix is bounded separately. It must have a 16-byte-aligned pointer and column stride. bounds names the logical extents when the array is padded: copies zero-fill loads and skip stores beyond them. The returned binding retains array and descriptor storage; keep it alive through all launches and graph replays. One binding serves loads and stores. Requires CUDACore and PTX. Source data may change; address and shape may not.

reduction_plan(ownership, axis) -> ReductionPlan or nothing

Derive the collective reduction of logical axis for an ownership, or nothing when no local-tree-plus-xor-shuffle recipe reaches every value sharing a kept coordinate with an identical instruction sequence on all threads.

reinterpret_tile(T, tile::TmemTile; dims)

View the same storage as FP32, BF16 or FP16, explicitly choosing the logical axis whose extent changes. This changes representation, not values or readiness. Currently requires a flat rectangular affine view with that axis traversing consecutive storage positions. 16-bit-to-FP32 views require complete aligned pairs.

relayout(target, fragment)

Reinterpret this thread's values under target ownership. Valid when every thread holds the same element set in both ownerships, so the change is a register permutation without communication. Other conversions are rejected.

relayout_permutation(from, to) -> Vector{Int} or nothing

For each slot of to, the slot of from holding the same coordinate in the same thread, when that permutation is identical on every thread.

Two ownerships assign the same coordinates to the same thread slots.

Multiply every register value by a scalar, preserving ownership.

shared_layout(plan)

The canonical storage a TMA tile occupies: rows of the swizzle width along the inner axis, packed, with the matching hardware swizzle, as Swizzle{B,M,3} composed with the row layout. Offsets are elements. The same encoding is what ldmatrix loads, wgmma_operand and tcgen05_operand describe (see swizzled_structure).

simplify_ownership(ownership)

The lane-local or warp-striped ownership with the same distribution, when one exists; otherwise the ownership itself.

Final log-sum-exp in the state's reduced ownership; empty rows return -Inf.

softmax_merge(left, right) -> (; state, left_rescale, right_rescale)

Combine independent summaries with identical reduced ownership. Combine their unnormalized weighted numerators using the returned factors. Floating-point merge order may change the result. Two empty summaries remain empty.

Normalize using one FP32 reciprocal per result and multiplication; empty rows return zero.

softmax_update(state, scores) -> (; state, weights, rescale)

Update running statistics with a score tile. weights are unnormalized exp(scores - new_maximum) in the original fragment distribution. Multiply an existing weighted numerator by rescale before adding this tile's weighted values. Statistics use FP32 weights, before any conversion for another MMA. Empty chunks produce zero weights; an empty previous state has zero rescale. Distributed reductions require all lanes to participate. The state retains the reduction axis selected by SoftmaxState(scores; dims).

store!(plan, destination, accumulator, thread)
store!(destination, fragment, thread)
store!(pointer, packed::PackedFragment)

Store an accumulator to a logical output view using its ownership. The tile form stores any fragment at its ownership's coordinates: scalar stores in general, stmatrix for a packed fragment whose ownership decomposes into CopyAtom blocks over a shared tile when compiling for sm_90 or later. The packed payload overload writes this thread's local words contiguously to a global typed BF16/FP16 pointer or a UInt16 bit pointer. It writes complete words, using 4-, 8- or 16-byte alignment for one, two/three or at least four words, and a sufficiently large distinct destination region for each thread.

store_async!(partition, fragment)

Issue a warp-collective store to TMEM. FP32 fragments store to FP32 partitions; packed BF16/FP16 fragments store to partitions of the same element type with matching ownership. Storage remains in use until wait_stores().

swizzled_structure(layout_type, T, kaxis) -> (; major, swizzle_bytes, leading_bytes, stride_bytes, row_elements, groups) or nothing

Recognize a static shared layout type as a canonical swizzled encoding of T elements (Swizzle{B,M,3} over packed rows of 16 << B bytes, M the chunk bits of 16 bytes) and read the descriptor fields for an operand whose logical K is axis kaxis: major is :K when K is the contiguous axis and :MN otherwise, swizzle_bytes the row width, leading_bytes the distance between row-width groups along the contiguous axis, stride_bytes the distance between eight-row core-matrix groups, row_elements one row and groups the number of rows placed side by side along the contiguous axis.

tcgen05_operand(atom, role, tile::SharedTile, origin=(Int32(0),Int32(0)))

Borrow an A (M,K) or B (K,N) operand from a canonical swizzled shared tile. The tile's static layout is recognized structurally (see swizzled_structure); the origin's coordinates are multiples of eight, and the non-K extent from the origin covers the atom's M or N. The root storage is aligned to eight swizzle rows (1024 bytes for 128-byte rows). Bounds checks may be elided with @inbounds.

tcgen05_operand(operand::Tcgen05Operand, origin)

The same tile at another logical origin, without recomputing the descriptor.

tma_load!(dst, binding, origin, barrier)

Issue one copy from a zero-based LOGICAL global origin (a pair of Int32 values, or a triple whose third coordinate selects the matrix of a three-dimensional binding) to the canonical shared tile. One elected thread executes. The caller initializes the 8-byte-aligned shared mbarrier, publishes its initialization to the async proxy, and supplies an expected transaction count including transfer_bytes(binding). Wait for that barrier phase before reading; release consumers before reusing storage. No implicit arrival, wait, thread rendezvous, or proxy fence is inserted. Out-of-bounds elements are zero-filled. The inner origin starts a 16-byte chunk (a multiple of 16 one-byte elements); the hardware rejects other origins.

tma_store!(binding, src, origin)

Issue one copy from the canonical shared tile to a zero-based LOGICAL global origin, from one elected thread. Origins are non-negative and the inner origin starts a 16-byte chunk; elements past the array's far edges are not written. The copy belongs to the issuing thread's current bulk group: commit_tma_stores closes the group, and wait_tma_reads or wait_tma_stores observes it. Shared data written through ordinary stores becomes visible to the copy only after fence.proxy.async.shared::cta and a rendezvous of the writing threads with the issuing thread; Tylo inserts neither.

tmem_allocate!(slot, Val(columns))
tmem_deallocate!(address, Val(columns))
tmem_relinquish_permit()

Warp-collective TMEM management. tmem_allocate! writes the base address of columns (a power of two from 32 to 512) to a shared slot; the CTA reads it after a barrier. tmem_deallocate! returns the columns and tmem_relinquish_permit lets other CTAs allocate. One warp executes each.

tmem_allocate!(slot, Val(columns))
tmem_deallocate!(address, Val(columns))
tmem_relinquish_permit()

Warp-collective TMEM management. tmem_allocate! writes the base address of columns (a power of two from 32 to 512) to a shared slot; the CTA reads it after a barrier. tmem_deallocate! returns the columns and tmem_relinquish_permit lets other CTAs allocate. One warp executes each.

tmem_location(tile, coordinate) -> (; address, bit_offset)

Resolve a logical coordinate to an ISA word address and subword bit offset. This is an address query, not a memory access. Coordinates are zero-based. Checks cover the view and hardware bounds, not the caller's allocation size.

tmem_allocate!(slot, Val(columns))
tmem_deallocate!(address, Val(columns))
tmem_relinquish_permit()

Warp-collective TMEM management. tmem_allocate! writes the base address of columns (a power of two from 32 to 512) to a shared slot; the CTA reads it after a barrier. tmem_deallocate! returns the columns and tmem_relinquish_permit lets other CTAs allocate. One warp executes each.

Expose the logical typed values of a packed fragment without numerical conversion.

validate_copy(ownership, T, layout)

Check on the host that the matrix copy derived for a 32-lane ownership over a shared layout addresses whole, contiguous, 16-byte aligned rows for every lane. Throws when no copy derives or a row is broken; returns nothing.

Check canonical operand geometry on the host; pointer alignment is checked at binding.

vector_plan(ownership, T, axis, align) -> (; width, groups) or nothing

Group each thread's slots into vectors of width bytes (16, 8 or 4) that are contiguous along logical axis and start at coordinates whose byte offset is a multiple of the width, identically on every thread, given a tile alignment of align bytes. The widest such grouping wins; nothing when even the narrowest multi-element vector does not fit.

Wait until at most N committed cp.async groups remain for this thread; no CTA barrier.

wait_load(pending)

Complete prior TMEM loads and return this load's typed registers. FP32 loads return Fragment; BF16/FP16 loads return PackedFragment. The wait also carries a compiler dependency through the returned registers. It waits for all prior loads of the executing threads, not just this handle.

Wait for ALL this warpgroup's committed WGMMA groups; return ready registers.

Complete all prior TMEM stores of the executing threads; warp collective.

wait_tma_reads(Val(N))

Wait until at most N of the issuing thread's committed bulk groups are still reading shared memory; the source tiles of the others may be reused.

wait_tma_stores(Val(N))

Wait until at most N of the issuing thread's committed bulk groups are still in flight; the global writes of the others are complete and visible to the issuing thread.

wgmma_operand(plan, role, tile, origin=(Int32(0),Int32(0)))

Borrow an operand from canonical K-major swizzled shared storage (see swizzled_structure). Origins are LOGICAL; A is (M,K), B is (K,N). Non-K origins must be multiples of eight, K origins multiples of 16; the plan's K fits one swizzle row and the whole plan fits the tile. Root storage is aligned to eight swizzle rows. Bounds/alignment checks may be elided with @inbounds once proven.

window(fragment, Val(origin), Val(shape))

Select a static logical register window, preserving all participating threads. The slots inside the window must be the same on every thread. Packed windows must retain complete word pairs. Runtime register indexing and implicit redistribution are not provided.

window_plan(ownership, origin, shape) -> (; slots, ownership) or nothing

The slots (1-based) falling inside a logical window, identical on every thread, and the window's own ownership with coordinates relative to origin.

@rtuple(f, range)
@rtuple(range) do i
    ...
end

Internal tuple construction for small, statically known integer ranges. Call f once per index, in range order, with inlining requested at each call site. Indices keep their integer type; empty ranges return ().

The callable and range expressions are evaluated once in the caller's scope. Range values become type parameters: literals and type-derived bounds work, but wrapping an unknown runtime range does not make it known to inference. This helper does not guarantee register residency.