[Ascend][PTO] Vectorize T.Parallel to PTO VMI: pipeline framework and element-wise support

[jimmychou0]

Background and Goal

Goal: users write scalar-parallel T.Parallel programs (element-wise, control flow, reduce) inside T.SimdVF(lanes=L); the compiler vectorizes them into PTO VMI instructions, which PTOAS compiles to run correctly on Ascend 950. The full capability lands in tile-ai mainline over several PRs; each PR submission is a milestone.

Division of responsibility:

	TileLang side	PTOAS / toolchain side
Responsibilities	Deriving logical lane layouts, the vectorize pass pipeline, VMI instruction generation, Fragment logical mapping	Physical register splitting, mask materialization, Fragment buffer and vector register allocation, mem_bar insertion, etc.

Roadmap Milestones

Milestone	PR	Content	Depends on
M1	PR-1	T.Parallel vectorization core framework: element-wise compute; extent divisible / not divisible by lanes	—
M2	PR-2	Control flow support: conditional branches, lane-uniform serial for loops	M1
M3	PR-3	T.reduce tileop support: max/min/sum reduction	M2
M4	PR-4	Fragment Buffer UB capacity planning & allocation, PTOAS vector register allocation	M1 + PTOAS
M5	PR-5	Non-contiguous access and PACK data types: gather/scatter, fp4 packed data type support	M4
M6	PR-6	reducer support: reduction via T.alloc_reducer + T.finalize_reducer	M2 + M4

Scope of This PR (PR-1)

1. VMI instruction support: tl.vmi.* instruction-entry registration and VMI codegen, covering data movement, arithmetic, broadcast and lane index, mask creation, type conversion, etc.; codegen translates each tl.vmi.* call into pto.vmi.* PTODSL text for PTOAS, which handles mask materialization, register allocation, and optimization.
2. Pass pipeline: LayoutInference for PTO (lane-layout inference and annotation), Verify (input-contract checking), Legalize (normalization; framework init only in this PR), Vectorize (uniform conversion of the whole SIMD_VF region).
3. 1D/2D Parallel: 1D maps directly onto lanes; 2D recognizes two tightly nested Parallels as one iteration space, vectorizes the selected dimension, and keeps the other as an outer serial loop.
4. Divisible and non-divisible extents: when extent is divisible by lanes, all loops in the region share one constant full mask; otherwise a scalar remaining drives a dynamic prefix mask per chunk.
5. Element-wise compute: binary (+, -, *, /), unary (fabs/exp/log/sqrt), type conversion (f32↔fp8, f16/bf16↔f32, integer width changes).

1. Overall Transformation Flow

flowchart LR
    A["T.Parallel source<br/>inside T.SimdVF(lanes=L)"] -->|"LayoutInference for PTO<br/>layout_inference.cc · parallel.cc"| B["TIR with lane-layout annotation<br/>pto_parallel_loop_layout"]
    B -->|"Verify (read-only)<br/>verify_parallel_to_pto.cc"| C["TIR passing the input contract"]
    C -->|"Legalize (normalize)<br/>legalize_parallel_to_pto.cc"| D["Normalized TIR<br/>(framework init only in this PR)"]
    D -->|"Vectorize (uniform conversion)<br/>vectorize_parallel_to_pto.cc<br/>+ parallel_to_pto_utils.cc"| E["Vectorized TIR<br/>tl.vmi.* calls + chunk loop + mask"]
    E -->|"PTO codegen<br/>codegen_pto.cc"| F["pto.vmi.* PTODSL text<br/>to PTOAS"]

Loading

Flow notes:

1. Entry and dispatch. T.SimdVF(lanes=L) marks the region in the frontend (writing the tl.simdvf_lanes annotation); a kernel built with target=pto enters Verify→Legalize→Vectorize, while anything else takes the original AscendSimdVFLowerParallel path, leaving the Ascend compilation path unaffected.
2. Stage contracts. Verify checks the input read-only and stops at the first error; Legalize normalizes (framework init only in this PR); Vectorize converts the whole SIMD_VF region uniformly, producing vectorized TIR made of tl.vmi.* calls.
3. Emission. PTO codegen (codegen_pto.cc) translates each tl.vmi.* call into pto.vmi.* PTODSL text, handed to PTOAS for mask materialization, register allocation, and optimization.

2. Modules

2.1 LayoutInference for PTO

Role: pick a lane layout for each Parallel inside a PTO SIMD_VF region and write it back as an IR annotation for Verify/Vectorize to consume.

Aspect	Content
Input → output	T.Parallel inside a SIMD_VF region with target=pto → pto_parallel_loop_layout annotation
Core implementation	parallel.cc adds a PTO branch to Parallel's layout candidates (builds a lane layout with vectorize_size=1 when in_pto_simd_vf); layout_inference.cc recognizes PTO regions and writes the annotation; the annotation key is independent of SIMT's parallel_loop_layout, so the two backends never misread each other
Key convention	The Fragment object only expresses the lane-to-iteration mapping; it introduces no Fragment storage. 1D/2D dimension selection is supported; in 2D only the selected dimension is padded

Mapping for 1D T.Parallel(E) (VL = lanes):

lane = i % VL
q    = i // VL        # i.e. i = q * VL + lane

A 2D T.Parallel(M, N) appears in TIR as two tightly nested Parallels and is recognized as one 2D iteration space; the layout annotation goes only on the outermost For. Only the selected dimension is vectorized; the other stays as an outer serial loop. With P_N = ceil(N/VL)×VL and P_M = ceil(M/VL)×VL:

Vectorized dim	Layout input shape	Flattened coordinate t	lane and chunk	Generated loops
j	(M, P_N)	t = i×P_N + j	lane = t % VL, q = t // VL	i = k (outer serial), j = q×VL + lane
i	(P_M, N)	t = j×P_M + i	lane = t % VL, q = t // VL	j = k (outer serial), i = q×VL + lane

2.2 Shared Implementation Layer (parallel_to_pto_utils + pto_index_analysis)

Verify and Vectorize must make the same judgments over the same program — whether it is convertible, what each address pattern is, how each cast lowers. These capabilities are implemented once in the shared layer and used by both passes, so the rules and their implementations cannot drift apart.

Group	Function	Role
Layout inversion	InvertPtoLaneLayout	Inverts an annotated Parallel into the lane/chunk/non-vectorized-dimension mapping, producing PtoInverseMapping (lanes, index_expr, chunk_count, tail_lanes, 2D selection)
Access classification	AnalyzeBufferAccess	Classifies one Buffer access as continuous / same address on every lane (lane-uniform) / unsupported, producing MemoryAccess (address regularity expressed with ElemOffset and proven equivalent by the arithmetic analyzer over the given value ranges)
Region scan	ScanPtoRegion / DiscoverUnitsShared	Identifies the Parallel loops to convert in a region, plus shared state
Lane use	AnalyzeLaneUse	Distinguishes indices used only for addresses from those used as computed values (the latter need vci)
Cast planning	PlanValueCast	Produces a conversion plan per Cast (target dtype, saturate, rounding), covering the supported conversions such as FP8↔f32 and f16/bf16↔f32; gives a reason when not convertible
Whitelists	IsSupportedElementDType / IsSupportedVectorBinaryOp / FindUnsupportedExprNode	Admission sets for element types, binary ops, and expression nodes

pto_index_analysis: Index Substitution and Padding Arithmetic

When the user writes x = i + 1 and then accesses A[x], judging the address regularity of A[x] requires substituting x back to i + 1. ResolveIndexAliases performs exactly this substitution: it walks Bind definition chains and expands index variables layer by layer until the expression contains only loop variables, constants, and BufferLoads, then hands it to the address-regularity check. Whether an address is continuous or lane-uniform decides which vectorization path applies, and real programs often chain several such indirections.

Rules of ResolveIndexAliases:

• Recursive substitution with cycle detection and a depth cap of 64 — exceeding the cap rejects rather than silently truncating;
• The expression node count is re-counted after each round and rejected past a budget of 4096 (preventing chains like x1=i+i; x2=x1+x1; … from exploding exponentially);
• A Bind whose definition contains a BufferLoad is a value, not an index expression, and is never substituted — substituting would duplicate a memory read.

The same file provides CheckedPtoPadding: an overflow-guarded version of "pad the logical extent up to a multiple of lanes". The lane loop indexes with int32, so an extent beyond int32 is rejected outright.

2.3 Verify: Input Contract Checking

Role: read-only checking of the input TIR; accepts forms that Legalize can normalize later; aborts compilation at the first error. Diagnostics carry the [VerifyParallelToPTO] prefix with a category and a location.

Main rejection scenarios, grouped by diagnostic category:

Category	Rejection scenarios (examples)
lanes	SIMD_VF containing Parallel without the tl.simdvf_lanes annotation; annotation not an Int64 constant; illegal lanes value
Region admission	Hand-written tl.vmi.* mixed with a converting Parallel; volatile_scope attribute; region metadata referencing lane-varying values
Loop structure	serial/pipelined/while loops, nested blocks, or conditional branches inside a Parallel (this round); ThreadBinding regions or wrappers; Vectorized wrappers
Statement forms	Buffer allocation/declaration inside a Parallel; assert statements
Expression	Cross-lane reduce; expression nodes outside the whitelist
Memory access	Address neither continuous nor lane-uniform; GM (global memory) access; address patterns that cannot be proven regular
dtype/cast	Element types outside the whitelist; FP8 casts outside the supported set; lane-uniform scalar FP8 casts
Layout	Missing pto_parallel_loop_layout; non-invertible layout

Verify and Vectorize share every analysis function from §2.2, so the two passes cannot disagree on the same input.

2.4 Legalize: Normalization

In PR-1 this pass only initializes the pass pipeline framework and rewrites nothing. It is introduced so that PR-2's control-flow support — normalization of IfThenElse/Min/Max→Select, etc. — lands here, keeping Vectorize's input contract unchanged.

2.5 Vectorize: Converting the Whole SIMD_VF Region Uniformly

Core class VectorizeParallelToPTOImpl. The entry VisitStmt_(const SBlockNode*) finds a PTO SIMD_VF region → ScanPtoRegion scans it → each Parallel loop is inverted by InvertPtoLaneLayout → converted by ConvertUnit → spliced back by ReplaceStmt.

ConvertUnit: the Main Flow for One Parallel Loop

flowchart TD
    M["Input: one Parallel loop + PtoInverseMapping"] --> S1["Reset per-loop state<br/>bind_env / value_map / temp counter"]
    S1 --> S2{"extent divisible by lanes?"}
    S2 -->|"divisible"| S3["Reuse the region-wide shared full mask"]
    S2 -->|"non-divisible"| S4["Create the remaining scalar (init: logical extent)"]
    S3 --> S5["ConvertBody: statement-by-statement<br/>(via ConvertExpr / ConvertBinary)"]
    S4 --> S5
    S5 --> S6["DeadBindElim local DCE"]
    S6 --> S7["Assemble the chunk loop (induction variable q):<br/>mask bind → body → remaining update<br/>2D: the non-vectorized dimension re-emitted as an outer serial loop"]
    S7 --> S8["ReplaceStmt splices over the original Parallel"]

Loading

Two key decisions:

• Mask strategy: divisible extents share one region-wide constant full mask; non-divisible extents create a remaining scalar (initialized to the loop's logical extent, decremented by lanes per chunk), and each chunk binds a dynamic prefix mask via create_mask(remaining, L), runs the body, then updates remaining. The mask variable must exist before the body is converted — the vstore and vector ops in the body reference it.
• 2D Parallel loops: the selected dimension is vectorized; the other dimension is re-emitted as an outer serial loop wrapping the chunk loop.

After the body is converted, a local DCE (DeadBindElim) removes leftover bindings that only served address derivation (dead vci/vbrc/vadd/vload chains).

ConvertExpr: Expression Conversion Dispatch

Input	Path	Product
BufferLoad, continuous address	Classified by AnalyzeBufferAccess, pointer via tl.access_ptr (extent=VL)	vload
BufferLoad, lane-uniform	Scalar BufferLoad kept (indices substitute lane=0, chunk=q), bound to a scalar Var	sload; each textual occurrence binds separately — no cross-access load reuse
loop_var used as a value	Substituted through the inverse mapping and simplified	vci
Cast	Plan taken from PlanValueCast; identity casts pass through	Vector vcvt with explicit saturate/rounding attributes — backend defaults are never relied on
Unary math Call (fabs/exp/log/sqrt)	Table-mapped to the matching VMI unary op	vabs/vexp/vln/vsqrt (with mask)
Add/Sub/Mul/Div	ConvertBinary	See the table below
Anything else	Unreachable (already rejected by Verify); reaching it is an internal error	—

Every intermediate value is materialized through EmitBind as a named Bind (vload_0, brc_1, …), keeping the output in a stable SSA form.

ConvertBinary: the Scalar/Vector 2×2 Dispatch

lhs \ rhs	Scalar	Vector
Scalar	Stays scalar arithmetic; no VMI emitted	Scalar side broadcast with vbrc, then vadd/vsub/vmul/vdiv (with mask)
Vector	Scalar side broadcast with vbrc, then vadd/vsub/vmul/vdiv (with mask)	vadd/vsub/vmul/vdiv directly (with mask)

[KurrinQu]

simt的修改，似乎不需要？

[jimmychou0]

补充说明（更正我上一条回复）：这个守卫是必需的，已在 2dbf83d 恢复。依据：本 PR 的验证方法论是 ptoas vmi-v0.1.9（released vmi 轨道，与 PR#1 基线一致），该轨道的 ptodsl 不带 ptodsl._allreduce；而生成的 kernel 源码导入 tilelang.contrib.ptodsl（其 init 顶层 import .simt），无守卫时任何 VMI kernel 会在导入期 ModuleNotFoundError。守卫只影响 vmi 轨道缺失该模块的场景，主线行为不变（238 上主线 ptoas-0.66 + 无守卫版与 vmi-0.1.9 + 守卫版均已实测通过）。若倾向不改此文件，替代方案是把 contrib/ptodsl/init.py 改为 PEP 562 懒加载，VMI 路径就不会被 SIMT 内部依赖连累——可另行提案。

[KurrinQu]

可能要根据主线上使用的PTOAS版本才能确定是否需要这里的修改。

[KurrinQu]

这个文件应该不需要上传？可以考虑替换成 PR #2 中的element-wise的example

[jimmychou0]

同意，已删除（9b5354c8）：example_simdvf_vecadd_lower.py 与依赖它的 test_simdvf_vecadd_pto_run.py 一并去掉。element-wise 的设备端覆盖由本 PR 自带的 test_vectorize_parallel_pto_run.py 承担（vecadd 1D/2D、tail mask、cast，auto-schedule 保持开启）；PR #2 合入后 t_parallel_st_suite 的 sv1 作为 element-wise 的权威 example。

[KurrinQu]

commit 标题和正文信息中存在很多fork仓中（deepseek/wenbo)的pr链接信息，因为这俩都是私仓，带上了现在指向的链接也不对，指向的是tile-ai仓的pr，是否需要带上？我建议是squash commit或rebase重写全部commit message。

[WenboCodes]

commit 标题和正文信息中存在很多fork仓中（deepseek/wenbo)的pr链接信息，因为这俩都是私仓，带上了现在指向的链接也不对，指向的是tile-ai仓的pr，是否需要带上？我建议是squash commit或rebase重写全部commit message。

+1

[WenboCodes]

tilelang/ascend/language/vmi.py: the T.vmi.* surface (exported via tilelang.ascend.language) — vector loads/stores, masks, arithmetic/bitwise/shift, vcvt (incl. packed FP4 sources), vinterpret_cast, vzip/vunzip, precision-aware vdiv, gather/scatter, histogram ops, register allocation helpers

T.vmi的语言前端在这个PR里面是否是必要的？

contrib/ptodsl/simt.py: optional ptodsl._allreduce import (absent from vmi-track releases and mainline ptoas ≥ 0.67; the contrib package is imported by every generated kernel)

这个是什么，为什么需要simt的支持？

2. feat(pto-parallel): vectorize T.Parallel inside SimdVF to PTO VMI

• Verify → Legalize (identity stub) → Vectorize pipeline ahead of AscendSimdVFLowerParallel, with routing tests; CF rewriting stays out of scope
• Verify: shared region scan + inverse lane mapping, PTO LayoutInference branch, 1D/2D lane mapping, T.SimdVF(lanes=...), contract tests
• Vectorize: chunk loops with dynamic prefix tail masks, remaining reset, dtype/cast codegen, element-wise/tail/FP8 device tests
• Integration: shared mask emitted inside the vecscope (VPTO rejects pset values crossing the boundary)

这部分的AI味过于浓重了，看不懂这个PR到底支持了什么特性

@jimmychou0 @KurrinQu

[jimmychou0]

针对 @WenboCodes 的三个问题，description 已重写（用户视角优先，附 before/after 示例）：

1. T.vmi 前端在这个 PR 是否必要？ 必要：T.Parallel 自动向量化降低到 tl.vmi.* TIR op，缺 builtin 注册和 pto.vmi.* 发射就无法落地；且向量化测试复用的手写 VMI 用例也依赖前端。这与 roadmap PR-1 的口径一致（"T.Parallel 核心向量化（含 VMI 语言面）"）。如果 reviewer 更倾向拆分（VMI 前端单独先行一个 PR），当前分支正好按两个 commit 切开，可以无损 rebase 拆分。

2. simt.py 那条不是"SIMT 支持"，是一行导入容错：每个生成的 kernel 都会 import tilelang.contrib.ptodsl（codegen 发射的固定头部），该包 init 顶层连带 import simt 子模块，后者 import ptodsl 内部模块 ptodsl._allreduce——vmi 轨道和主线 ≥0.67 的 ptodsl 没有这个模块，无守卫时任何 kernel（包括纯 VMI）在导入期就挂。守卫在 ≤0.66 上零行为变化；缺模块时只有 SIMT allreduce 辅助函数在被调用时报清晰错误。本 PR 其余部分不使用 SIMT。

3. AI 味/看不懂：已把 description 改成"先说用户能写什么"——开头就是一个 T.Parallel(300) 在 SimdVF(lanes=128) 里的完整例子（生成 vload/vmul/vadd/vstore + 44-lane 动态尾掩码），然后列本期支持（1D/2D、整除/非整除、cast）与不支持（控制流/规约/gather，后续 PR）。pass 管线等实现细节降级到后面的 Implementation notes 小节。

[WenboCodes]

针对 @WenboCodes 的三个问题，description 已重写（用户视角优先，附 before/after 示例）：

1. T.vmi 前端在这个 PR 是否必要？ 必要：T.Parallel 自动向量化降低到 tl.vmi.* TIR op，缺 builtin 注册和 pto.vmi.* 发射就无法落地；且向量化测试复用的手写 VMI 用例也依赖前端。这与 roadmap PR-1 的口径一致（"T.Parallel 核心向量化（含 VMI 语言面）"）。如果 reviewer 更倾向拆分（VMI 前端单独先行一个 PR），当前分支正好按两个 commit 切开，可以无损 rebase 拆分。

2. simt.py 那条不是"SIMT 支持"，是一行导入容错：每个生成的 kernel 都会 import tilelang.contrib.ptodsl（codegen 发射的固定头部），该包 init 顶层连带 import simt 子模块，后者 import ptodsl 内部模块 ptodsl._allreduce——vmi 轨道和主线 ≥0.67 的 ptodsl 没有这个模块，无守卫时任何 kernel（包括纯 VMI）在导入期就挂。守卫在 ≤0.66 上零行为变化；缺模块时只有 SIMT allreduce 辅助函数在被调用时报清晰错误。本 PR 其余部分不使用 SIMT。

3. AI 味/看不懂：已把 description 改成"先说用户能写什么"——开头就是一个 T.Parallel(300) 在 SimdVF(lanes=128) 里的完整例子（生成 vload/vmul/vadd/vstore + 44-lane 动态尾掩码），然后列本期支持（1D/2D、整除/非整除、cast）与不支持（控制流/规约/gather，后续 PR）。pass 管线等实现细节降级到后面的 Implementation notes 小节。

1. 不对吧。需要tl.vmi.*不代表需要T.vmi前端，没看到这个PR有手写VMI用例。
2. 导入容错是个啥？