When the work tile is not special for reduction, it's valid. Otherwise need to skip global loading that producer warpgroup do, also math computation t

Wait on a stage to be unlocked

Correctness helper (1 iteration, ref-checked).

Add an edge src -> dst to dag ir with weight

Adds a tile offset in the unit of tile.

Advance to the next work tile and return its info.

Allocate an array of elements in shared memory. :param element_type: The type of elements to allocate :type element_type: Union[Type[

Compare GPU results against CPU reference implementation. Evaluates the epilogue computation on CPU and validates that: - Output tens

Computes the grid shape

Get the original pipeline this resource handle belongs to.

Returns ImplicitGemm tensor C size in number of elements

Check if dtype is a float8 type that doesn't support dlpack. params dtype: The cutlass numeric type to check type dtype: Type[cutlass.Numeric]

Check if the value is statically known. :return: True if the value is not a dynamic expression :rtype: bool

Used for initialization

Overrides the internal iteration index

Store a fragment to memory

Returns the sum of diagonal elements

Returns a SyncObject corresponding to an agent's PipelineOp.

Compares whether two plans are equal :param other_plan: plan to compare against the default GEMM :type other_plan: cutlass_c

Compares whether two plans are equal :param other_plan: plan to compare against the default Conv2d :type other_plan: cutlass

Insert a new compiled device module

Adds a pointer offset in units of Element

Increment and return an instance to self.

Returns an unsorted list of alignments supported by this data type combination :param operand: identifier of operand in question (e.

Helper to invoke the output functor over each vector of output

Returns whether the block assigned this work should compute the epilogue for the corresponding output tile. For the basic tile scheduler, this is alwa

Create a conditional branch. Parameters ---------- cond : ir.Value The condition value (i1 type). true_bl

Creates an abstraction for a circular buffer of synchronization primitives indexed by stage count. Args: stage: Number of pipeli

Return the corresponding torch.dtype per the given DSL type

The DSL element type (e.g., Float32, Int32).

Get the string value of an environment variable. Note that the value is cached after the first call.

Obtains the threadblock offset (in units of threadblock-scoped tiles)

Loads a fragment from memory with additional logical offset

Initializes an array of mbarriers using warp 0.

Compute one k-tile of mma for Q*K^T. Updates the tiled MMA QK and pipeline states. :param qk_params: The qk parameters :type qk_param

prepare the MMA accumulators when initialization or zeroing is required.

Sorts each list of kernels in `kernels_by_alignment` in descending order of threadblock shape

Replace address for the global tensor (to be done by single thread)

Override ArithValue.with_signedness for keyword-only __init__.

Helper method to convert an argument to its corresponding MLIR types. This method converts numeric meta types and types convertible via `get

Adds a pointer offset in units of Element