[Likai-19]

Background and Goal

The PTO vectorization framework (PR-1, tile-ai/tilelang#3512) provides:
LayoutInference → Verify → Legalize → Vectorize → VerifyPTOCodegenLegality → codegen, with element-wise
T.Parallel support (1D/2D, divisible and tail extents).

This PR is milestone M2 (control flow): programs with data-dependent behavior — conditional branches,
Min/Max, boolean expressions, per-lane accumulators and loop-carried values — can now be written as
scalar-parallel T.Parallel and are lowered to the same VMI instruction set. Program shape is unchanged:
everything still lives inside T.SimdVF(lanes=L) with T.Parallel units.

	TileLang side (this PR)	PTOAS / toolchain side
Responsibilities	Control-flow normalization, admission rules, carry promotion decisions, region splitting, per-lane scalar promotion	Mask materialization, register allocation, mem_bar insertion, fragment/slot allocation

Relationship to PR-1

Milestone	PR	Content	Depends on
M1	PR-1	element-wise core framework	—
M2	PR-2 (this PR)	control flow: conditional branches, lane-uniform serial loops, per-lane accumulators	M1
M3	PR-3	T.reduce tileop	M2
M4	PR-4	Fragment Buffer UB capacity planning & allocation	M1 + PTOAS
M5	PR-5	gather/scatter, non-contiguous access, PACK types	M4

Scope of This PR (PR-2)

1. Control-flow normalization (Legalize): if/else writing the same destination → Select;
   single-sided writes → Select(cond, new, old); Min/Max → compare + Select; and/or/not →
   nested Select. Legalize only rewrites expressions — it never emits VMI, so Vectorize keeps a single
   expression-dispatch path.
2. Control-flow vectorization (Vectorize): Select → vsel; comparisons → vcmp (vector/vector) or
   vcmps (vector/scalar); boolean ops → mask_and/mask_or/mask_not. A lane-uniform condition is
   materialized through vbrc + vcmp, so a condition that only depends on a non-vectorized coordinate
   is supported.
3. Loop shells: a constant-bound T.serial inside a T.Parallel unit is admitted and preserved;
   only its body is vectorized. When the body is the canonical carry pattern (seed store immediately
   before the loop, then same-index self read-modify-write), the carried value is promoted to a
   local.var vreg, so the carry becomes an SSA value (scf.for iter_args) and no UB round-trip
   is needed.
4. Region-level carry (Plan B): a serial loop wrapping units, whose carried buffer is exactly one
   lane vector wide with a lane-invariant seed, is promoted the same way, with one broadcast at region
   level and one write-back after the loop.
5. Vector-scope splitting (Plan A): when the carried value is not a single full lane vector (e.g.
   per-warp slices), each iteration gets its own vecscope: the loop is hoisted out of the region and the
   region is emitted around the loop body, so a scope transition — the only legal synchronization point —
   separates consecutive iterations. Nested carriers are hoisted level by level.
6. Per-lane scalar accumulators (T.alloc_var): a single-element local.var inside a unit is a
   private register per lane, promoted to a local.var vector; scalar seeds broadcast, lane-varying
   seeds are stored as-is.
7. UB round-trip elimination (new pass): eliminate_simdvf_roundtrip.cc removes region-local
   "staging" buffers (T[i] = expr followed by OUT[i] = T[i]) inside a SIMD_VF region before
   vectorization, because a vstore/vload pair inside one vecscope has no ordering guarantee.
8. Region-local allocation hoisting: when a region is split, every region-local UB allocation is
   hoisted to the enclosing block so all parts share one allocation. Otherwise each part would
   re-declare the allocation and MergeUBAllocations would hand out several merged slots for the same
   buffer.
9. Fail-closed policy: any shape whose carry cannot be promoted or split (ambiguous multi-carrier,
   sibling carriers, cross-statement binds) is rejected with an actionable diagnostic; there is no fallback to the plain
   conversion, which has no ordering guarantee.

1. Overall Transformation Flow

flowchart TD
    A["T.Parallel source<br/>inside T.SimdVF(lanes=L)"]
    A -->|"LayoutInference for PTO"| B["TIR + lane-layout annotation"]
    B -->|"Verify (read-only)<br/>control flow / shells / local.var admitted"| C["TIR passing the input contract"]
    C -->|"Legalize<br/>if→Select, Min/Max→Select, boolean expansion"| D["Normalized TIR<br/>(no control flow left in Parallel)"]
    D -->|"EliminateSimDVFRoundTrip<br/>drop staging buffers inside a vecscope"| D2["TIR without UB round-trips"]
    D2 -->|"Vectorize<br/>carry promotion / vecscope split / unit conversion"| E["Vectorized TIR<br/>tl.vmi.* + vregs + one vecscope per iteration"]
    E -->|"VerifyPTOCodegenLegality"| E2["Final TIR"]
    E2 -->|"PTO codegen"| F["pto.vmi.* PTODSL text → PTOAS"]

Loading

Notes:

1. Nothing new on the Ascend (non-PTO) path: every added pass and rule is gated on TargetIsPTO.
2. Hand-written VMI regions are untouched: regions without T.Parallel are skipped by all passes.
3. Order matters: the round-trip elimination runs before Vectorize (it must see the plain TIR), and
   the allocation hoisting runs inside Vectorize (it must see the split).

2. Modules

2.1 Legalize: control-flow normalization

ControlFlowMutator rewrites only expressions and statements that can be expressed as values.

Input form	Rewrite	Why
if (c) T[addr] = a else T[addr] = b (same destination)	T[addr] = Select(c, a, b)	same address on both arms → a value choice
if (c) T[addr] = a (single-sided)	T[addr] = Select(c, a, T[addr])	preserves the "keep old value" semantics
Min(a, b) / Max(a, b)	Select(a < b, a, b) / Select(a > b, a, b)	reuses the compare+select path
a and b / a or b / not a	nested Select	boolean is a lane mask

Verify only admits the forms Legalize can normalize (same-destination stores, nested conditionals of
the same shape, empty branches). Anything else is rejected with a reason.

2.2 Verify: admission rules for control flow

Category	Newly admitted	Still rejected
Conditional branch	if/else whose arms store the same buffer+indices (or are empty / nested), condition value-classifiable	divergent control flow, lane-varying stores
T.serial in a unit	constant-bound, step 1, induction used as a lane-uniform value	runtime bound, while, unsupported wrappers
local.var	single-element T.alloc_var inside a unit	local.var at region level, non-single-element
Region-local allocation + split	handled by hoisting (see §2.5)	—

2.3 Vectorize: control-flow ops

IR node	Emitted VMI
Select	vsel(cond, a, b)
EQ/NE/LT/LE/GT/GE	vcmp (both sides vector) / vcmps (one side scalar)
And/Or/Not on masks	mask_and / mask_or / mask_not
lane-uniform condition	vbrc the operands, then vcmp (all-true / all-false mask)

2.4 Vectorize: loop shells and carry promotion

AnalyzeLoopShell recognizes the canonical carry inside a preserved shell:

seed store (immediately before the loop, must not read the carried buffer)
for t in serial(const N):
    B[lane] = f(B[lane], ...)        # exactly one store, ≥1 load, same lane index

and rewrites it to

vreg = vbrc(seed)                    # region-level broadcast
for t in serial(N):
    vreg = f(vreg, ...)              # plain SSA update → scf.for iter_args in PTODSL
vstore(vreg, B)                      # one write-back after the loop

Shapes outside this pattern are rejected with [VectorizeParallelToPTO] loop shell ….

2.5 Vectorize: Plan A (vecscope split) and Plan B (region-level carry)

Plan B promotes a region-level carried buffer when it is exactly one lane vector wide, has a
lane-invariant constant seed, and no other unit touches it:

SIMD_VF { <seed unit>; for t { <unit RMW B> }; <tail> }
  →  SIMD_VF { <seed unit → vreg = vbrc(seed)>; for t { <unit RMW vreg> }; <tail reads vreg> }
     + one vstore(vreg, B) after the loop

Plan A covers everything else (e.g. per-warp slices). Each iteration gets its own vecscope:

SIMD_VF { pre; for t { U(t) }; post }   →   SIMD_VF { pre };  for t { SIMD_VF { U(t) } };  SIMD_VF { post }

decisions:

• exactly one promotable carrier → Plan B; more than one → Plan A splits every carrier into
  its own region;
• sibling carriers inside a loop are each split into their own region;
• region-local UB allocations are hoisted to the enclosing block before the parts are built, so the
  parts share one allocation. The hoist set covers every shared/shared.dyn allocation of the
  region, not just the carrier, so a bystander buffer written before the loop and read inside it is
  not duplicated;
• unsplittable shapes fail closed: a carrier that reads a value bound outside the loop cannot
  cross a region boundary and is rejected with a diagnostic.

2.6 Vectorize: per-lane scalar promotion (T.alloc_var)

A single-element local.var buffer is a register per lane. Its accesses are promoted to a local.var
vector (<name>_lane), so m = T.max(m, v) inside a shell becomes a vector update of vreg and the
value stays in RF. LayoutInference skips the lane-address classification for this buffer
(IsLaneScalarBuffer), and Verify admits it only inside a converting unit.

2.7 EliminateSimDVFRoundTrip: a new pass

Removes region-local staging buffers before vectorization:

Rule	Rewrite	Preconditions
forward re-addressing	T[i] = expr + adjacent pure copy OUT[i] = T[i] → OUT[i] = expr	T region-local, single writer, single consumer unit; OUT untouched otherwise
backward copy propagation	pure copy T[i] = S[i] + adjacent consumer → consumer reads S[i]	T single write/read; S never written in the region

Iterated to a fixpoint, then dead allocations are pruned. Conservatively skipped: serial wrappers,
region-external readers, multiple writers, extent mismatch, tails, 2D units, predicated access,
non-identity indices.

3. Acceptance and verification

3.1 Repository tests

Suite	Result
Structural (testing/ascend/language/pto/, no device)	312 passed
Device (testing/ascend/language/pto/, NPU)	509 passed, 0 failed

New structural tests: test_simdvf_parallel_control_flow.py (10), test_simdvf_parallel_loop_shell.py (6),
test_simdvf_roundtrip_elimination.py (10), test_simdvf_parallel_lane_scalar.py (4),
test_simdvf_region_local_carry.py (3).

4. Boundaries

• Not in this PR: fragment/slot storage (M4), gather/scatter and stride ≠ 1 (M5), T.reduce (M3).
• Float Min/Max semantics: Min/Max → Select is not IEEE fmin/fmax — NaN is not propagated
  and -0/+0 selection follows the comparison. Kernels that rely on NaN propagation must not use
  Min/Max.
• Loop shells are admitted for constant bounds only; runtime bounds remain rejected.
• Split cost: each iteration becomes its own vecscope, so a split trades scope transitions for
  ordering guarantees. Carry promotion (vreg) avoids that cost where it applies.

5. Support matrix (at a glance)

5.1 Newly supported by this PR

Form	Lowering
if/else writing the same destination (multi-store arms included; unequal arm counts resolve to "keep old value")	Select
single-sided if	Select(cond, new, old)
nested if	nested Select
explicit T.Select	vsel
== != < <= > >=	vcmp / vcmps
and/or/not on masks	mask_and / mask_or / mask_not
Min/Max	Select (float NaN/±0 caveat, see §4)
lane-uniform condition	vbrc + vcmp
constant-bound T.serial inside a unit	preserved shell, body vectorized
canonical carry inside a shell	local.var vreg
T.serial wrapping units	Plan B promotion or Plan A split
nested / sibling carried loops	promoted or split level by level
T.alloc_var inside a unit	local.var vector
region-local T.alloc_shared surviving a split	hoisted, one shared allocation
region-local staging buffer across units	removed by EliminateSimDVFRoundTrip

5.2 Rejected with a diagnostic

Form	Diagnostic
while in a unit	[VerifyParallelToPTO] loop structure: while loops inside T.Parallel are not yet supported
runtime-bound T.serial	… must have constant bounds and a positive extent
T.Pipelined / non-serial loop	… only T.serial loops inside T.Parallel are supported
nested block inside a unit	nested blocks inside T.Parallel are not yet supported
branch arm containing a non-store statement	If … must be buffer stores, nested conditionals or no-ops
divergent destinations	[LegalizeParallelToPTO] paired then/else stores must write the same buffer indices
T.Parallel under if/while	T.Parallel under a conditional branch is not yet supported
ThreadBinding in a region	thread bindings are not allowed inside a converting SIMD_VF block
hand-written tl.vmi.* mixed with T.Parallel	hand-written tl.vmi.* calls cannot be mixed …
3+ layer T.Parallel	[PTO layout inference] a T.Parallel nest with N layers is not supported
region-level T.alloc_var	only a single-element local.var inside a converting T.Parallel unit …
non-canonical carry / carrier reading an outer bind	[VectorizeParallelToPTO] loop shell: … (fail-closed)

5.3 Out of scope (owned by a later PR)

Capability	Owned by	Behaviour today
cross-lane reduce (T.reduce*)	M3	cross-lane reduce is not supported
fragment buffers / slot storage	M4 (PR-4)	fragment access rejected on this branch; lifted by PR-4
gather / scatter / stride ≠ 1	M5	no vectorizable dimension … neither continuous nor lane-uniform
GM access from a region	—	must not access GM directly (stage through UB)
volatile buffers	—	volatile accesses are not supported

5.4 Limitations

• Staging buffer inside one unit (tmp[i] = expr; y[i] = tmp[i]) compiles but leaves a
  vstore → vload pair inside a single vecscope, where there is no ordering point. Workaround: write
  the expression directly, or stage it across two units (the elimination pass then removes it).
• T.serial(0, 8, 2) is normalized by the frontend to a step-1 loop of extent 4, so the step != 1
  rejection is unreachable from the public API — this is correct behaviour, not a gap.
• Min/Max NaN and −0/+0 semantics differ from IEEE fmin/fmax (see §4).

[KurrinQu]

这个问题主要处理的问题和场景是？看注释怎么像是处理同步问题的。PTOAS主线已经合入mem_bar插入特性，需要确认这个pass是否需要。

[Likai-19]

这个pass已删除，PTOAS主线的mem_bar合入后已不需要该pass

[KurrinQu]

两侧 store 配对时，请完整保留内层分支“不满足条件就保持旧值”的语义。

Legalize 先递归处理内层 if。内层单侧写会成为带 predicate 的 BufferStore，但外层两侧写同一地址时，这里只取 th->value 和 el->value，最后输出无 predicate 的 Select store。这样会把原本“不满足内层条件就不写”的 lane 也覆盖掉。

可以在旧值可合法读取、写回的场景下，将内层单侧写正规化为 Store(y, Select(d, new, Load(y, i)), i)，再递归合成外层 Select；若继续沿用当前 predicated-store 路径，则必须合并内外层 predicate。现有 nested-if 测试的内层分支都有完整的 then/else 写入，覆盖不到这条路径，请补“两侧都是嵌套单侧写”的测试。

[Likai-19]

确认是问题，已修。现在按三种情况分流：只要两个分支中有任意一侧带 predicate，就不合并成 Select，而是两个分支各自发一条互补的谓词 store（pred = c / ¬c，并与内层 predicate 合取）；只有两边都无 predicate、且没有写后读时才走 Select 。

[KurrinQu]

** 按地址配对 store 时，需要保留两个分支各自的读写依赖。**

当前实现按 then 分支的 store 顺序输出结果，再查找 else 中同地址的 store。若两个分支的写入顺序不同，而且后一个表达式读取前一个写入的值，else 的语义会改变。例如 then 是“先写 y、再读 y 写 z”，else 是“先写 z、再读 z 写 y”；配对后固定按 y、z 输出，else 对 z 的读取就跑到了写入前面。

建议 Verify 检查配对和输出计划能否保留两侧各自的写入顺序，Legalize 使用同一计划生成 Select。允许两侧 store 数量不同；单侧多出的写入Legalize在合适位置补旧值，不能一律移到末尾。共同写入发生相对顺序冲突时明确拒绝。这样无需先实现完整依赖分析，也能保留常见的单侧额外更新。

[KurrinQu]

这里如果当前只支持常数次循环，可以接受。但是我没理解，如果是变量次循环且证明lane unifrom，因为目前的方案不依赖循环展开，支持有特殊的地方吗？

[KurrinQu]

Min/Max -> Select 的普通有限数值路径容易验证，建议再补 NaN 的两个操作数位置以及 +0/-0 两种顺序的测试。
PR 正文已经提示特殊值限制，但“NaN is not propagated”不准确：Select(a>b,a,b) 在 a=NaN,b=1 时选择 1，在 a=1,b=NaN 时仍会返回 NaN。建议把行为写具体。

[Likai-19]

已处理。
1、说明已经修改为：Min(a,b) / Max(a,b) 展开为 Select(a<b,a,b) / Select(a>b,a,b)，左操作数为 NaN 时比较为假、返回右操作数（NaN 被丢弃）；右操作数为 NaN 时原样返回（NaN 保留）；-0/+0 比较相等，因此右操作数获胜、符号随操作数顺序。
2、新增测试用例 test_simdvf_parallel_minmax_semantics_run.py：NaN 的左右两个位置 + (+0,-0)/(-0,+0) 两种顺序，min/max 各一遍，共 6 个用例，都验证通过了。

[KurrinQu]

新增 tl.vmi.mask_and/or/not 已有 op 注册和 PTODSL 名称映射，但 ClassifyVmiCallResult_ 仍只把 create_mask/vcmp/vcmps 标为 kVmiMask，三个新 op 会落到 kUnknown。建议同步补上结果分类，并覆盖直接结果和纯 alias 的查询。

[Likai-19]

已补充，ClassifyVmiCallResult_ 现在把 tl.vmi.mask_and / mask_or / mask_not 也归为 kVmiMask

[KurrinQu]

Select 的标量 bool 条件会在 Vectorize 阶段触发内部断言，而不是得到 PR 正文描述的 lane-uniform mask。

例如：

flag: T.bool

with T.SimdVF(lanes=64):
    for i in T.Parallel(128):
        if flag:
            y[i] = a[i]
        else:
            y[i] = b[i]

Legalize 后是：

BufferStore(y, Select(flag, BufferLoad(a, [i]), BufferLoad(b, [i])), [i])

flag 仍是 scalar bool，ConvertExpr(flag) 不会生成 lane mask，cond.dtype().lanes() > 1 失败。

请让标量 bool 条件走明确的广播路径，生成全真或全假的 lane mask；或者在 Verify 阶段明确拒绝这种条件并给出稳定诊断。当前实现既没有完成正文宣称的支持，也没有在输入检查阶段拒绝该形状。

问题场景

TileLang 源码：

@T.prim_func
def main(
    A: T.Tensor((128,), "float32"),
    B: T.Tensor((128,), "float32"),
    Y: T.Tensor((128,), "float32"),
    flag: T.bool,
):
    with T.Kernel(1):
        a = T.alloc_shared((128,), "float32")
        b = T.alloc_shared((128,), "float32")
        y = T.alloc_shared((128,), "float32")
        T.copy(A, a)
        T.copy(B, b)
        with T.SimdVF(lanes=64):
            for i in T.Parallel(128):
                if flag:
                    y[i] = a[i]
                else:
                    y[i] = b[i]
        T.copy(y, Y)

Legalize 已经把分支改写成 Select，但 Select 条件仍是 scalar bool。Vectorize 的 ConvertExpr(flag) 返回 scalar bool，之后在 vectorize_parallel_to_pto.cc:2207 的 lane 检查处失败。

支持该输入的一个建议：对 scalar bool 先生成 vbrc(flag)，再用一个与 mask 类型匹配的比较生成全真/全假 mask；或者提供一个明确的 scalar-bool-to-mask helper，避免用不适合 bool dtype 的通用数值比较。

[KurrinQu]

arms_ordered 目前只要发现一个 arm 先写 buffer、再读它，就以“vecscope 内没有 ordering point”为由拒绝。PTOAS 主线已合入 PR307，默认分析 vecscope 内的 UB 读写依赖并插入 mem_bar。TileLang 不应因为存在先写后读（RAW）或先读后写（WAR）就限制输入，也不需要在 Legalize/Vectorize 中重复实现这类同步分析和 barrier 插入。

TileLang 这里需要检查的是：将 if/else 合并为 Select 或 predicated stores 后，是否仍保留两侧原有的访存语义。当前按地址配对的路径采用 then 分支的 store 顺序，如果 else 的依赖顺序与它相反，就不能直接套用。例如 then 是 y=1; z=y+1，else 是 z=2; y=z+1；若统一按 y、z 输出，else 就会先读旧 z，再写 z=2，结果已经改变。统一按 else 顺序也会破坏 then。

建议将拒绝条件改为具体的变换合法性判断：能够保持两侧语义时正常转换；不能安全配对时，保留各分支原有顺序并使用正确的分支 mask，或明确诊断当前不支持这种分支合并。生成正确的访问序列之后，由 PTOAS 完成依赖分析和 mem_bar 插入。请同时更新“vecscope 内没有 ordering point”的注释、报错和测试预期。