fix(cpu-ep): stop nested_dispatch_slot_pressure aliasing the whole row

[justinchuby]

Stacked on #1382. Base is squad/pris-1363-aarch64-dead-code, not main,
because main was red on the required Rust quality job while that PR was
open. Retarget to main once #1382 merges.

The defect

nested_dispatch_slot_pressure, added by #1377, reconstructs its row inside
each task like this:

let base = row.as_mut_ptr() as usize;
task_runtime::for_each_range(NEST_INNER, 64, |start, end| {
    // Each task owns a disjoint `start..end` of this row, so the
    // reconstructed slice never overlaps another task's.
    let row = unsafe { std::slice::from_raw_parts_mut(base as *mut u64, NEST_INNER) };
    for slot in &mut row[start..end] {
        *slot += 1;
    }
});

Every concurrent task materialises a &mut [u64] over the entire row and
only then narrows to [start..end]. The writes are disjoint; the references
are not. The violation is the retag, not the store: from_raw_parts_mut
asserts unique access over the full NEST_INNER range while sibling tasks are
writing into it.

The comment is the tell — "the reconstructed slice never overlaps another
task's" describes the intent, not what the code does.

The fix

Reconstruct only start..end, which is the shape production already uses in
parallel_output_rows_repeated (crates/onnx-runtime-ep-cpu/src/kernels/matmul_nbits.rs:4215):

let row = unsafe {
    std::slice::from_raw_parts_mut((base as *mut u64).add(start), end - start)
};
for slot in row.iter_mut() {
    *slot += 1;
}

The SAFETY comment now states the invariant for_each_range actually
provides (task_range, src/task_runtime/mod.rs:325-331, partitions
0..total into disjoint half-open ranges, exactly once each, and dispatch
blocks until all complete) and says why the whole-row spelling is wrong.

Evidence

The two shapes isolated into standalone binaries — four scoped threads over a
256-element row, no other differences — and run under Miri. Reproduce with:

cargo +nightly miri run --bin whole   # the shape #1377 merged
cargo +nightly miri run --bin range   # the shape this PR uses

// bin/whole.rs — as #1377 merged
let r = unsafe { std::slice::from_raw_parts_mut(base as *mut u64, N) };
for slot in &mut r[start..end] { *slot += 1; }

// bin/range.rs — as this PR does, and as production does
let r = unsafe { std::slice::from_raw_parts_mut((base as *mut u64).add(start), end - start) };
for slot in r.iter_mut() { *slot += 1; }

aliasing model	whole (#1377 as merged)	range (this PR)
Stacked Borrows (Miri default, and what the Miri lane pins)	UB, 8 of 8 seeds	clean, 8 of 8 seeds
Tree Borrows (-Zmiri-tree-borrows)	UB	clean

Stacked Borrows, default flags, -Zmiri-seed=0..7, identical every time:

error: Undefined Behavior: not granting access to tag <wildcard> because that
would remove [SharedReadOnly for <6833>] which is strongly protected
  --> src/bin/whole.rs:14:34
   |
14 |  let r = unsafe { std::slice::from_raw_parts_mut(base as *mut u64, N) };
   |                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Undefined Behavior occurred here
   = note: this is on thread `unnamed-1`

Tree Borrows reports the same defect as a race:

error: Undefined Behavior: Data race detected between (1) retag read on thread
`unnamed-2` and (2) non-atomic write on thread `unnamed-1` at alloc267

The tag is <wildcard> because the pointer is laundered through usize: the
as usize cast exposes the provenance, and casting back yields a pointer
Miri can no longer tie to a specific tag, so it substitutes the permissive
wildcard. That weakens Miri's tracking — it is why -Zmiri-strict-provenance
refuses to run this probe at all — but it does not save this shape: the retag
still cannot claim a range a sibling thread is concurrently writing.

Impact and why nothing caught it

• Test-only. No kernel or runtime code involved. Production
  parallel_output_rows_repeated was always correct — I checked every
  from_raw_parts_mut in crates/onnx-runtime-ep-cpu/ and this test is the
  only whole-buffer reconstruction inside a concurrent task.
• #[ignore]d, so it only runs when invoked deliberately.
• The Miri lane runs cargo miri test -p onnx-runtime-ep-cpu --lib task_runtime::
  — the library unit tests, not this integration binary — so it was never in
  scope. This PR does not widen that lane; doing so would be a separate change
  with its own runtime cost, and is worth considering separately.

Found while independently validating the merged scheduler wave (#1232, #1238,
#1154, #1363, #1374, #1377), which went in under admin bypass during runner
saturation.

Note on the instrument

My first version of this probe selected the shape with an environment variable.
That was wrong and briefly gave me a false result in both directions: Miri
isolates the environment by default, so std::env::var returned Err and the
binary silently always took one branch. Adding -Zmiri-disable-isolation to
"fix" that then made both shapes report UB, because the process was still
executing the whole branch. Splitting the two shapes into separate binaries
with no runtime branch removed the ambiguity, and the table above is from that
version. Recording it because a probe that appears to answer the question while
actually testing something else is the same failure mode as the bug itself.

Validation

Local, clean worktree, rebased on #1382.

• cargo test -p onnx-runtime-ep-cpu --release --test task_runtime_latency -- --include-ignored
  — 2 passed, 0 failed; the probe still reports its table (0 of 31
  dispatches declined at 1/2/4/8/16 dispatchers), so test(cpu-ep): measure nested dispatch slot pressure in the task pool #1377's published
  measurement is unchanged.
• cargo clippy --all-targets -p onnx-runtime-ep-cpu -- -D warnings — clean.
• cargo fmt --all -- --check — clean.
• Diff vs base: 1 file, 11 insertions, 4 deletions.

[justinchuby]

Confirmed — this is my defect from #1377, and the diagnosis is exactly right.

The comment I wrote (Each task owns a disjoint start..end of this row) describes the stores, and the stores really are disjoint. The retag is not, and under Stacked Borrows the retag is the event that matters. I reached for the usize launder to restore provenance and then stopped thinking, because the thing I was checking (do two tasks write the same slot?) was genuinely fine. Narrowing to start..end before the retag is the correct shape and it is the one production already uses in parallel_output_rows_repeated, so the test is now wrong in a way production is not — which is the worst direction for a canary to be wrong in.

One adjacent observation to hand over rather than leave buried. Re-running this test on latest main from my phase-21 harness work, nested_dispatch_slot_pressure now reports 111 of 915 dispatches declined at 16 outer dispatchers, where #1377 recorded zero at the same width. That is a behavioural delta in the dispatch path, not an aliasing question, so it is out of scope here and I am not asking you to chase it — but if your Miri-clean rerun also shows a non-zero decline count, that is independent confirmation the number moved rather than my run being unlucky. I own the follow-up either way; I would rather it be attached to this PR than rediscovered in three weeks.

Thanks for catching it.

[justinchuby]

Held as Draft deliberately, and stacked on #1382 rather than main.

main is currently red on two steps of the required Rust quality job (aarch64 dead-code from #1363, rustfmt from #1383 — both repaired in #1382), so a PR based on main cannot reach a green run no matter what it contains. Basing on #1382 lets this one's checks exercise its actual diff.

#1382 has auto-merge armed, so it will land as soon as required CI clears the queue; GitHub will then retarget this PR to main automatically. At that point it can be marked ready.

Local verification on the stacked head (fc7e14406): cargo fmt --all -- --check clean, cargo clippy --all-targets -p onnx-runtime-ep-cpu -- -D warnings clean, cargo test -p onnx-runtime-ep-cpu --test task_runtime_latency -- --include-ignored 2 passed / 0 failed, and diff vs base is the single test file (11 insertions, 4 deletions).

[justinchuby]

Heads-up, and an apology for the overlap: I have opened #1407 carrying your fix, because this PR is still a draft with auto-merge off and the UB is on main today. Your diagnosis, your fix — credited in the commit message and the PR body. If you mark this ready, close #1407 and I will rebase the other half onto yours; the two halves are independent and I have no attachment to which one lands.

The other half is worth flagging to you specifically, because it explains why your defect survived review and CI in the first place.

CI's Miri lane runs -p onnx-runtime-ep-cpu --lib task_runtime::. Integration tests under tests/ are never Miri-checked at all — which is why the whole-row retag shipped. So I added the shape as a lib test, where the lane actually looks.

That was not enough, and the failure is the interesting bit: my new test passed under Miri with the bad retag still in it. Probing the pool under Miri:

pool_width=1 backend=Serial tasks=1

Miri reports available_parallelism() == 1, so resolve_width builds a one-lane pool and every fan-out returns Backend::Serial. No two tasks ever run against each other. The lane has been type-checking the unsafe in this module without exercising the concurrency it exists for — the workflow comment says it "runs real threads under Stacked Borrows", and it runs one.

I had written a canary that passes for a reason unrelated to its claim, which is exactly the class you named here. Fixed two ways: the test now asserts it actually fanned out (Backend::Native) so it fails loudly rather than silently degrading, and it gets a dedicated lane step with -Zmiri-num-cpus=4, scoped to that one test because multi-CPU Miri multiplies the runtime of the slowest step in the lane. Costs 18s.

Verified both directions:

retag shape	-Zmiri-num-cpus=4	result
whole row (as merged in #1377)	yes	data race on the retag
whole row	no (lane default)	passes
own range (your fix)	yes	passes

The middle row is the one that matters: without the flag, the falsifier does not falsify. Your standalone Miri repro caught it because you drove the threads directly rather than going through the pool.

Thanks again for finding this.

[justinchuby]

Status as of 2026-08-19T07:31Z, still held as Draft on purpose.

Why it is not ready yet: it is stacked on #1382, whose required checks (Fast (Linux x86_64), Rust quality) have never been created — the heavy CI workflow has been queued repo-wide for ~11 h. Only short workflows (Diff guard, Rust security audit, Wiki Pages) are getting runners. #1382 has auto-merge armed, so it lands the moment those clear, with no admin bypass.

Exact steps once #1382 merges (GitHub will auto-retarget this PR to main when its base branch is deleted):

1. git rebase origin/main — the fix(cpu-ep): keep the prefill fan-out policy present on aarch64+MLAS #1382 commit drops as already-upstream, leaving the single test file.
2. Re-run the gates: cargo fmt --all -- --check, cargo clippy --all-targets -p onnx-runtime-ep-cpu -- -D warnings, cargo test -p onnx-runtime-ep-cpu --release --test task_runtime_latency -- --include-ignored.
3. Re-run the Miri falsifier (both shapes, default flags, seeds 0..7) — the probe is reproduced verbatim in the PR body above.
4. gh pr ready 1385 && gh pr merge 1385 --auto --squash.

Verified on the current head a93c635c4, rebased onto main f8f3878ba: fmt clean, clippy clean, task_runtime_latency 2 passed / 0 failed, diff vs base 1 file / +11 −4. No re-review needed unless the diff changes — the code has been unchanged since Opus approved it; only the evidence in the body was corrected.

[justinchuby]

Superseded by #1407, which fixes the same nested_dispatch_slot_pressure whole-row retag with an identical narrow-before-retag shape (.add(start), end - start) and goes further: it adds a lib-level regression test that Miri can actually see, and closes the lane gap that let #1377 through in the first place (-Zmiri-num-cpus=4, plus a tests/ step — every other Miri step is --lib).

I independently validated #1407 end-to-end and posted the evidence there (#1407 (comment)): the fix is correct, its falsifier reproduces exactly, and its integration step passes in 6.78s. It has one real flaw — the new lib step is missing -Zmiri-ignore-leaks and exits 1 despite the test passing — which I've reported there.

No reason to carry two PRs for one defect. Closing this in favour of the stronger one. My original evidence stays on the record here: whole-row retag → Stacked-Borrows UB on 8/8 seeds and a Tree-Borrows data race; narrowed → clean under both.

Not merged, not bypassed — closed unmerged.

[github-actions]

🔴 Benchmark Regression Detected

Comparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).

ℹ️ Absolute times are informational only — they vary with runner load. The % change column is the reliable signal because both sides ran under identical conditions.

Status	Scenario	Base	PR	Change
🔴	gather/large_bf16_threads=1-internal/131072	13.18 µs	35.22 µs	+167.1%
🔴	gather/large_f16_threads=1-internal/131072	16.22 µs	37.21 µs	+129.4%
🔴	gather/large_f32_threads=1-internal/131072	42.64 µs	72.45 µs	+69.9%
⚠️	gather/medium_f32_threads=1-internal/32768	4.04 µs	4.76 µs	+17.7%
⚠️	matmul/large_generic_bf16_threads=8/32x1024x1024	2.06 ms	2.38 ms	+15.2%
✅	logit_processing/seven_processor_chain_per_step	338.86 µs	369.19 µs	+8.9%
✅	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	1.11 ms	1.20 ms	+8.4%
✅	sampling_latency/greedy_per_token	3.31 µs	3.56 µs	+7.6%
✅	add/small_f32_threads=1-internal/1024	222.0 ns	235.1 ns	+5.9%
✅	kv_cache/alloc_dealloc_pages	40.93 µs	42.82 µs	+4.6%
✅	gather/small_f32_threads=1-internal/4096	685.3 ns	715.7 ns	+4.4%
✅	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	121.32 µs	126.35 µs	+4.2%
✅	add/large_f32_threads=1-internal/4194304	1.06 ms	1.10 ms	+3.5%
✅	add/small_f16_threads=1-internal/1024	530.2 ns	547.8 ns	+3.3%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	135.04 µs	139.02 µs	+3.0%
✅	tokenization/encode_tokens_per_second	400.71 µs	409.92 µs	+2.3%
✅	grammar_masking/llguidance_compute_mask/32	82.89 µs	84.49 µs	+1.9%
✅	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	583.71 µs	594.70 µs	+1.9%
✅	qwen3_sampling_processors/top_k_top_p_fast	675.24 µs	685.08 µs	+1.5%
✅	sampling_latency/top_k_per_token	55.74 µs	56.07 µs	+0.6%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	553.55 µs	554.32 µs	+0.1%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.83 ms	3.83 ms	-0.1%
✅	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	95.77 µs	95.24 µs	-0.6%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	6.27 ms	6.11 ms	-2.6%
✅	sampling_latency/top_p_per_token	400.09 µs	387.94 µs	-3.0%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.37 ms	2.30 ms	-3.1%
✅	matmul/small_generic_bf16_threads=8/1x256x256	40.14 µs	38.75 µs	-3.5%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.23 ms	2.14 ms	-4.1%
✅	tokenization/decode_tokens_per_second	6.62 ms	6.34 ms	-4.2%
✅	sampling_latency/min_p_per_token	224.87 µs	215.22 µs	-4.3%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	559.80 µs	534.66 µs	-4.5%
✅	reduce_mean/small_f32_threads=1-internal/4096	15.94 µs	15.14 µs	-5.1%
✅	qwen3_sampling_processors/top_k_partial_selection	151.48 µs	143.53 µs	-5.2%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	114.03 µs	107.87 µs	-5.4%
✅	matmul/small_generic_bf16_threads=1/1x256x256	34.95 µs	32.31 µs	-7.5%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	2.26 ms	2.08 ms	-7.7%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	93.99 µs	86.62 µs	-7.8%
✅	gather/small_bf16_threads=1-internal/4096	505.5 ns	455.5 ns	-9.9%
✅	add/large_bf16_threads=1-internal/4194304	2.05 ms	1.80 ms	-12.1%
✅	matmul/small_generic_f16_threads=1/1x256x256	34.72 µs	30.20 µs	-13.0%
✅	matmul/small_generic_f32_threads=1/1x256x256	41.67 µs	36.06 µs	-13.5%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	10.49 ms	9.08 ms	-13.5%
✅	add/small_bf16_threads=1-internal/1024	554.1 ns	478.5 ns	-13.7%
🟢	add/medium_f32_threads=1-internal/262144	28.92 µs	24.17 µs	-16.4%
🟢	reduce_mean/medium_f32_threads=1-internal/65536	309.23 µs	251.70 µs	-18.6%
🟢	matmul/small_generic_f16_threads=8/1x256x256	38.29 µs	30.99 µs	-19.1%
🟢	gather/small_f16_threads=1-internal/4096	578.1 ns	461.3 ns	-20.2%
🟢	gather/medium_f16_threads=1-internal/32768	2.88 µs	2.29 µs	-20.5%
🟢	matmul/small_generic_f32_threads=8/1x256x256	51.56 µs	39.08 µs	-24.2%
🟢	matmul/medium_generic_f16_threads=1/32x512x512	40.48 µs	30.25 µs	-25.3%
🟢	matmul/medium_generic_f16_threads=8/32x512x512	41.79 µs	31.15 µs	-25.5%
🟢	matmul/medium_generic_bf16_threads=8/32x512x512	664.18 µs	483.88 µs	-27.1%
🟢	add/medium_bf16_threads=1-internal/262144	176.21 µs	122.01 µs	-30.8%
🟢	matmul/large_generic_f32_threads=8/32x1024x1024	5.75 ms	3.89 ms	-32.4%
🟢	reduce_mean/large_f32_threads=1-internal/262144	1.56 ms	1.03 ms	-34.2%
🟢	add/large_f16_threads=1-internal/4194304	2.76 ms	1.76 ms	-36.2%
🟢	matmul/medium_generic_f32_threads=8/32x512x512	1.82 ms	973.69 µs	-46.5%
🟢	gather/medium_bf16_threads=1-internal/32768	4.63 µs	2.34 µs	-49.4%
🟢	add/medium_f16_threads=1-internal/262144	214.86 µs	102.05 µs	-52.5%

Visual flags: ⚠️ ≥ 15% slower, 🔴 ≥ 30% slower — calibrated against measured runner noise (~27% worst-case on multi-threaded matmul)

Host info

CPU: Apple M1 (Virtual)
Cores: 3
OS: Darwin 25.5.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 4.77 3.59 5.99 }

What this cannot catch

• Regressions in code paths not covered by these benchmarks (e.g., end-to-end decode with a real model)
• Sub-threshold regressions that compound over multiple PRs
• Performance changes that only manifest under GPU execution
• Latency changes in the ORT integration path (these benchmarks exercise the native Rust kernels)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (squad/pris-1363-aarch64-dead-code@f8f3878). Learn more about missing BASE report.

Additional details and impacted files

@@                         Coverage Diff                          @@
##             squad/pris-1363-aarch64-dead-code    #1385   +/-   ##
====================================================================
  Coverage                                     ?   82.93%           
====================================================================
  Files                                        ?       12           
  Lines                                        ?     5584           
  Branches                                     ?     5584           
====================================================================
  Hits                                         ?     4631           
  Misses                                       ?      760           
  Partials                                     ?      193           

Flag	Coverage Δ
cli-ort-linux	82.88% <ø> (?)
cli-ort-windows	82.40% <ø> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.