Skip to content

feat(vpto): VPTOSoftPostUpdate Step 2: accumulator analysis + stride analysis unification#981

Open
yexiaosu wants to merge 8 commits into
hw-native-sys:mainfrom
yexiaosu:vmi-post-update-2
Open

feat(vpto): VPTOSoftPostUpdate Step 2: accumulator analysis + stride analysis unification#981
yexiaosu wants to merge 8 commits into
hw-native-sys:mainfrom
yexiaosu:vmi-post-update-2

Conversation

@yexiaosu

Copy link
Copy Markdown
Contributor

Implement the accumulator analysis path for VPTOSoftPostUpdate (design doc step 2), and refactor the overall stride computation architecture.

Changes

  • Unified stride analysis: replace the mutually-exclusive accumulator/delta two-path design with getStride, which independently analyzes each operand — accumulator (iter_arg) first, delta (IV/affine) fallback — then combines via computeFinalStride.
  • Recursive decomposeLinear: decompose yield expressions along addi/subi/muli/index_cast/addptr chains into blockArg * coeff + increment, requiring coeff == 1. Replaces single-level addi/addptr matching.
  • getIterArgIncrement look-through: trace operands back through addi/subi/addptr with loop-invariant offset and index_cast to find the underlying iter_arg, applying collected type casts to the increment.
  • Three-state return for getIterArgIncrement: nullopt (not iter_arg, try delta), Value() (iter_arg but decomposition failed, give up), Some(value) (success). Prevents treating unknown delta as zero.
  • Zero constant for unchanged iter_arg: yield = blockArg correctly returns delta=0 instead of signaling failure.

Tests

10 new lit tests covering: non-constant stride, base via addptr, both iter_args, cross-iter_arg stride, vsstb with iter_arg repeat_stride, multi-level yield chain, look-through offset, mixed accumulator+delta, unchanged iter_arg, non-additive yield (negative).

@yexiaosu

Copy link
Copy Markdown
Contributor Author

Review findings

1. Pass is a silent no-op on inferred vecscopes (blocking)

runOnOperation only does walk(pto::VecScopeOp), but the pass is registered before PTOInferVPTOVecScope (ptoas.cpp:2659) — and pto.vecscope is created by exactly that pass (PTOInferVPTOVecScope.cpp:432), the only such call site in the pipeline.

Every lit test and test/vpto/cases/*/kernel.pto writes pto.vecscope by hand, so the suite stays green. Removing the vecscope from soft_postupdate_delta-iv-vlds.pto and re-running: neither vlds nor vsts is rewritten. A pto.vecscope does appear in the output (line 18, added by the later pass), but the vlds keeps its non-post form.

Design doc 5.1 is self-contradictory on this: it places the pass before PTOInferVPTOVecScope while claiming the position guarantees "vecscope structure is intact".

Fix — two options. Either move the pass after PTOInferVPTOVecScope, or drop the vecscope restriction and walk the func body directly.

2. vsstb weight = 32 is only correct for 8-bit element types (blocking)

repeat_stride is measured in 32-byte blocks — the pseudocode at docs/isa/micro-isa/03-vector-load-store.md:503 is UB_block[base + repeat_stride + blk * block_stride], with the comment explicitly stating "one 32B block = 8 lanes". But pto.addptr's offset is in elements (it lowers to a GEP, whose element type is the pointer's element type; f32 is not reduced to i8).

So the weight in element units should be 32 / sizeof(T) — 8 for f32, 16 for f16/bf16, and 32 only for i8. Both design doc 4.2.1's table and the implementation hardcode 32.

Both directions measured wrong:

  • Starting repeat_stride at 4 blocks, the pass emits %2 = pto.addptr %1, %c128, i.e. ub_out + 128 f32 elements. Correct is 4 blocks × 32 B = 128 B = 32 f32 elements. Overshoots by 4×.
  • Advancing base by pto.addptr +32 f32 elements per iteration (= 128 B = 4 blocks) with repeat_stride fixed, the pass emits repeat_stride = 1. Correct is 4. Undershoots by 4×.

The existing soft_postupdate_delta-vsstb.pto happens to use rs_0 = 0 together with delta(base) = 0, where stride_new = (0 + w·Δrs) / w = Δrs — the weight cancels itself out, so the test cannot detect this.

3. Minor issues

  • materializeAtLoopEntry has two latent traps: it clones a def-chain out of the loop with no isPure check, and it unconditionally returns getResult(0) (hoistBefore handles the multi-result case correctly; this one does not).
  • materializeConst does no range check when building the i16 stride_new constant; a value exceeding i16 yields invalid IR.
  • Doc/implementation mismatch (not a bug). For weight ≠ 1, combineStride requires the entire total to be a compile-time constant. But by 4.2.5 check 3's own reasoning, only delta(base) needs to be a constant divisible by the weight — delta(strideOperand) could be a symbolic expression. This is a deliberate conservative choice; doc and implementation should just state the same thing.

yexiaosu added 8 commits July 24, 2026 17:47
Implement iter_arg-based accumulator analysis as a priority path before
delta analysis. This handles patterns where base or offset is an iter_arg
with explicit yield increment (arith.addi or pto.addptr), including
non-constant, loop-varying strides that delta analysis cannot handle.
Also fix computeInitialPtr to materialize base at loop entry for iter_arg
base pointers.
Major refactoring of VPTOSoftPostUpdate based on review:

- Replace two-path (accumulator/delta) mutual exclusion with unified
  getStride that independently analyzes each operand: accumulator
  (iter_arg) first, delta (IV/affine) fallback.
- Remove tryAccumulatorAnalysis; merge logic into processForOp.
- Recursive decomposeLinear for yield expressions: walk addi/subi/
  muli/index_cast/addptr chains instead of single-level matching.
- getIterArgIncrement look-through: trace through addi/subi/addptr
  with loop-invariant offset and index_cast to find underlying
  iter_arg.
- Distinguish "not iter_arg" (nullopt, try delta) from "iter_arg but
  decomposition failed" (Value(), give up) in getIterArgIncrement.
- Return zero constant for unchanged iter_arg (yield=blockArg) instead
  of treating it as failure.
- Require both operands to be successfully analyzed (!deltaBase ||
  !deltaOffset → skip) to prevent treating unknown delta as zero.
Analysis (decomposeLinear, getIterArgIncrement, computeDelta) is now pure and
returns a symbolic StrideExpr, cached per Value. Materialization runs once at a
single insertion point, after all legality checks pass.

Fixes three defects: a stride whose increment is defined after the candidate op
produced invalid IR or crashed; shared subexpressions were re-decomposed
exponentially, taking 58s for a 28-op loop body; and scaling a non-constant i16
repeat_stride by an index-typed weight built an invalid arith.muli.
Two defects in the loop rewrite, both dating from the initial pass.

Grouping: rewrites were grouped by (base, stride), but whether two ops can
share an iter_arg depends on them walking the same address sequence, and the
start address init_ptr is derived from base *and* strideOperand. Two ops on the
same base with the same stride but different offsets -- %ub[%iv] and
%ub[%iv + 64] -- were merged, and only the first one's init_ptr was kept, so
the second silently started 64 elements early. strideOperand is now part of the
key. Keying on the original operands rather than on init_ptr keeps the Value
identity comparison meaningful, since computeInitialPtr may materialize a
separate pto.addptr per candidate; this can split groups that could have been
merged, but never merges groups that must stay apart.

Traversal order: processVecScope reversed the walk result to get "inner to
outer", but Operation::walk already defaults to post-order, so the reverse
produced outer-to-inner. Rewriting a loop erases it, which also destroys every
loop nested inside it, leaving the already-collected inner ForOp handles
dangling -- a nested scf.for with a candidate in the outer loop aborted the
compiler. Dropping the reverse restores the intended order.
negative-non-additive-yield was vacuous: FileCheck matches substrings, so
`-> !pto.vreg<64xf32>` also matched the transformed form, and the CHECK-NOT
region was bounded by the following `CHECK: scf.yield`, never reaching the
yield's use of updated_base. Anchor on end-of-line with {{$}} and move
CHECK-NOT after the yield check. It now fails against transformed IR, where
previously it passed.

accum-multilevel-yield never reached the subtraction path in decomposeLinear:
the canonicalizer folded subi(addi(%off, 64), 4) into addi(%off, %c60) before
the pass ran. Subtract the induction variable instead so the chain cannot be
constant-folded, confirmed with --mlir-print-ir-before.

All 26 soft_postupdate lit tests pass.
The pass only walks pto.vecscope regions, but it was scheduled before
PTOInferVPTOVecScope -- the only pass that creates them. Kernels that do not
already carry a hand-written vecscope therefore reached the pass with nothing
to walk, and the rewrite was silently skipped.

Every lit test wrote pto.vecscope by hand, so the suite stayed green. Removing
the vecscope from soft_postupdate_delta-iv-vlds.pto reproduces it: the vlds
keeps its non-post form even though a vecscope appears in the output, added by
the later inference pass.

Moving the pass after the inference fixes it. The inference splices whole ops
into the scope, so scf.for bodies are not restructured and the "directly in
loop body" check is unaffected.

Adds soft_postupdate_inferred-vecscope.pto, whose input has no hand-written
vecscope; it fails on the previous ordering.
vsstb's repeat_stride counts 32-byte blocks while pto.addptr takes an
element offset; the pass bridged them with a constant 32, which is only
correct when sizeof(T) == 1. Derive the scale from the pointer element
type via an explicit per-op stride unit.

Rescaling applies only to delta(base), so Element-class ops (E == W) are
unchanged and a non-constant repeat_stride increment is now supported.
materializeAtLoopEntry cloned a def-chain in front of the loop without checking
the ops are pure, and always returned getResult(0), which is the wrong value for
a multi-result op. hoistBefore already handled both correctly.

Nothing checked that the computed stride fits the i16 repeat_stride operand.
arith.constant truncates silently, so a stride of 100000 became -31072 and the
store walked wrong addresses with no diagnostic. constantsFitType now checks
this before any IR is created, so the candidate is rejected instead.

Design doc 4.2.5 gains the range check, and check 3 now records that the
weight != 1 path requires the whole total_delta to be constant -- stricter than
necessary, but relaxing it needs structural divisibility on StrideExpr.
@yexiaosu
yexiaosu force-pushed the vmi-post-update-2 branch from 209aa49 to adad5b5 Compare July 24, 2026 10:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant