|
tvm
|
Functions | |
| Pass | VerifyGPUCode (ffi::Map< ffi::String, PrimExpr > constraints) |
| Pass to verify GPU code constraints. | |
| Pass | VerifyVTCMLimit (ffi::Optional< Target > default_target=std::nullopt) |
| Pass to check if VTCM usage is within limit. | |
| Pass | OOBChecker () |
| Statically check TIR code for out of bounds array access. | |
| Pass | ConvertSSA () |
| De-duplicate definitions, including schedulable block iterators, across Functions. | |
| Pass | StmtSimplify () |
| Simplify schedulable TIR using block iteration constraints and shared simplifier options. | |
| Pass | CanonicalizeLoop () |
| Canonicalize loop to start from zero . | |
| Pass | LowerCrossThreadReduction () |
| Lower cross-thread reduction from thread bindings to intrinsic function calls. | |
| Pass | LowerInitBlock () |
| Lower block init stmt into If stmts. | |
| Pass | 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. | |
| Pass | ConvertBlocksToOpaque () |
| Substitute all the block vars with the PrimExprs they are bound to, indicated by the corresponding iter_values in BlockRealize, for opaque blocks by removing all . the iter_values in BlockRealize and iter_vars in Block. | |
| Pass | LiftThreadBinding () |
| Lift the same thread bindings to their LCA loops. | |
| Pass | CompactTensorAllocation (bool is_strict=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. | |
| Pass | LowerMatchTensor () |
| Remove match tensors inside the block. Also, it will validate the binding. | |
| Pass | InjectPermutedLayout () |
| Inject permuted layout for shared memory. | |
| Pass | TransformMmaTensorLayout () |
| Transform Mma scope (m16n8k8.matrixA/B/C) to local scope with layout transformation. | |
| Pass | 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. | |
| Pass | LowerThreadBinding () |
| Lower thread-binding loops to fresh lexical launch regions. Run after LoopPartition and LowerOpaqueBlock, while retaining loops for scheduling and feature extraction before this boundary. | |
| Pass | InjectSoftwarePipeline () |
| This pass transforms annotated loops into pipelined ones where producers and consumers are overlapped with the information provided in loop annotations, which enables optimization techniques like prefetching and pipeline parallelism. | |
| Pass | LowerAutoCopy () |
| Automatically do memory optimizations for auto copy blocks. | |
| Pass | ManifestSharedMemoryLocalStage () |
| Add the explicit local stage for the shared memory access on GPU. | |
| Pass | AnnotateIrregularLoop () |
| Annotate irregular loop mark. | |
| Pass | LoopPartition () |
| partition loops in the stmt. | |
| Pass | InjectVirtualThread () |
| Inject virtual thread loops. | |
| Pass | RewriteUnsafeSelect () |
| Detect and rewrite unsafe select that contains memory access. | |
| Pass | LowerVtcmAlloc () |
| Lower VTCM allocations. | |
| Pass | ThreadSync (tvm::ffi::String storage_scope) |
| Insert sync between parallel read/write of shared buffers. | |
| Pass | LowerSynchronization () |
| Lower asynchronous queue operations and remove synchronization regions. Run after ThreadSync and optional InjectPTXAsyncCopy. | |
| Pass | InferFragment () |
| Infer the TensorCore fragment information using tensor intrinsics. | |
| Pass | LowerThreadAllreduce () |
| Lower cross thread allreduce. | |
| Pass | InjectPTXAsyncCopy () |
| Rewrite global to shared memory copy on CUDA with asynchronous copy. | |
| Pass | MergeSharedMemoryAllocations () |
| Merge multiple TIR-level shared memory allocations into one. | |
| Pass | DefaultGPUSchedule () |
| Set default thread bindings for GPU Functions. | |
| Pass | RemoveWeightLayoutRewriteBlock (bool skip_tensor_rewrite=false) |
| Remove weight layout rewrite block before benchmark. | |
| Pass | RemoveStoreUndef () |
| Remove stores of tirx::undef_op. | |
| Pass tvm::s_tir::transform::AnnotateIrregularLoop | ( | ) |
Annotate irregular loop mark.
| Pass tvm::s_tir::transform::CanonicalizeLoop | ( | ) |
Canonicalize loop to start from zero .
| Pass tvm::s_tir::transform::CompactTensorAllocation | ( | bool | is_strict = 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.
Before narrowing, B is a [16, 16] tensor, but only a skinny vector B[i, 0:16] is accessed.
This pass narrows the tensor shape and adjust its accessed region accordingly. In this particular case, because only a 1 * 16 vector of B is accessed, the pass narrows B to shape [1, 16], and changes the access to B[i, j] to B[0, j].
| is_strict | ensure the compacted shape always smaller than the original shape. otherwise it allows to grow the shape to match actual accessed tensor regions. |
| Pass 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, for opaque blocks by removing all . the iter_values in BlockRealize and iter_vars in Block.
| Pass tvm::s_tir::transform::ConvertSSA | ( | ) |
De-duplicate definitions, including schedulable block iterators, across Functions.
| Pass tvm::s_tir::transform::DefaultGPUSchedule | ( | ) |
Set default thread bindings for GPU Functions.
| Pass tvm::s_tir::transform::InferFragment | ( | ) |
Infer the TensorCore fragment information using tensor intrinsics.
| Pass tvm::s_tir::transform::InjectPermutedLayout | ( | ) |
Inject permuted layout for shared memory.
| Pass tvm::s_tir::transform::InjectPTXAsyncCopy | ( | ) |
Rewrite global to shared memory copy on CUDA with asynchronous copy.
| Pass tvm::s_tir::transform::InjectSoftwarePipeline | ( | ) |
This pass transforms annotated loops into pipelined ones where producers and consumers are overlapped with the information provided in loop annotations, which enables optimization techniques like prefetching and pipeline parallelism.
The pipeline scope consists of the direct children of the annotated loop (ignoring SBlockRealize, SBlock, SeqStmt), and the number of children is denoted by n in the documentation.
The following annotations are used to guide the loop transformation:
1) Loop annotation software_pipeline_stage defines the pipeline stage. An array of n integers, and each element should be in range [0, max_stage], where max_stage is the maximum (inclusive) stage. 2) Loop annotation software_pipeline_order defines the pipeline order. An array of n integers, a permutation of [0, 1, ..., num_components - 1]; 3) SBlock annotation double_buffer_scope controls certain tensor sizes to allow decoupling of read/write dependency. It's an integer index of the write regions of the block.
Every annotated loop is transformed into a loop with three blocks as its direct children:
1) Prologue block, where components whose stage is less than max_stage is executed;
2) Body block, where all the components are executed;
3) Epilogue block, where only components whose stage is greater than 0 will be executed. The execution order is controlled by the annotation software_pipeline_order, and thus could be different than the original order.
Note: For nested software pipelines, the inner software pipeline will be generated first, which may affect the number of the direct children of the outer loop. In this case, the annotations for the outer software pipeline should include the result of the inner software pipeline, which is the three blocks as discussed above.
| Pass tvm::s_tir::transform::InjectVirtualThread | ( | ) |
Inject virtual thread loops.
| Pass tvm::s_tir::transform::LiftThreadBinding | ( | ) |
Lift the same thread bindings to their LCA loops.
| Pass tvm::s_tir::transform::LoopPartition | ( | ) |
partition loops in the stmt.
| Pass tvm::s_tir::transform::LowerAutoCopy | ( | ) |
Automatically do memory optimizations for auto copy blocks.
| Pass tvm::s_tir::transform::LowerCrossThreadReduction | ( | ) |
Lower cross-thread reduction from thread bindings to intrinsic function calls.
| Pass tvm::s_tir::transform::LowerInitBlock | ( | ) |
Lower block init stmt into If stmts.
| Pass tvm::s_tir::transform::LowerMatchTensor | ( | ) |
Remove match tensors inside the block. Also, it will validate the binding.
| Pass 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.
| Pass tvm::s_tir::transform::LowerSynchronization | ( | ) |
Lower asynchronous queue operations and remove synchronization regions. Run after ThreadSync and optional InjectPTXAsyncCopy.
| Pass tvm::s_tir::transform::LowerThreadAllreduce | ( | ) |
Lower cross thread allreduce.
| Pass tvm::s_tir::transform::LowerThreadBinding | ( | ) |
Lower thread-binding loops to fresh lexical launch regions. Run after LoopPartition and LowerOpaqueBlock, while retaining loops for scheduling and feature extraction before this boundary.
| Pass tvm::s_tir::transform::LowerVtcmAlloc | ( | ) |
Lower VTCM allocations.
| Pass tvm::s_tir::transform::ManifestSharedMemoryLocalStage | ( | ) |
Add the explicit local stage for the shared memory access on GPU.
| Pass tvm::s_tir::transform::MergeSharedMemoryAllocations | ( | ) |
Merge multiple TIR-level shared memory allocations into one.
| Pass tvm::s_tir::transform::OOBChecker | ( | ) |
Statically check TIR code for out of bounds array access.
| Pass 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.
| Pass tvm::s_tir::transform::RemoveStoreUndef | ( | ) |
Remove stores of tirx::undef_op.
| Pass tvm::s_tir::transform::RemoveWeightLayoutRewriteBlock | ( | bool | skip_tensor_rewrite = false | ) |
Remove weight layout rewrite block before benchmark.
| skip_tensor_rewrite | Whether to skip tensor rewrite. |
| Pass tvm::s_tir::transform::RewriteUnsafeSelect | ( | ) |
Detect and rewrite unsafe select that contains memory access.
| Pass tvm::s_tir::transform::StmtSimplify | ( | ) |
Simplify schedulable TIR using block iteration constraints and shared simplifier options.
| Pass tvm::s_tir::transform::ThreadSync | ( | tvm::ffi::String | storage_scope | ) |
Insert sync between parallel read/write of shared buffers.
| storage_scope | The storage scope considered. |
| Pass tvm::s_tir::transform::TransformMmaTensorLayout | ( | ) |
Transform Mma scope (m16n8k8.matrixA/B/C) to local scope with layout transformation.
Pass to verify GPU code constraints.
| constraints | The dict to specify constraints. |