Print a formatted message using the target’s debug print facilities. The format string is target
specific, but Vulkan and CUDA both use the C++ conventions. WGSL isn’t currently supported.
Print a formatted message using the target’s debug print facilities. The format string is target
specific, but Vulkan and CUDA both use the C++ conventions. WGSL isn’t currently supported.
GPU version of asm. Currently parses all the same options, but most are not
applicable to GPU assembly architectures. Should validate and give proper errors at some point.
Also adds a new register spec: the inferred register specifier (_). This is because the
specifier isn’t actually meaningful in PTX and is currently ignored, with constraints being
inferred from the value type. The reason it’s still present is because we may want to add an
explicit mem specifier to allow fine-grained memory clobbering, or support for other assembly
formats that do use different register types.
An atomic numerical type wrapping a normal numeric primitive. Enables the use of atomic
operations, while disabling normal operations. In WGSL, this is a separate type - on CUDA/SPIR-V
it can theoretically be bitcast to a normal number, but this isn’t recommended.
Im2col indexing. Loads a “column” (not the same column as im2col) of pixels into shared
memory, with a certain offset (kernel position). The corners are the bounds to load pixels
from at offset 0, so the top left corner of the kernel. The offset is added to the
corner offsets, so a (-1, -1) corner will stop the bounding box at (1, 1) for kernel
offset (2, 2).
The scope represents the region currently being parsed, as well as the current insertion point.
It is created from scratch for the initial codegen phase, or rebuilt from an existing scope and
a rewriter for passes/conversions.
Hacky solution for getting comptime properties into the scope.
Allows querying certain target-specific properties at compile time, rather than at runtime.
Review on how to better solve this and delegate to the compiler if possible.
A CUDA CUtensorMap object. Represents a tensor encoded with a lot of metadata, and is an
opaque packed object at runtime. Does not support retrieving any shapes or strides, nor does
it give access to the pointer. So these need to be passed separately in an aliased Tensor if needed.
Grid constant tensor map, currently only maps to CUDA tensormap. May be interleaved or swizzled,
but last dimension must be contiguous (since strides don’t include the last dimension).
WhileOp builder. Needs to work like this for borrow checker reasons.
Rust requires all closures passed to a function to be live at once, so we use a builder pattern
and separate the cond and body into separate functions to allow the borrows to work.
The type used for addressing storage types in a kernel.
This is the type usize maps to when used in a kernel, with isize being mapped to the signed
equivalent.
What a compiled kernel does with one buffer binding, stamped on the entry
function’s arguments by the visibility analysis and carried to the launch
site, where the write scope stages exactly what the kernel writes and
checks exactly what it reads.
Data are organized in a specific order in global memory; however, this may not match the order
in which the application accesses data in shared memory. This difference in data organization
may cause bank conflicts when shared memory is accessed. In order to avoid this problem, data
can be loaded to shared memory with shuffling across shared memory banks. When interleave is
TensorMapInterleave::B32, swizzle must be TensorMapSwizzle::B32.
Other interleave modes can have any swizzling pattern.
A type that can be used as a kernel comptime argument.
Note that a type doesn’t need to implement CubeComptime to be used as
a comptime argument. However, this facilitate the declaration of generic cube types.
Read a value into registers. Should only be implemented for types that can exist outside of
memory like primitives or Array, but not for things like Barrier that must exist behind a
pointer.
Marker trait for scalar primitives. Should be implemented for all scalar CubePrimitives, but
not for Vector or non-standard primitives like Barrier. Alternatively, treat these as
types that can be stored in a [Vector]
Perform a ballot operation across all units in a plane.
Returns a set of 32-bit bitfields as a Vector, with each element containing the value from 32
invocations.
Note that vector size will always be set to 4 even for PLANE_DIM <= 64, because we can’t
retrieve the actual plane size at expand time. Use the runtimePLANE_DIM to index appropriately.
Broadcasts the value from the specified plane unit at the given index
to all active units within that plane. Requires a constant index. For non-constant indices,
use plane_shuffle().
Perform an exclusive product operation across all units in a plane.
This multiplies all values to the “left” of the unit, excluding this unit’s value. The 0th unit
will be set to E::one().
Also known as “exclusive prefix product” or “exclusive scan”.
Perform an exclusive sum operation across all units in a plane.
This sums all values to the “left” of the unit, excluding this unit’s value. The 0th unit will
be set to E::zero().
Also known as “exclusive prefix sum” or “exclusive scan”.
Perform an inclusive product operation across all units in a plane.
This multiplies all values to the “left” of the unit, including this unit’s value.
Also known as “prefix product” or “inclusive scan”.
Perform an inclusive sum operation across all units in a plane.
This sums all values to the “left” of the unit, including this unit’s value.
Also known as “prefix sum” or “inclusive scan”.
Perform a shuffle down operation across the plane.
Each unit reads the value from a unit with a higher lane ID (current_id + delta).
A unit whose lane_id + delta >= plane_dim reads an indeterminate value that differs
by backend, so mask those units at the call site.
Perform a shuffle up operation across the plane.
Each unit reads the value from a unit with a lower lane ID (current_id - delta).
A unit with lane_id < delta reads an indeterminate value that differs by backend, so
mask those units at the call site.
Perform a shuffle XOR operation across the plane.
Each unit exchanges its value with another unit at an index determined by XOR with the mask.
This is useful for butterfly reduction patterns.
sync_async_proxy_shared is a synchronization fence for the experimental SM 9.0+ copy
functions, applying bidirectionally between the async proxy (i.e. TMA) and shared memory.
Should be used after initializing the barriers, and before the copy operation.
PTX: fence.proxy.async.shared::cta
Experimental and subject to change.
Copy a tile from a shared memory src to a global memory dst, with the provided
offsets. Should be combined with memcpy_async_tensor_commit and
memcpy_async_tensor_wait_read.
Copy a tile from a shared memory src to a global memory dst, with the provided
offsets. Should be combined with memcpy_async_tensor_commit and
memcpy_async_tensor_wait_read.
Copy a tile from a shared memory src to a global memory dst, with the provided
offsets. Should be combined with memcpy_async_tensor_commit and
memcpy_async_tensor_wait_read.
Copy a tile from a shared memory src to a global memory dst, with the provided
offsets. Should be combined with memcpy_async_tensor_commit and
memcpy_async_tensor_wait_read.
Copy a tile from a shared memory src to a global memory dst, with the provided
offsets. Should be combined with memcpy_async_tensor_commit and
memcpy_async_tensor_wait_read.
Used for let x; local bindings. These are technically immutable in Rust, but because they can
be visible at a higher level scope than their assignment we need to allocate a slot of them at
that level and treat them as mutable.
Barrier, then load reference with the result marked workgroup-uniform —
mirrors WGSL’s workgroupUniformLoad. Lets a workgroup-shared value gate
control flow that contains barriers. Non-WGSL backends lower it to
[sync_cube] plus a plain load.