tilelang.transform

Wrapping transformations.

Submodules

• tilelang.transform.add_bufstore_wrapper
• tilelang.transform.decouple_type_cast
• tilelang.transform.hoist_broadcast_values
• tilelang.transform.pass_config
• tilelang.transform.simplify

Attributes

DEFAULT_SIMT_THREADS

Functions

get_pass_context()	Get the current pass context
PipelinePlanning()	infer the fragment/shared memory layout
LayoutInference()	LayoutInference
LowerTileOp()	LowerTileOp
InjectSoftwarePipeline()	InjectSoftwarePipeline
LegalizeNegativeIndex()	Legalize negative indices in buffer loads.
InjectAssumes()	Inject Assumes for natural shape boundary conditions. And convert Assumes in Evaluate(Call(...)) form
VerifyParallelLoop()	VerifyParallelLoop
VerifyBufferInit()	Warn when a non-global-scope buffer is read before anything writes it.
ThreadSync(storage_scope)	Insert sync between parallel read/write of shared buffers.
IfStmtBinding()	IfStmtBinding
MergeIfStmt()	MergeIfStmt
LoopUnswitching()	LoopUnswitching: Hoist loop-invariant if statements out of loops.
LegalizeVectorizedLoop()	LegalizeLoopVectorize
LegalizeSafeMemoryAccess()	LegalizeLoopVectorize
LowerAccessPtr()	Lower TileLang frontend tl.access_ptr to tir.builtin.tvm_access_ptr.
MakePackedAPI()	MakePackedAPI
MaterializeKernelLaunch([lower_thread_binding, ...])	Materialize the target-neutral kernel launch nest emitted by T.Kernel
AnnotateDeviceRegions()	AnnotateDeviceRegions
SplitHostDevice()	Split host/device functions even for empty kernels.
AnnotateReadOnlyParams()	Annotate read-only handle parameters for PrimFuncs.
VectorizeLoop([enable_vectorize])	VectorizeLoop
ConfigIndexBitwidth()	Config index bitwidth.
FlattenBuffer()	FlattenBuffer
MergeSharedMemoryAllocations([...])	MergeSharedMemoryAllocations
PlanAndUpdateBufferAllocationLocation()	Plan and update buffer allocation locations within PrimFuncs.
HoistGlobalBufferAllocations()	Hoist global buffer allocations to the top of the block (host side).
HoistNonRestrictParams()
StorageRewrite()	StorageRewrite
LowerOpaqueBlock()	LowerOpaqueBlock
LowerThreadAllreduce()	LowerThreadAllreduce
LowerIntrin()	LowerIntrin
LowerDeviceKernelLaunch()	Create and return a transform pass that lowers device kernel launch constructs to target-specific IR.
CanonicalizeLegacyReducer()	Rewrite legacy (v1) reducer syntax into first-class reducer v2 ops.
VerifyReducerEpoch()	Verify lifecycle and access rules of reducer v2 epochs.
ReducerPlanAndMaterialize()	Plan physical storage/communication for reducer v2 epochs.
VerifyReducerConsumed()	Assert no reducer v2 construct survives past materialization.
UnrollLoop()	Unroll loops as in Halide pipeline.

Package Contents

tilelang.transform.get_pass_context()

Get the current pass context

tilelang.transform.PipelinePlanning()

infer the fragment/shared memory layout

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.LayoutInference()

LayoutInference

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.LowerTileOp()

LowerTileOp

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.InjectSoftwarePipeline()

InjectSoftwarePipeline

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.LegalizeNegativeIndex()

Legalize negative indices in buffer loads.

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.InjectAssumes()

Inject Assumes for natural shape boundary conditions. And convert Assumes in Evaluate(Call(…)) form (tvm builtin assume call) to AttrNode form.

Returns:

fpasstvm.transform.Pass

The result pass

tilelang.transform.VerifyParallelLoop()

VerifyParallelLoop

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.VerifyBufferInit()

Warn when a non-global-scope buffer is read before anything writes it.

Returns:

fpass – The registered pass. It inspects the IR and returns it unchanged.

Return type:

tvm.transform.Pass

tilelang.transform.ThreadSync(storage_scope)

Insert sync between parallel read/write of shared buffers.

Parameters:

storage_scope (str) – The target storage scope.

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.IfStmtBinding()

IfStmtBinding

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.MergeIfStmt()

MergeIfStmt

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.LoopUnswitching()

LoopUnswitching: Hoist loop-invariant if statements out of loops.

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.LegalizeVectorizedLoop()

LegalizeLoopVectorize

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.LegalizeSafeMemoryAccess()

LegalizeLoopVectorize

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.LowerAccessPtr()

Lower TileLang frontend tl.access_ptr to tir.builtin.tvm_access_ptr.

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.MakePackedAPI()

MakePackedAPI

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.DEFAULT_SIMT_THREADS = 128

tilelang.transform.MaterializeKernelLaunch(lower_thread_binding=True, default_threads=DEFAULT_SIMT_THREADS, unsupported_annotations=None)

Materialize the target-neutral kernel launch nest emitted by T.Kernel into a backend-specific form. Each backend pipeline decides the mode for itself; this is where the target-dependent parts of a launch (whether threads exist and how many run by default) are decided.

Parameters:

• lower_thread_binding (bool) – If True (SIMT backends, e.g. CUDA/ROCm/Metal), lower the blockIdx.* grid loops into thread_extent AttrStmts and bind the thread placeholders as threadIdx.* thread_extent scopes. If False (backends without SIMT, e.g. CPU), lower blockIdx.* loops into plain serial For loops and drop the thread placeholders. A body that references a thread index is rejected on such targets.

• default_threads (int | list[int] | tuple | None) – Thread-block extents used by SIMT backends when T.Kernel was called without threads=. Ignored when lower_thread_binding is False. None means the backend has no default and threads= is required.

• unsupported_annotations (list[str] | None) – Launch annotations (keys on the tilelang_root block, e.g. cluster_dims) that have no meaning on this backend. A launch carrying one is rejected here instead of being silently ignored by later passes.

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.AnnotateDeviceRegions()

AnnotateDeviceRegions

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.SplitHostDevice()

Split host/device functions even for empty kernels.

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.AnnotateReadOnlyParams()

Annotate read-only handle parameters for PrimFuncs.

Adds attribute tl.readonly_param_indices listing param indices that are never written, enabling CUDA codegen to emit const qualifiers to unlock read-only cache loads.

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.VectorizeLoop(enable_vectorize=True)

VectorizeLoop

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

Parameters:

enable_vectorize (bool)

tilelang.transform.ConfigIndexBitwidth()

Config index bitwidth.

Returns:

• fpass (tvm.transform.Pass) – The result pass

• —-

tilelang.transform.FlattenBuffer()

FlattenBuffer

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.MergeSharedMemoryAllocations(enable_aggressive_merge=False, align_bytes=16, disable_reuse=False)

MergeSharedMemoryAllocations

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

Parameters:

• enable_aggressive_merge (bool)

• align_bytes (int)

• disable_reuse (bool)

tilelang.transform.PlanAndUpdateBufferAllocationLocation()

Plan and update buffer allocation locations within PrimFuncs.

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.HoistGlobalBufferAllocations()

Hoist global buffer allocations to the top of the block (host side).

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.HoistNonRestrictParams()

tilelang.transform.StorageRewrite()

StorageRewrite

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.LowerOpaqueBlock()

LowerOpaqueBlock

tilelang.transform.LowerThreadAllreduce()

LowerThreadAllreduce

tilelang.transform.LowerIntrin()

LowerIntrin

tilelang.transform.LowerDeviceKernelLaunch()

Create and return a transform pass that lowers device kernel launch constructs to target-specific IR.

This pass transforms high-level device kernel launch and related intrinsics into lower-level IR suitable for backend code generation and device-side lowering.

Returns:

The transform pass that performs device kernel launch lowering.

Return type:

tvm.transform.Pass

tilelang.transform.CanonicalizeLegacyReducer()

Rewrite legacy (v1) reducer syntax into first-class reducer v2 ops.

Deprecation shim: T.clear + read-modify-write stores + in-place T.finalize_reducer(acc) become reducer_init/reducer_update/ out-of-place finalize with a fresh destination fragment. Unrecognized access patterns are compile errors, never silently accepted.

Returns:

The canonicalization pass.

Return type:

tvm.transform.Pass

tilelang.transform.VerifyReducerEpoch()

Verify lifecycle and access rules of reducer v2 epochs.

Enforces that every T.alloc_reducer has exactly one T.reducer_init, updates only inside T.Parallel between init and finalize, exactly one out-of-place T.finalize_reducer(acc, dst), and no ordinary reads/writes/aliasing of the reducer handle.

Returns:

The verification pass.

Return type:

tvm.transform.Pass

tilelang.transform.ReducerPlanAndMaterialize()

Plan physical storage/communication for reducer v2 epochs.

Runs after LayoutInference (loop layouts are read-only inputs) and materializes the first-class reducer ops into ordinary fragment storage, guarded read-modify-write updates, and an explicit finalize plan.

Returns:

The planning/materialization pass.

Return type:

tvm.transform.Pass

tilelang.transform.VerifyReducerConsumed()

Assert no reducer v2 construct survives past materialization.

Returns:

The verification pass.

Return type:

tvm.transform.Pass

tilelang.transform.UnrollLoop()

Unroll loops as in Halide pipeline.

This pass unrolls loops based on configuration options including: - auto_max_step: Threshold of number of steps to be automatically unrolled - auto_max_depth: Maximum nested level of loops that can be automatically unrolled - auto_max_extent: Maximum extent of loop that will be unrolled - explicit_unroll: Whether to explicitly unroll instead of setting a pragma - unroll_local_access: Whether to always unroll local access

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass