tilelang.transform

Wrapping transformations.

Submodules

• tilelang.transform.pass_config
• tilelang.transform.simplify

Functions

get_pass_context()	Get the current pass context
ClusterPlanning()	ClusterPlanning
PipelinePlanning()	infer the fragment/shared memory layout
LayoutInference()	LayoutInference
LowerTileOp()	LowerTileOp
InjectSoftwarePipeline()	InjectSoftwarePipeline
FrontendLegalize()	FrontendLegalize
LowerHopperIntrin()	LowerHopperIntrin
WarpSpecializedPipeline()	WarpSpecializedPipeline
RewriteWgmmaSync()	RewriteWgmmaSync
ThreadSync(storage_scope)	Insert sync between parallel read/write of shared buffers.
ThreadPartialSync(storage_scope)	Insert partial sync.
IfStmtBinding()	IfStmtBinding
MergeIfStmt()	MergeIfStmt
MultiVersionBuffer()	WarpSpecializedPipeline
WarpSpecialized()	WarpSpecializedPipeline
InjectTmaBarrier()	InjectTmaBarrier
InjectFenceProxy()	InjectFenceProxy
LegalizeVectorizedLoop()	LegalizeLoopVectorize
LegalizeSafeMemoryAccess()	LegalizeLoopVectorize
MakePackedAPI()	MakePackedAPI
AnnotateDeviceRegions()	AnnotateDeviceRegions
VectorizeLoop([enable_vectorize])	VectorizeLoop
InjectPTXAsyncCopy()	Rewrite global to shared memory copy on CUDA with asynchronous copy.
LowerDeviceStorageAccessInfo()	Lower attached storage access information on device.
LoopVectorizeDynamic()	Try to vectorize loop with dynamic shape.
ConfigIndexBitwidth()	Config index bitwidth.
FlattenBuffer()	FlattenBuffer
EliminateStorageSyncForMBarrier()	EliminateStorageSyncForMBarrier
MergeSharedMemoryAllocations([...])	MergeSharedMemoryAllocations
LowerL2Persistent()	LowerL2Persistent
PersistThreadblock()	PersistThreadblock
AlignDynamicSharedMemoryAllocations([align_bytes])	AlignDynamicSharedMemoryAllocations
LowerSharedBarrier()	LowerSharedBarrier
StorageRewrite()	StorageRewrite
LowerOpaqueBlock()	LowerOpaqueBlock
LowerThreadAllreduce()	LowerThreadAllreduce
LowerDeviceKernelLaunch()	LowerDeviceKernelLaunch

Package Contents

tilelang.transform.get_pass_context()

Get the current pass context

tilelang.transform.ClusterPlanning()

ClusterPlanning

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

FrontendLegalize

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.LowerHopperIntrin()

LowerHopperIntrin

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.WarpSpecializedPipeline()

WarpSpecializedPipeline

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.RewriteWgmmaSync()

RewriteWgmmaSync

Returns:

fpass – The result pass

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.ThreadPartialSync(storage_scope)

Insert partial sync.

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

WarpSpecializedPipeline

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

InjectTmaBarrier

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.InjectFenceProxy()

InjectFenceProxy

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

MakePackedAPI

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

VectorizeLoop

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

Parameters:

enable_vectorize (bool)

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

Try to vectorize loop with dynamic shape.

Returns:

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

• —-

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

EliminateStorageSyncForMBarrier

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

MergeSharedMemoryAllocations

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

Parameters:

• enable_aggressive_merge (bool)

• align_bytes (int)

tilelang.transform.LowerL2Persistent()

LowerL2Persistent

tilelang.transform.PersistThreadblock()

PersistThreadblock

tilelang.transform.AlignDynamicSharedMemoryAllocations(align_bytes=16)

AlignDynamicSharedMemoryAllocations

Parameters:

align_bytes (int) – The alignment bytes.

tilelang.transform.LowerSharedBarrier()

LowerSharedBarrier

tilelang.transform.StorageRewrite()

StorageRewrite

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

tilelang.transform.LowerOpaqueBlock()

LowerOpaqueBlock

tilelang.transform.LowerThreadAllreduce()

LowerThreadAllreduce

tilelang.transform.LowerDeviceKernelLaunch()

LowerDeviceKernelLaunch