tvm.s_tir.transform#
Namespace of all S-TIR transformations
- tvm.s_tir.transform.AnnotateIrregularLoop()
Annotate irregular loop mark.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.CanonicalizeLoop()
Canonicalize the loop to start from zero and use trivial step
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.CompactTensorAllocation(is_strict: bool = True)
Compact the tensor access region by removing the tensor regions that are not accessed, i.e. narrowing the tensor shape and adjust the access region if necessary.
- Parameters:
is_strict (bool) – Ensure the compacted shape to be always smaller than the original shape. Otherwise it allows to grow the shape to match actual accessed tensor regions.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.ConvertBlocksToOpaque()
Substitute all the block vars with the PrimExprs they are bound to, indicated by the corresponding iter_values in BlockRealize, and then convert the blocks into opaque ones by removing all the iter_values in BlockRealize and iter_vars in Block.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.ConvertSSA()
De-duplicate definitions, including schedulable block iterators, across Functions.
- tvm.s_tir.transform.DefaultGPUSchedule()
Set default thread bindings for GPU Functions.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.InferFragment()
Infer the TensorCore fragment information using tensor intrinsics.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.InjectPTXAsyncCopy()
Rewrite global to shared memory copy on CUDA with asynchronous copy.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.InjectPermutedLayout()
Inject permuted layout in mma
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.InjectSoftwarePipeline()
Transform annotated loops into pipelined one that parallelize producers and consumers
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.InjectVirtualThread()
Inject virtual thread loops.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.LiftThreadBinding()
Lift the same thread bindings to their LCA loops.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.LoopPartition()
Partition loops in the stmt.
Consumes the owning For’s
loop_partition_hintannotation. A provably true value enables partitioning even without likely tags. Run this pass before opaque lowering when hinting thread-binding loops.- Returns:
fpass – The result pass
- Return type:
- class tvm.s_tir.transform.LoopPartitionConfig(partition_const_loop=<MISSING>, no_unroll_loop_with_extent_one=<MISSING>, unroll_loop_with_partition_hint_no_interval=<MISSING>)
Config for loop partition pass
- property no_unroll_loop_with_extent_one
Don’t unroll loops with extent 1
- property partition_const_loop
Split constant loop
- property unroll_loop_with_partition_hint_no_interval
Unroll loops with loop_partition_hint and no interval
- tvm.s_tir.transform.LowerAutoCopy()
Automatically do memory optimizations for auto copy blocks
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.LowerCrossThreadReduction()
Lower cross-thread reduction from thread bindings to intrinsic function calls.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.LowerInitBlock()
Lower block init stmt into If statements.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.LowerMatchTensor()
Remove match tensors inside the block. Also, it will validate the binding.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.LowerOpaqueBlock()
Remove the block to ensure that the TIR can not be scheduled again.
Run LoopPartition first when thread-binding loops carry loop_partition_hint.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.LowerSynchronization()
Lower async queue operations and erase regions after synchronization planning.
- tvm.s_tir.transform.LowerThreadBinding()
Lower thread-binding loops to fresh lexical launch regions.
Run after LoopPartition and LowerOpaqueBlock. Each loop keeps its own lexical binding, including loops that use the same hardware axis.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.LowerVtcmAlloc()
Lower vtcm allocation.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.ManifestSharedMemoryLocalStage()
Add the explicit local stage for the shared memory access on GPU.
- Returns:
fpass – The result pass
- Return type:
- class tvm.s_tir.transform.MemoryScope(*args: Any, **kwargs: Any)
A registered memory-space identity.
MemoryScope(name)looks up an existing scope and raisesValueErrorfor an unknown name. Registration alone does not enable backend support.Examples
custom = MemoryScope.register("my_accelerator.scratch") assert MemoryScope("my_accelerator.scratch") == custom
- property name: str
Return the canonical registered spelling.
- classmethod register(name: str) MemoryScope
Register a new canonical name, raising ValueError for a duplicate.
- tvm.s_tir.transform.MergeSharedMemoryAllocations()
This pass merges multiple TIR-level shared memory allocations into one allocation.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.PlanAndUpdateTensorAllocationLocation()
Locate the tensor allocation to the exact position (usually is the lca of tensor access). This pass will inject opaque block with alloc_tensors at the allocation site.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.RemoveStoreUndef()
Remove stores of undefined values from the Stmt.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.RemoveWeightLayoutRewriteBlock(skip_tensor_rewrite=False)
Remove weight layout rewrite block before benchmarking during tuning stage.
- Parameters:
skip_tensor_rewrite (bool) – If True, exact rewrite of Tensor, according to the given index map, will be skipped.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.RewriteUnsafeSelect()
Detect and rewrite unsafe select that contains memory access.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.StmtSimplify()
Simplify schedulable TIR with block constraints.
- tvm.s_tir.transform.ThreadSync(storage_scope: str | MemoryScope)
Insert sync between parallel read/write of shared tensors.
- Parameters:
storage_scope (str or MemoryScope) – The target storage scope.
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.TransformMmaTensorLayout()
Transform mma tensor layout
- Returns:
fpass – The result pass
- Return type:
- tvm.s_tir.transform.VerifyVTCMLimit(default_target=None)
Verify if the size of the allocated vtcm memory satisfies the limit.
The limit is determined from the “vtcm-capacity” attribute of the target.
- Parameters:
default_target (Optional[tvm_tirx.target.Target]) – The default target to use if a Function does not have a target attribute.
- Returns:
fpass – The result pass
- Return type: