Check the arch enum by criterion, raise DSLRuntimeError if the arch enum does not satisfy the criterion

Combined consumer release + advance with automatic elect_one using internal state.

Wait for a consumer stage and return its stage token/index. When `token` is provided, reuse the preceding `consumer_try_wait()` resul

Create a GPU tensor and return its pointer, torch tensor, cute tensor, CPU tensor, and strides.

Create GPU tensor from either a new or existing CPU tensor. :param torch_tensor_cpu: Optional existing CPU tensor to reuse. If None, creates a ne

Returns the most preferable alignment for a given shape and layout :param shape: extent of each dimension of the tensor :typ

Free the barrier memory on the GPU. This function frees the device memory for the barrier. It must be called after the barri

Given the `obj`, recursively go through it to extract all contained MLIR types

Get B extents. fprop: B extents array contains [K,T,R,S,C]. Turn that into ((K), (C,S,R,T)) dgrad: B extents array contains [K,T,R,S,C]. Turn that int

Get all users of the current node

Get the initial work tile info.

Initializes data structures

Returns true if source is needed based on state of runtime arguments

Decorator to mark a function for JIT compilation for Host code.

Creates a tensor by composing an engine (iterator/pointer) with a layout. A tensor is defined as T = E ∘ L, where E is an engine (array, pointer,

Perform a Consumer Epilogue to release all buffers

Returns a name identifying the tile description :return: name identifying the tile description :rtype: int

Reshape self.layout to shape

Dynamically-generated function that constructs a GEMM operation and verifies it against multiple test cases.

Execute RMSNorm and optionally benchmark performance. :param M: Number of rows (batch size * sequence length) :type M: int :param N:

Sets the predicate mask, overriding value stored in predicate iterator

Stores a fragment to memory at the location pointed to by the iterator

After smem async fence, before TMA store commit. Aux stores usually performed here Upon exit, all TMA stores for this subtile must have been issued Be

Transform MMA layout from ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, ...rest) to ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), ...rest). This

The factory function generating cute::Tuple with input tuple :param input_tuple: the input tuple :type input_tuple: tuple :param dtyp

Create FP8 tensor. - correctness mode: randint {-1, 0, 1} via bf16 cast - perf mode: random valid fp8 bit patterns via uint8

Returns a SyncObject corresponding to an agent's PipelineOp.

Normalize a Tensor or sequence of Tensors to a list of Tensors. Helper function for operations with variadic operands.

Compute element-wise arc cosine of the input tensor. :param a: Input tensor :type a: Union[TensorSSA, Numeric] :param fastmath: Enable fa

Adds a pointer offset in units of Element

Adds a tile offset

Adds a tile offset

Keeps references to external variables (e.g., Torch tensors when taking a view) in the scope of the lifetime of the JitArguments obje

Start of subtile store iteration

Is a thread sync needed after begin(). Allows chaining async copies across multiple nodes

Compute the number of contiguous elements needed to store a tensor with the given size

Clears the predicate set efficiently

Streams the result to global memory

Compute element-wise cosine of the input tensor. :param a: Input tensor (in radians) :type a: Union[TensorSSA, Numeric] :param fastmath:

Create the CUTE-format scale factor tensor on CUDA based on the reference tensor.

Create input, weight, and output tensors for RMSNorm.

Gets a pointer to the device data imaginary part

Convert a CUTLASS dtype object to its clean string name. This is needed to format dtype objects into CLI arguments without full modul

add mask for small tiles in ELL

Clears the predicate set efficiently

End of subtile store iteration

Find an existing extern function or declare it if it doesn't exist. This method checks if a function with the given name already exists in th

Execute the given function for each supplemental tensor.

Returns a pointer

Maps a CUTLASS numeric type to a cuDNN cudnnDataType_t

Get the hierarchical coordinate corresponding to a linear index. This method maps from a linear index back to the logical coordinate

Gets the mask

Gets the mask

Obtains calling linear threadblock index of the first block to work on the given tile

Get the non-zero stride mnl. This is used in argument construction