tilelang.transform package

Submodules

• tilelang.transform.pass_config module
  • PassConfigKey
    • PassConfigKey.CUDA_KERNELS_OUTPUT_DIR
    • PassConfigKey.TIR_ADD_LOWER_PASS
    • PassConfigKey.TIR_DISABLE_CSE
    • PassConfigKey.TIR_DISABLE_STORAGE_REWRITE
    • PassConfigKey.TIR_DISABLE_VECTORIZE
    • PassConfigKey.TIR_ENABLE_DEBUG
    • PassConfigKey.TIR_ENABLE_EQUIV_TERMS_IN_CSE
    • PassConfigKey.TIR_MERGE_STATIC_SMEM
    • PassConfigKey.TIR_NOALIAS
    • PassConfigKey.TIR_SIMPLIFY
    • PassConfigKey.TIR_USE_ASYNC_COPY
    • PassConfigKey.TL_CONFIG_INDEX_BITWIDTH
    • PassConfigKey.TL_DEBUG_MERGE_SHARED_MEMORY_ALLOCATIONS
    • PassConfigKey.TL_DISABLE_DYNAMIC_TAIL_SPLIT
    • PassConfigKey.TL_DISABLE_SAFE_MEMORY_ACCESS
    • PassConfigKey.TL_DISABLE_TMA_LOWER
    • PassConfigKey.TL_DISABLE_WARP_SPECIALIZED
    • PassConfigKey.TL_DYNAMIC_ALIGNMENT
    • PassConfigKey.TL_SIMPLIFY
• tilelang.transform.simplify module
  • Simplify()
  • apply_simplify()
  • simplify_prim_func()

Module contents

Wrapping transformations.

tilelang.transform.AnnotateDeviceRegions()

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.ClusterPlanning()

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.ConfigIndexBitwidth()

Config index bitwidth.

Returns:

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

• —-

tilelang.transform.EliminateStorageSyncForMBarrier()

tilelang.transform.FlattenBuffer()

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.FrontendLegalize()

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.IfStmtBinding()

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.InjectFenceProxy()

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.InjectPTXAsyncCopy()

Rewrite global to shared memory copy on CUDA with asynchronous copy.

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.InjectSoftwarePipeline()

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.InjectTmaBarrier()

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.LayoutInference()

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.LegalizeVectorizedLoop()

LegalizeLoopVectorize

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.LoopVectorizeDynamic()

Try to vectorize loop with dynamic shape.

Returns:

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

• —-

tilelang.transform.LowerDeviceStorageAccessInfo()

Lower attached storage access information on device.

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

Note

Run this pass after all storage access analysis finish.

tilelang.transform.LowerHopperIntrin()

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.LowerL2Persistent()

tilelang.transform.LowerTileOp()

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.MakePackedAPI()

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.MergeIfStmt()

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.MergeSharedMemoryAllocations()

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.MultiVersionBuffer()

WarpSpecializedPipeline

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.PipelinePlanning()

infer the fragment/shared memory layout

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.RewriteWgmmaSync()

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.ThreadPartialSync(storage_scope: str)

Insert partial sync.

Parameters:

storage_scope (str) – The target storage scope.

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.ThreadSync(storage_scope: str)

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.VectorizeLoop(enable_vectorize: bool = True)

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.WarpSpecialized()

WarpSpecializedPipeline

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.WarpSpecializedPipeline()

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.get_pass_context()

Get the current pass context