tilelang.ascend.transform

Ascend-only TIR transform passes.

Submodules

• tilelang.ascend.transform.z3_scheduler

Functions

InsertNd2Nz()	Rewrite layout-driven UB ND->UB NZ and UB->L1 ND->NZ copies.
NormalizeControlFlowForSchedule()	Normalize control flow ahead of AutoSchedule.
RestoreWhileLoops()	Rewrite synthetic_while-tagged for loops back into while(true).
AnnotateMultiBufferEligible()	Annotate scheduled For loops with buffers that can be multi-buffered.
NormalizeConflictHints()	Normalize T.assume_no_conflict and T.assume_conflict markers.
EstimateLatency()	Annotate every Ascend AutoSchedule task with latency and II.
AscendRemoveNoOp()	Remove no-op statements using equality-safe store analysis.
LegalizeSimdMerging()	Lower MODE_MERGING SIMD assignments to explicit read-write calls.
UnrollLoopSkipVF()	Expand explicit unroll loops outside SIMD_VF / SIMT_VF blocks.
AutoSchedule()	Schedule materialized internal schedule units for downstream lowering.
AssignCore()	Annotate each scheduled task with its widest legal AIV/AIC mask.
PrepareMultiBuffer()	Choose multi-buffer clocks and add counter init/advance tasks.
ResolveCore()	Narrow core candidates from actual scalar and counter consumers.
InsertSync()	Insert synchronization into scheduled Ascend TIR.
MaterializeMultiBuffer()	Lower logical multi-buffer accesses to physical buffer versions.
LowerScheduledTIR()	Lower scheduled TIR into final AIV/AIC core bodies.
RewriteDualCopy()	Lower dual-copy operations using the enclosing vector-core sid.
NormalizeBufferVersion()	Lower staged manual/automatic versions to scoped internal attributes.
RewriteAscendBufferVersionLayout()	Align each physical version of a multi-buffered UB allocation to 32 bytes.
NormalizeAscendFractalStorage()	Normalize Cube storage and aliases before OOB padding and scheduling.
AscendSimdVFLowerParallel()	Lower T.Parallel inside SIMD_VF blocks to MicroAPI register-level vector ops.
AscendThreadSync(storage_scope)	Insert thread-storage synchronization independently within each SIMT_VF.
AscendLayoutInference()	Ascend fork of LayoutInference: VF regions are opaque to the worklist
AscendLowerTileOp()	Ascend fork of LowerTileOp: VF region scopes, buffer-version key remap,
AscendInsertOOBPadding()	Clamp DMA copy OOB tails and emit GM->L1 padding fills before AutoSchedule.
InferBufferAliases()	Infer which on-chip buffers in a manually scheduled kernel may reuse storage.
MergeUBAllocations([align_bytes, disable_reuse])	Merge Ascend on-chip allocations from a pairwise reuse contract.
MarkScalarDcacheBypass()	Rewrite scalar global buffer accesses to use dcache bypass intrinsics
RewriteFp4ToFp4x2()	Retype genuine fp4 (float4_e2m1fn, lanes=1) storage into the 1-byte
RewriteFlagToBuf()	Normalize each hard_event's synchronization into its 8-slot flag namespace.

Package Contents

tilelang.ascend.transform.InsertNd2Nz()

Rewrite layout-driven UB ND->UB NZ and UB->L1 ND->NZ copies.

tilelang.ascend.transform.NormalizeControlFlowForSchedule()

Normalize control flow ahead of AutoSchedule.

Rewrites each while cond: BODY into a bounded serial for loop (tagged synthetic_while) so AutoSchedule can pipeline the body, then extracts complex if conditions and buffer-dependent serial/unrolled loop bounds into temporary Bind variables. RestoreWhileLoops() undoes the while-rewrite after scheduling.

tilelang.ascend.transform.RestoreWhileLoops()

Rewrite synthetic_while-tagged for loops back into while(true).

tilelang.ascend.transform.AnnotateMultiBufferEligible()

Annotate scheduled For loops with buffers that can be multi-buffered.

Requires NormalizeControlFlowForSchedule() followed by MaterializeScheduleUnits(), so buffer-dependent control expressions are schedulable tasks and manual stages and flattened scheduling guards are available through the shared IRStructure codec.

tilelang.ascend.transform.NormalizeConflictHints()

Normalize T.assume_no_conflict and T.assume_conflict markers.

tilelang.ascend.transform.EstimateLatency()

Annotate every Ascend AutoSchedule task with latency and II.

tilelang.ascend.transform.AscendRemoveNoOp()

Remove no-op statements using equality-safe store analysis.

tilelang.ascend.transform.LegalizeSimdMerging()

Lower MODE_MERGING SIMD assignments to explicit read-write calls.

tilelang.ascend.transform.UnrollLoopSkipVF()

Expand explicit unroll loops outside SIMD_VF / SIMT_VF blocks.

This pre-AutoSchedule pass only materializes T.unroll(..., explicit=True) loops. Non-explicit unroll hints, serial loops, and all VF block interiors are preserved for their later lowering stages.

tilelang.ascend.transform.AutoSchedule()

Schedule materialized internal schedule units for downstream lowering.

tilelang.ascend.transform.AssignCore()

Annotate each scheduled task with its widest legal AIV/AIC mask.

tilelang.ascend.transform.PrepareMultiBuffer()

Choose multi-buffer clocks and add counter init/advance tasks.

tilelang.ascend.transform.ResolveCore()

Narrow core candidates from actual scalar and counter consumers.

tilelang.ascend.transform.InsertSync()

Insert synchronization into scheduled Ascend TIR.

tilelang.ascend.transform.MaterializeMultiBuffer()

Lower logical multi-buffer accesses to physical buffer versions.

tilelang.ascend.transform.LowerScheduledTIR()

Lower scheduled TIR into final AIV/AIC core bodies.

tilelang.ascend.transform.RewriteDualCopy()

Lower dual-copy operations using the enclosing vector-core sid.

tilelang.ascend.transform.NormalizeBufferVersion()

Lower staged manual/automatic versions to scoped internal attributes.

tilelang.ascend.transform.RewriteAscendBufferVersionLayout()

Align each physical version of a multi-buffered UB allocation to 32 bytes.

tilelang.ascend.transform.NormalizeAscendFractalStorage()

Normalize Cube storage and aliases before OOB padding and scheduling.

tilelang.ascend.transform.AscendSimdVFLowerParallel()

Lower T.Parallel inside SIMD_VF blocks to MicroAPI register-level vector ops.

tilelang.ascend.transform.AscendThreadSync(storage_scope)

Insert thread-storage synchronization independently within each SIMT_VF.

Parameters:

storage_scope (str)

tilelang.ascend.transform.AscendLayoutInference()

Ascend fork of LayoutInference: VF regions are opaque to the worklist and SIMT_VF bodies infer against the region’s own lane scope.

tilelang.ascend.transform.AscendLowerTileOp()

Ascend fork of LowerTileOp: VF region scopes, buffer-version key remap, and no CUDA async-copy post-processing.

tilelang.ascend.transform.AscendInsertOOBPadding()

Clamp DMA copy OOB tails and emit GM->L1 padding fills before AutoSchedule.

For every supported DMA copy this rewrites the copy to its in-bounds ranges; for a padded GM->L1 copy it also appends semantic T.fill operations with exact destination regions. Runs after InsertNd2Nz and before AutoSchedule so each fill is scheduled as its own MTE2 task and ordered against L1->L0 consumers. LowerTileOp later converts the fills to ascend_fill_l1, which codegen emits as asc_fill_l1.

tilelang.ascend.transform.InferBufferAliases()

Infer which on-chip buffers in a manually scheduled kernel may reuse storage.

The resulting pairwise compatibility contract is consumed by MergeUBAllocations().

tilelang.ascend.transform.MergeUBAllocations(align_bytes=16, disable_reuse=False)

Merge Ascend on-chip allocations from a pairwise reuse contract.

AutoSchedule writes this contract during synchronization insertion. Manual schedules must run InferBufferAliases() first.

Buffer reuse is always performed unless disable_reuse is set.

Returns:

fpass – The result pass

Return type:

tvm.transform.Pass

Parameters:

• align_bytes (int)

• disable_reuse (bool)

tilelang.ascend.transform.MarkScalarDcacheBypass()

Rewrite scalar global buffer accesses to use dcache bypass intrinsics for buffers that have writes (scalar BufferStore or MTE copy). Pure-read buffers keep normal BufferLoad (dcache path).

tilelang.ascend.transform.RewriteFp4ToFp4x2()

Retype genuine fp4 (float4_e2m1fn, lanes=1) storage into the 1-byte packed-pair form (float4_e2m1fnx2, lanes=2), halving fp4 element counts / offsets / indices so pointer arithmetic addresses bytes correctly. fp4 views over non-fp4 storage are left untouched.

tilelang.ascend.transform.RewriteFlagToBuf()

Normalize each hard_event’s synchronization into its 8-slot flag namespace. Sparse flag_ids that extend outside [0,8) and occupy at most 8 slots are compacted into [0,8) without spilling. When more than 8 slots are required, a 0/1 knapsack (capacity 8) keeps the subset of sync-point blocks that fills the flag slots best (renumbered into [0,8)); the rest spill to the shared 32-slot get_buf/rls_buf mutex pool. Hard_events whose flag_ids are already within [0,8) are untouched.

Must run after InferBufferAliases for manual schedules because alias inference uses set_flag/wait_flag as its liveness-graph anchors. The standard pipeline places this after MergeUBAllocations.