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.
NormalizePTOAllocMarkers()	Move VMI allocation kinds onto scoped carriers in the owning VF.
ImportPTOAllocMarkers()	Attach normalized VMI allocation kinds to final AllocBuffer nodes.
LegalizePTOVMIMemoryAccess()	Build and verify the explicit VMI memory descriptor before scheduling.
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.
VerifyParallelToPTO()	Read-only verification of the PTO Parallel input contract.
LegalizeParallelToPTO()	Placeholder for the PTO Parallel legalization stage.
VerifyPTOCodegenLegality()	Verify PTO requirements before code generation without modifying IR.
VectorizeParallelToPTO()	Rewrite verified T.Parallel loops inside PTO SIMD_VF blocks into
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.

Each cloned VMI allocation gets a fresh data Var, and the corresponding tl.vmi.alloc_kind AttrStmt carrier emitted by NormalizePTOAllocMarkers is remapped to that Var through the same mapping used for buffer loads and stores. A carrier whose allocation was not rebuilt by the clone keeps its original Var.

tilelang.ascend.transform.NormalizePTOAllocMarkers()

Move VMI allocation kinds onto scoped carriers in the owning VF.

Consumes each SBlock’s tl.vmi.alloc_kind map and emits one AttrStmt(data_var, "tl.vmi.alloc_kind", StringImm(kind)) per marked allocation in the block that actually owns it, in alloc_buffers order. Carriers for a SIMD_VF block go inside its existing outer tl.simdvf_scope statement; carriers for a nested block go in that block’s body. Allocations, data Vars, statement order, and unrelated annotations stay untouched, and marked allocations owned outside a SIMD_VF are rejected.

The pass is idempotent: existing carriers and final allocation annotations are not duplicated. Run it after binding the Ascend PTO target and before UnrollLoopSkipVF. Functions without a PTO target pass through unchanged.

tilelang.ascend.transform.ImportPTOAllocMarkers()

Attach normalized VMI allocation kinds to final AllocBuffer nodes.

Must run immediately after LowerOpaqueBlock on the PTO target. It collects the temporary tl.vmi.alloc_kind AttrStmt carriers and the actual AllocBuffer statements before rewriting, matches them by data Var identity, writes the kind to the matched allocation’s annotations, and removes the consumed carriers while preserving their bodies and order. An orphan carrier or conflicting kinds are rejected. The final allocation need not be lexically inside tl.simdvf_scope, because lowering materializes it before the scope statement. Functions without a PTO target pass through unchanged.

tilelang.ascend.transform.LegalizePTOVMIMemoryAccess()

Build and verify the explicit VMI memory descriptor before scheduling.

Every ordinary T.vmi.vload / T.vmi.vstore address (a scalar-origin BufferLoad) is replaced with tl.access_ptr(BufferLoad, extent, rw_mask) plus a rewritten instruction scalar offset. The pass proves, in one shared acceptance exit, that the descriptor preserves the source address, contains the complete physical access envelope and covers the declared extent under the existing logical pointer-region analysis. An ordinary caller-written explicit descriptor is preserved and only verified; there is no whole-allocation fallback for an ordinary access. Indexed vgather / vscatter addresses, including one reached through a direct handle binding, are instead rebuilt as an owner-origin descriptor with capacity prod(owner.shape) after the owner identity, layout, capacity, stable-coverage and bound-address snapshot proofs. Functions without a PTO target pass through unchanged; the pass is a no-op on IR without VMI accesses.

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

Read-only verification of the PTO Parallel input contract.

Runs on the PTO target only: SIMD_VF blocks whose T.Parallel units fall outside the supported mapping contract are rejected, while functions without a PTO target pass through unchanged.

tilelang.ascend.transform.LegalizeParallelToPTO()

Placeholder for the PTO Parallel legalization stage.

Only the registration exists for now, so that the formal PTO pipeline order is fixed; the safe scalar If/Select/Min/Max rewrites are left to a later change. The placeholder preserves the input IR and rejects nothing.

tilelang.ascend.transform.VerifyPTOCodegenLegality()

Verify PTO requirements before code generation without modifying IR.

Consumes the final device TIR after LowerDeviceKernelLaunch and the final layout and UB rewrites. Currently checks destination alignment for default/continuous explicit-mask VMI stores of b8/b16/b32 data: 32 bytes for ordinary stores, or natural element alignment for recognized 1PT writeback. Bool data is excluded. Shapes, layouts, masks and addresses are unchanged. Functions without a PTO target pass through unchanged.

tilelang.ascend.transform.VectorizeParallelToPTO()

Rewrite verified T.Parallel loops inside PTO SIMD_VF blocks into logical VMI code.

Runs on the PTO target only; functions without a PTO target pass through unchanged.

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.