feat(vmi): generalize VMIPredicateFold (R1–R6) + neutral reduce skip - #1194
Open
peanutchan wants to merge 2 commits into
Open
feat(vmi): generalize VMIPredicateFold (R1–R6) + neutral reduce skip#1194peanutchan wants to merge 2 commits into
peanutchan wants to merge 2 commits into
Conversation
Fold statically proven all-true/all-false vcmp into identity vsel and DCE dead pad work so frontends can always emit expert-pad masking when E is compile-time known (see hw-native-sys#1150). Co-authored-by: Cursor <cursoragent@cursor.com>
Expand pad vsel DCE into AllTrue demask, AllFalse hist/compute folds, neutral vmax/vadd peeps, and skip post-reduce neutral combines in VMIToVPTO. Document camodel A/B (CCE-parity R4; pad/peep micros) and honor PTO_FLAGS in ptodsl native_build for fold A/B. Co-authored-by: Cursor <cursoragent@cursor.com>
3 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
VMIPredicateFoldbeyond expert-padvcmp/vselinto R1–R6 (AllTrue demask, AllTrue/AllFalse pad, neutral splat peeps, AllFalsevdhist).vadd/vmax/vminaftervcadd/vcmax/vcmininVMIToVPTO(R4, CCE-parity on reduce dsl).vmi_predicate_fold_general.pto+ updated reduce FileChecks;PTO_FLAGShonored in ptodslnative_buildfor A/B.docs/vmi/vmi_predicate_fold_generalization_report.mdRelated: #580, earlier pad-only PR mouliangyu/PTOAS#581 (
feature-vmi).Camodel verification (Ascend950PR_9599 / CANN 9.1.0-beta.3)
R4 — neutral reduce skip (stock 0.53 → local 0.56)
R1–R3 / R5–R6 (fold ON vs
--disable-vmi-predicate-fold)RV_VSEL/RV_VCMPafter fold-infstore; partial keeps vselvmax(-inf,x)microVDUPS/VMAXvdhistmicroDHISTFull
topk_gatecamodel still blocked onwait_flaguncovered sections under local 0.56 — pad/peep covered by VL64 micros matching the same IR shapes.Test plan
pto-test-opt test/lit/vmi_new/vmi_predicate_fold_pad.pto -vmi-predicate-fold | FileCheckpto-test-opt test/lit/vmi_new/vmi_predicate_fold_general.pto -vmi-predicate-fold | FileCheckvmi_to_vpto_reduce_*)PTO_FLAGS=--disable-vmi-predicate-foldon pad/neutral microsVMI Predicate / Neutral-Element Fold — Implementation Report
Date: 2026-08-08
Scope: Generalize
VMIPredicateFold+ skip compiler-synthesized neutral reduce combines inVMIToVPTO.Binary-equal: folds are algebraically identity (AllTrue demask, AllFalse passthru,
max(x,-inf)=x,add(x,0)=x,vdhist(acc,*,F)=acc). Camodel ACL runs of the four R4 reduce kernels all PASSnp.allclosevs reference. Device NPU (torch.npu) was not available on this host.What shipped
create_mask(VL)on computeVMIPredicateFoldvsel→ true armVMIPredicateFold(existing + kept)vsel→ false armVMIPredicateFold(existing + kept)vadd/vmax/vmin(reduced, 0/-inf/+inf)aftervcadd/vcmax/vcminVMIToVPTOvmax(-inf,x)/vmin(+inf,x)/vadd(0,x)→xVMIPredicateFoldvdhist(acc,src,F) → accVMIPredicateFoldneed_pad)topk_gate_vmi_w128.pyShared analysis:
VMIMaskUtils(IntRange,MaskLattice, affine/vci ranges,mask_and/or/xor/not).Lit green:
vmi_predicate_fold_pad.pto,vmi_predicate_fold_general.pto, all updatedvmi_to_vpto_reduce_*.ptoFileChecks.Top-3 IR gains per rule
Gain ranking is static IR cost (ops removed on the hot path), ordered by corpus hit rate / documented cycle gaps where known.
R1 — AllTrue demask (largest corpus surface)
Hits: quant
create_mask(VL)on almost every VF iter; dsl ~362 full-mask sites.#1
vmulfull-mask (quant scale / swiglu-style)Before
After
#2
vaddfull-mask (dsl elementwise)Before
After
#3
vmaxfull-mask (amax / tree reduce leaves)Before
After
R2 — AllTrue pad
vsel(topk whenE % VL == 0or affine-proven)Documented in this week’s pad fold; always-emit pad + fold removes Python
need_padon w128.#1
pad_all_true_vci(E = VL = 64)Before:
vci+vbrc(E)+vcmp lt+vsel(m, score, -inf)After:
return %score(cmp/sel DCE’d)#2
pad_affine_multipass_all_true(multipass index affine)Before: loop body
vci/vadds/vcmp/vselper iterAfter: loop yields identity; pad ops removed
#3
vsel_same_arms(degenerate pad)Before:
vsel %m, %x, %xAfter:
return %xR3 — AllFalse pad
vsel(tail / empty expert chunk)#1
pad_all_false_vci(E = 0)Before:
vcmp+vsel(m, score, -inf)After:
return %-inf#2
pad_vadds_all_false(index past VL)Before:
vadds+vcmp+vselAfter:
return %-inf#3 topk rem chunk with proven empty window
Same rewrite as #1/#2 when range analysis proves all lanes fail
idx < E.R4 — Skip neutral reduce combine (RowMax / SoftmaxGrad / EuclideanNorm)
From
performance_analysis_0804.md: RowMax +33%, SoftmaxGrad +17%, EuclideanNorm +21% attributed in part to DSLvmax/vaddcleanup aftervcmax/vcadd.#1
vcaddf16 (SoftmaxGrad / EuclideanNorm style)Before (old lowering)
After
#2
vcaddi32 / f32 single-chunkSame shape as #1:
return %vcaddwith novadd(reduced, 0).#3
vcmaxmultichunk (RowMax across physical parts)Before:
vcmax× N →vmax(reduced, -inf)then inter-chunkvmaxAfter: seed acc from first
vcmax; only inter-chunkvmaxremains (no combine with-inf)R5 — Unrolled first-iter neutral splat
Hits: topk
acc_max = -inftree; quantacc = 0sum trees.#1
vmax(vbrc(-inf), x)→xBefore
After:
return %x#2
vadd(vbrc(0), x)→xBefore
After:
return %x#3
vmin(vbrc(+inf), x)→xBefore
After:
return %xR6 — AllFalse
vdhistHits: rem/empty histogram updates in group-count style kernels.
#1
vdhistwithcreate_mask(0)Before
After:
return %acc#2 / #3
Same rewrite whenever the mask lattice proves AllFalse (range-
vcmpormask_andwith AllFalse). Additional call sites collapse identically to #1; corpus density is lower than R1/R4.Camodel A/B (R4) — Ascend950PR_9599 / CANN 9.1.0-beta.3
Measured 2026-08-08 on
edgexpert-59a6(aarch64).Before = stock
ptoas0.53; After = localptoas0.56 with neutral-combine skip.All cases:
*_real_float_Rows_128_Cols_64.py, ACL path,PASS.RV_VMAXRV_VMINRV_VADDRV_VADDAfter R4, all four match the CCE
vf_real_execute_timefromperformance_analysis_0804.md(the prior +17…+33% DSL gaps attributed to neutralvadd/vmax/vmincleanup).Small-shape sanity (
RowMaxKernel.case1_float_Rows_2_Cols_64): vf 74→71, RVECEX 10→5,RV_VMAX2→0.Camodel A/B (R1–R3, R5–R6) — Ascend950PR_9599 / CANN 9.1.0-beta.3
Measured after rebuilding
libPTOASCompiler.so(fold was inpto-test-opt/.obut not linked into the CLI shared lib until 2026-08-08).Harness: local
ptoas0.56; Before =PTO_FLAGS=--disable-vmi-predicate-fold(now honored byptodslnative_build); After = fold on. Fulltopk_gatecamodel still blocked onwait_flaguncovered-section normalize under 0.56 — pad / peep coverage uses VL=64 micros with the samevcmp+vsel/vmax(-inf)/vdhistshapes.R1 — AllTrue demask (dsl real)
RV_VADD);RV_PSET3→1RV_VADD;RV_PSET3→1R1 alone is typically ADD-pipe / setup; largest corpus surface, smaller tick delta than R4.
R2 — AllTrue pad
vsel(E covers lanes)VCI+VCMP+VSELRV_VSEL/RV_VCMP(load→store)R3 — AllFalse pad
vselVDUPS(-inf)+store; 0 vsel/vcmpVCMP+VSEL(correct non-fold)R5 — Neutral
vmax(-inf,x)vmax(vbrc(-inf), x)VDUPS+VMAX; After: 0 both (identity)-infseed peep site;RV_PSET2→1 onlyR6 — AllFalse
vdhist → accvdhist(acc,src,create_mask(0))RV_DHIST; After: 0 (acc passthrough)DHISTboth sides);PSET3→2DHIST;PSET5→3 (R1-ish)Env notes
ninja PTOASCompilerafter editingVMIPredicateFold—python/pto/ptoas.socan be stale; CLI usespython/ptoas/mlir/_mlir_libs/libPTOASCompiler.so.PTO_FLAGSis forwarded byptodsl/_runtime/native_build.pyand included in the compile-config cache key.Verification checklist
vmi_predicate_fold_pad.ptovmi_predicate_fold_general.ptovmi_to_vpto_reduce_{addf,addf_f16,addi,minf,*_multichunk}.ptoVSEL/VCMPVDUPS+VMAXgone on microDHISTon AllFalse sitetopk_gatecamodelwait_flaguncovered tile section under local 0.56test_topk_gate.py)torch.npuhereFollow-ups (optional)
vdup(0/-inf)after R4 skips the combine (init still materialized by LowerUnified; already not in VF body for these cases).env_npu.sh, CANN 9.1b3).topk_gatecamodel (PTONormalizeUncoveredTileSections/wait_flag) for production E shapes.