Skip to content

perf(plugin): cut the per-Run dispatch floor for single-node subgraphs - #1244

Merged
justinchuby merged 3 commits into
mainfrom
squad/resch-small-node-dispatch
Aug 19, 2026
Merged

justinchuby merged 3 commits into
mainfrom
squad/resch-small-node-dispatch

Conversation

@justinchuby

@justinchuby justinchuby commented Aug 18, 2026 •

Copy link
Copy Markdown
Owner

What this is

The per-Run cost of dispatching a single-node fused subgraph through the plugin path, with the kernel removed from the picture.

On a 4 Ki f32 elementwise node our arm spends ~3.5 µs per Run where plain ORT spends ~2.6 µs, and the kernels themselves account for well under a microsecond of that — the same kernel on a 1 Mi tensor runs at 0.51x–1.03x of ORT. The difference is fixed dispatch overhead, and it is why every cheap elementwise op sits at 1.1x–1.5x of ORT at 4 Ki while the identical kernel wins at 1 Mi.

This PR removes four pieces of that overhead. No kernel is touched, no numerics change, and nothing about assignment or execution ownership changes: the CPU EP still executes every node it claims, locally, with no ORT CPU fallback anywhere.

The four cuts

1. device_mem_info is no longer resolved eagerly (6 ORT FFI calls per Run).

compute_execute resolved scratch_mem_info at the top of every call:

let scratch_mem_info =
    unsafe { device_mem_info(api_ref, kernel_context, exported.device_staging.as_ref()) };

device_mem_info calls KernelContext_GetInputCount, then for each input KernelContext_GetInput + GetTensorMemoryInfo + GetMemoryInfoDeviceType, then falls through to two more calls for input 0. For a one-input node that is six FFI calls.

Who reads it? Two consumers. intermediate_scratch — routed multi-node path only. And PlacementSources::subgraph_fallback, which operand_mem_info reads only when the node binds no ORT operands at all, and which prepare_workspace only reaches past its zero-byte and lifetime gates. An elementwise node has operands and needs no workspace, so it reached neither.

Now the routed path resolves it once per Run (unchanged) and passes SubgraphFallback::Resolved; the single-node path passes SubgraphFallback::Deferred(staging) and operand_mem_info calls device_mem_info itself, with the same arguments, if it ever actually needs the answer. Deferring is safe: the function only reads input memory info, so running it after outputs are allocated cannot see a different answer.

2. allocate_output takes want_mem_info (1 ORT FFI call per output per Run).

OwnedOutput::mem_info has exactly one consumer in the tree — stage_host_boundary_inputs, called under if let Some(staging) = exported.device_staging.as_ref(). A host EP has no staging context, so every Run made a GetTensorMemoryInfo call per output and dropped the result. Both call sites now pass exported.device_staging.is_some(), so a device EP is bit-for-bit unchanged and a host EP stops making the call.

3. staging_log(&format!(..)) → staging_log!(..) (one String per Run).

staging_log checks ONNX_GENAI_PLUGIN_TRANSFER_TRACE inside the function, so the format! was evaluated whether or not the trace was on — and one of these sites is at the top of compute_execute, formatting three fields into a heap String on every dispatch. The macro checks the gate first. All 12 sites converted, so the footgun is gone rather than papered over at one site.

4. Absent-output bookkeeping is built only when there are absent outputs (5 allocations per Run).

let absent_shapes: Vec<Vec<usize>> = output_shapes.clone();
let absent_strides_storage: Vec<Vec<i64>> =
    absent_shapes.iter().map(|s| contiguous_strides(s)).collect();
let mut ort_views: Vec<TensorMut<'_>> =
    owned_outputs.iter_mut().map(|o| o.view_mut()).collect();
let mut ort_view_iter = ort_views.drain(..);

That storage exists solely to back the TensorMuts of absent output slots. A node with no absent outputs — every elementwise op — cloned every output shape and built a stride vector per output to lend them to nobody. It is now conditional on !absent_bufs.is_empty(). The ort_views Vec was collected and immediately drained; the iterator is taken directly from owned_outputs.iter_mut().

Related: placement operands are now described by an OrtOperands enum, so the single-node path lends entry.input_slots (&[Option<usize>], flattened lazily) instead of collecting a fresh Vec<usize> per call for a consumer that almost never runs. The routed path passes its already-resolved slice.

Per Run on a 1-in/1-out elementwise node this is 7 fewer ORT FFI calls (16 → 9) and 7 fewer heap allocations.

A/B, one thread

taskset -c 8-15, this commit vs its parent, interleaved (A,B,A,B,A,B with a rebuild before each arm, so drift hits both arms equally), 400 iterations after 50 warmup, three rounds. Ratio is ours/ORT, lower is better; the p50 column is the median of the three rounds' p50 ratios, p90 likewise.

case before p50 after p50 before p90 after p90
sqrt_f32_4k 1.132 1.026 1.153 1.033
sigmoid_f32_4k 1.484 1.346 1.517 1.362
tanh_f32_4k 1.534 1.406 1.595 1.434
erf_f32_4k 1.760 1.510 1.779 1.512

Best-of-three (the contention-robust statistic on this shared box) agrees: sqrt 1.130 → 1.023, sigmoid 1.481 → 1.344, tanh 1.503 → 1.387, erf 1.570 → 1.480. 12 of 12 arm-pairs favour the change; there is no round in which any case regressed.

In absolute terms tanh_f32_4k goes 0.0046 ms → 0.0042 ms, i.e. ~0.4 µs off a ~0.9 µs gap.

A/B, threaded

Same protocol, taskset -c 0-15, two rounds, NXRT_MM_BENCH_THREADS = ONNX_GENAI_MLAS_THREADPOOL_THREADS = RAYON_NUM_THREADS.

case 4t before 4t after 16t before 16t after
sqrt_f32_4k 1.257 1.129 1.262 1.156
sigmoid_f32_4k 1.520 1.355 1.499 1.340
tanh_f32_4k 1.551 1.418 1.548 1.455
erf_f32_4k 1.776 1.671 1.777 1.624

The overhead is per call, not per element or per worker, so the gain is the same absolute number of microseconds at every thread count.

Drift control: the 1 Mi grid is unchanged

Same protocol, _f32_1m, three interleaved rounds, one thread. A large tensor amortises the per-call cost away, so these must not move — and they don't:

case before after
relu_f32_1m 1.032 1.034
exp_f32_1m 1.021 1.022
sigmoid_f32_1m 1.066 1.063
tanh_f32_1m 1.117 1.115
gelu_tanh_f32_1m 1.243 1.241
gelu_exact_f32_1m 1.424 1.420
fastgelu_f32_1m 1.239 1.224
erf_f32_1m 1.461 1.439
quickgelu_f32_1m 0.809 0.811
sqrt_f32_1m 0.514 0.511

Ten of ten within ±0.5 %, which is this box's noise floor. That is the shape of a per-call fix.

Correctness

New test, with a verified falsifier. output_memory_info_is_queried_only_when_the_caller_asked_for_it drives allocate_output against a hand-built OrtApi whose GetTensorMemoryInfo counts its calls, and asserts 0 calls for want_mem_info: false and 1 for true. Falsifier: making allocate_output ignore the flag fails it with left: 1, right: 0 — checked by breaking the code, not by inspection.

Behaviour preserved, argued per cut. (1) device_mem_info is called with identical arguments, only later and only when read; it reads inputs, which output allocation cannot change. (2) The gate on the memory-info query is the same condition as the gate on its only consumer. (3) The macro's only difference is when the format! runs. (4) The absent storage is only ever indexed for absent slots.

Suites. -p onnx-runtime-ep-plugin: 247 unit tests pass (246 before, +1 new). -p onnx-runtime-ep-cpu-plugin with NXRT_REQUIRE_ORT_TESTS=1: the full e2e suite passes, including all 54 plugin_ort_e2e cases — the routed multi-node fixtures (conformance_chain_add_mul, ..._repeated_runs_do_not_leak_stale_intermediates, conformance_mixed_partition) exercise the SubgraphFallback::Resolved arm, and every_assigned_node_is_also_executed_by_this_ep passes, so every node this EP claims is still executed here with ORT CPU fallback disabled. cargo fmt --all and cargo clippy --release --all-targets -p onnx-runtime-ep-cpu -p onnx-runtime-ep-cpu-plugin -p onnx-runtime-ep-plugin are clean.

Build identity. Pure native CPU EP: no MLAS at runtime, no ORT CPU EP fallback, no new dependency. AVX2/FMA host (avx2 fma f16c, no AVX-512), so ORT/MLAS and we are on the same instruction footing.

What is left

The remaining ~0.5 µs is, as far as I can attribute it without CPU counters (perf_event_paranoid is 4 on this box, so perf record is not available and everything here is A/B attribution):

  • ~9 ORT FFI calls that are genuinely needed. read_inputs costs 7 per input — KernelContext_GetInput, GetTensorTypeAndShape (which allocates an ORT object we then release), GetTensorElementType, GetDimensionsCount, GetDimensions, ReleaseTensorTypeAndShapeInfo, GetTensorData — and there is no cheaper spelling in the stable C API. ORT's own CPU kernels reach the same data through OpKernelContext with no FFI at all, which is a structural part of what a plugin EP pays.
  • ~7 remaining allocations: OwnedInput's shape and strides per input, kernel_inputs, infer_shapes's Vec<Vec<usize>>, slot_map, output_views, prepare_workspace's metadata, allocate_output's dims, and the Box<HostPool> in host_pool::install. Each is worth ~25–35 ns. Removing them needs either an inline-capacity vector type or per-session caching of the parts that cannot change between Runs; both are worth doing and neither belongs in this PR.

I deliberately did not touch host_pool::install — the per-call Box is one allocation, and that file is @sebastian's 16-thread scheduling work; a change there should come from him or after his PRs land.


Refreshed against main (2026-08-18)

The branch was behind main and its whole red CI wall came from that, not from
this change: crates/onnx-runtime-session/src/executor/mod.rs:175 failed
-D dead-code on current stable, which ca32b3adf ("fix(ci): unbreak the Rust
quality lane on current stable", #1239) fixed on main after this branch forked.
origin/main (c55a3fab3) is merged in — no rebase, no force-push — and the
diff this PR owns is unchanged at 2 files, +285/-79.

Revalidated on the merge commit, AVX2/FMA host, no AVX-512:

  • cargo test --release -p onnx-runtime-ep-plugin — 247 passed, 0 failed.
  • NXRT_REQUIRE_ORT_TESTS=1 cargo test --release -p onnx-runtime-ep-cpu-plugin
    — every suite green, including all 55 plugin_ort_e2e cases. The ones
    that matter to this change all pass on the merge:
    every_assigned_node_is_also_executed_by_this_ep,
    no_supported_node_is_ever_left_to_the_ort_cpu_ep,
    no_matmul_family_node_escapes_to_the_ort_cpu_ep, and
    every_fixture_loads_with_cpu_fallback_disabled — so assigned still equals
    executed with ORT CPU fallback off.
  • cargo fmt clean for the crates this PR touches. The one cargo fmt --all
    hunk on this tree is in onnx-runtime-ep-cuda/src/kernels/standard_attention.rs,
    which arrived from main untouched by this PR and is a local rustfmt-version
    difference, not a branch defect.

Independent review

Reviewed by Claude Opus 4.8, read-only, with the four cuts and the
absent-slot history stated as the priority list. Verdict APPROVE, no
blockers. It independently confirmed:

  • SlotKind::Absent(idx) is pushed into slot_map only in the same branch that
    pushes into absent_bufs, so has_absent == false implies slot_map holds no
    Absent and the skipped storage is never indexed; and the surviving indices
    are the same full-slot indices as before.
  • ort_view_iter from owned_outputs.iter_mut() has the identical borrow
    structure as the old collect() + drain(..), so nothing borrows a temporary.
  • OrtOperands::Slots(..).indices() yields the same elements in the same order
    as the removed .iter().flatten().copied().collect(), None skipped.
  • operand_mem_info is the only reader of the deferred value, only reachable
    when the node binds zero ORT inputs, which makes the timing of the deferred
    device_mem_info moot rather than merely argued.
  • Both allocate_output call sites gate on the same predicate as the only
    reader of OwnedOutput::mem_info, so a device EP is unchanged.

Its one correction is applied above: the prose said 14 staging_log sites; the
real count in origin/main is 12, and all 12 are converted.


Refreshed against main @ 6a855d5e0, and measured as a stack

origin/main moved a long way while this sat in the CI queue (#1346, #1352 and
#1361 on the quality lane; #1154, #1232, #1238 on the CPU side). Merged in
normally — no rebase — and re-measured from scratch against the new baseline.

Production pure-native A/B, plain ORT as the control arm. No MLAS, no ORT CPU
fallback, no deferral. taskset -c 8-15, one thread, 400 iterations, five
interleaved rounds out of two worktrees, started only once cores 8-15 were

=93% idle. Ratio is ours/ORT, lower is better. before is main at
6a855d5e0; after is #1244 + #1246 together, since #1246 is stacked on #1244
and the pair is what a user gets.

case ratio p50 main ratio p50 stack Δ ratio p90 main ratio p90 stack ours us ORT drift rounds won
thresholdedrelu_f32_4k 1.512 1.216 -19.6% 1.519 1.216 3.6 → 2.8 -4.2% 5/5
tanh_f32_4k 1.511 1.279 -15.4% 1.513 1.284 4.6 → 3.9 +0.0% 5/5
sigmoid_f32_4k 1.475 1.248 -15.4% 1.484 1.256 4.7 → 4.0 +0.0% 5/5
erf_f32_4k 1.470 1.363 -7.3% 1.489 1.364 7.5 → 7.0 +0.0% 5/5
hardsigmoid_f32_4k 1.416 1.127 -20.4% 1.426 1.133 3.5 → 2.8 +0.0% 5/5
leakyrelu_f32_4k 1.361 1.094 -19.6% 1.371 1.104 3.5 → 2.8 +0.0% 5/5
sqrt_f32_4k 1.141 0.946 -17.1% 1.150 0.955 4.0 → 3.4 -2.8% 5/5
log_f32_4k 0.767 0.698 -9.0% 0.776 0.703 7.9 → 7.2 +0.0% 5/5
selu_f32_4k 0.505 0.440 -12.9% 0.512 0.444 5.6 → 4.9 +0.0% 5/5
elu_f32_4k 0.480 0.416 -13.3% 0.488 0.421 5.3 → 4.6 +0.0% 5/5
celu_f32_4k 0.470 0.411 -12.6% 0.478 0.418 5.7 → 5.0 +0.0% 5/5
mish_f32_4k 0.276 0.264 -4.3% 0.278 0.268 17.3 → 16.6 +0.0% 5/5

Every case, every round. The two rows with a moving control (sqrt -2.8%,
thresholdedrelu -4.2%) are reported rather than dropped; both won 5/5 anyway
and their absolute time fell by the same ~0.7 us as everything else.

That constant ~0.7 us is the point. It is not proportional to tensor size — the
same absolute amount comes off hardsigmoid (3.5 -> 2.8 us) as off mish
(17.3 -> 16.6 us) — which is what a fixed per-Run cost looks like when you
remove some of it. It moves the cheap ops the most because they had the least
to hide it behind, and sqrt crosses from 1.141 to 0.946, from a loss to a
win.

Where the remaining time goes

Measured directly, by instrumenting compute_execute segment by segment on top
of this stack (temporary probe, not committed; perf is unavailable on this
host — perf_event_paranoid=4). Per Run, one-in/one-out elementwise node,
4096 f32, microseconds:

segment us note
KernelContext_GetOutput 0.35 ORT's own API — ours to call, not to optimise
read_inputs 0.15 4 ORT FFI calls, already one shape call after #1246
rest of allocate_output 0.13 GetTensorMutableData + strides
prepare_workspace 0.09 metadata vector + plan-cache lookup, for a kernel needing 0 bytes
host_pool::install 0.05 @sebastian's, not touched
infer_shapes 0.05
kernel_inputs 0.04
output_views 0.04

Non-kernel node cost is ~1.25 us and near-constant across all twelve
operators
(0.28 to 15.2 us of kernel time), which is the direct confirmation
that small-node ratios on this EP are dispatch-bound rather than kernel-bound.
There is no single large item left — the biggest, KernelContext_GetOutput, is
ORT's. The rest is a long tail of 0.04-0.15 us items, which is what #1358
(InlineVec) starts on.

And nothing breaks at 1 Mi

Same harness, 1048576 elements, 120 iterations, 3 rounds. A fixed per-Run
cost should be invisible here, and it is:

case ratio p50 main ratio p50 stack ours us ORT drift
celu_f32_1m 0.141 0.140 382.1 → 377.9 -0.0%
elu_f32_1m 0.138 0.136 347.7 → 343.6 +0.1%
erf_f32_1m 0.671 0.670 595.2 → 594.6 +0.1%
exp_f32_1m 0.606 0.590 247.0 → 240.5 +0.1%
fastgelu_f32_1m 0.643 0.648 415.3 → 413.2 -1.7%
gelu_exact_f32_1m 0.572 0.592 706.1 → 710.6 -2.8% ⚠
gelu_tanh_f32_1m 0.657 0.651 414.8 → 410.5 +0.3%
hardsigmoid_f32_1m 0.376 0.372 88.0 → 69.6 -0.9%
leakyrelu_f32_1m 0.421 0.430 90.1 → 87.3 -2.2% ⚠
log_f32_1m 0.270 0.274 616.4 → 612.4 +1.7%
mish_f32_1m 0.105 0.105 1626.0 → 1626.0 -0.2%
quickgelu_f32_1m 0.455 0.453 321.5 → 320.0 -0.5%
relu_f32_1m 1.034 1.022 131.7 → 130.2 +0.0%
selu_f32_1m 0.147 0.148 374.1 → 371.8 -1.5%
sigmoid_f32_1m 0.471 0.606 231.6 → 229.0 -23.4% ⚠
sqrt_f32_1m 0.314 0.302 148.1 → 143.9 +1.1%
tanh_f32_1m 0.644 0.631 226.2 → 221.9 -0.2%
thresholdedrelu_f32_1m 0.500 0.486 70.9 → 68.6 -0.3%

Flat, as predicted — 0.7 us against 70-1626 us of work. Absolute time is equal
or better in 16 of 18 cases. The two ⚠ rows had the control move more than the
effect: sigmoid is unusable (ORT itself moved -23.4%; our own absolute went
231.6 -> 229.0 us), and gelu_exact's +0.6% absolute sits inside its -2.8%
control. Reported rather than dropped.

This is the coverage claim for the change: it buys ~0.7 us at every size, which
is 20% of a small node and nothing at all of a large one, and it costs nothing
anywhere.

Every `Run` of a single-node fused subgraph paid a fixed cost that has
nothing to do with the kernel: six ORT FFI calls to resolve a memory info
nobody read, one more per output for the same reason, a `String` formatted
for a trace nobody enabled, and five heap allocations lent to nobody. On a
4 Ki f32 elementwise node that overhead is most of the gap to ORT.

Four cuts, all behaviour-preserving:

* `device_mem_info` is no longer resolved eagerly. Its single-node
  consumer is `operand_mem_info`'s no-ORT-operands fallback, which a node
  with operands never reaches, so `PlacementSources` now carries
  `SubgraphFallback::Deferred` and resolves on demand. The routed path
  still resolves once per `Run` and passes `Resolved`. That is
  `KernelContext_GetInputCount` + three calls per input + two more, gone
  from every elementwise dispatch.
* `allocate_output` takes `want_mem_info`. `OwnedOutput::mem_info` is read
  only by `stage_host_boundary_inputs`, under `device_staging.is_some()`,
  so a host EP was making one `GetTensorMemoryInfo` call per output per
  `Run` and dropping the answer.
* `staging_log(&format!(..))` became `staging_log!(..)`, which checks the
  trace gate before formatting. One of those sites is on the dispatch
  path and built a `String` on every call.
* The absent-output shape/stride storage is built only when there are
  absent outputs (it exists solely to back their `TensorMut`s), and the
  ORT output views are taken from an iterator instead of a `Vec` that was
  collected and immediately drained.

Placement operands are described by `OrtOperands`, which lets the
single-node path lend `entry.input_slots` directly instead of flattening
it into a fresh `Vec` per call for a consumer that almost never runs.

Nine fewer ORT FFI calls and seven fewer allocations per `Run` on a
one-in, one-out elementwise node.

A/B against the parent commit, one thread, `taskset -c 8-15`, three
interleaved rounds of 400 iterations, ratio = ours/ORT (lower is better),
each cell the median of the three rounds:

| case | before p50 | after p50 | before p90 | after p90 |
|---|---|---|---|---|
| `sqrt_f32_4k` | 1.132 | 1.026 | 1.153 | 1.033 |
| `sigmoid_f32_4k` | 1.484 | 1.346 | 1.517 | 1.362 |
| `tanh_f32_4k` | 1.534 | 1.406 | 1.595 | 1.434 |
| `erf_f32_4k` | 1.760 | 1.510 | 1.779 | 1.512 |

The same absolute microseconds come off at 4 and 16 threads. The 1 Mi
grid is the drift control and is unchanged across all ten cases (`relu`
1.032 -> 1.034, `sqrt` 0.514 -> 0.511, `quickgelu` 0.809 -> 0.811): the
cost is per call, so it disappears into a large tensor.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@justinchuby
justinchuby force-pushed the squad/resch-small-node-dispatch branch from 55aed5e to 66c9835 Compare August 18, 2026 12:42
@justinchuby
justinchuby marked this pull request as ready for review August 18, 2026 12:42
@justinchuby
justinchuby enabled auto-merge (squash) August 18, 2026 12:42
@github-actions

github-actions Bot commented Aug 18, 2026 •

Copy link
Copy Markdown

🔴 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
🔴 block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024 58.06 µs 123.63 µs +112.9%
🔴 block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024 762.85 µs 1.58 ms +106.9%
🔴 matmul/small_generic_f32_threads=8/1x256x256 35.22 µs 60.51 µs +71.8%
🔴 gather/large_f16_threads=1-internal/131072 10.85 µs 15.46 µs +42.4%
🔴 matmul/medium_generic_f32_threads=8/32x512x512 1.20 ms 1.68 ms +39.2%
🔴 gather/large_bf16_threads=1-internal/131072 11.08 µs 14.52 µs +31.0%
🔴 matmul/small_generic_f16_threads=8/1x256x256 30.71 µs 40.20 µs +30.9%
⚠️ matmul/large_generic_bf16_threads=8/32x1024x1024 1.75 ms 2.25 ms +28.7%
⚠️ matmul/small_generic_f32_threads=1/1x256x256 36.33 µs 46.39 µs +27.7%
⚠️ matmul/small_generic_f16_threads=1/1x256x256 31.15 µs 39.12 µs +25.6%
⚠️ gather/medium_f32_threads=1-internal/32768 3.76 µs 4.71 µs +25.3%
⚠️ matmul/large_generic_f32_threads=8/32x1024x1024 4.90 ms 5.76 ms +17.7%
✅ add/large_f32_threads=1-internal/4194304 678.01 µs 778.08 µs +14.8%
✅ gather/medium_bf16_threads=1-internal/32768 2.58 µs 2.92 µs +13.1%
✅ matmul/medium_generic_f32_threads=1/32x512x512 2.32 ms 2.63 ms +13.1%
✅ add/large_bf16_threads=1-internal/4194304 1.65 ms 1.87 ms +13.0%
✅ add/large_f16_threads=1-internal/4194304 1.73 ms 1.95 ms +12.8%
✅ gather/small_bf16_threads=1-internal/4096 490.0 ns 548.4 ns +11.9%
✅ gather/medium_f16_threads=1-internal/32768 2.49 µs 2.78 µs +11.5%
✅ reduce_mean/medium_f32_threads=1-internal/65536 249.70 µs 275.67 µs +10.4%
✅ block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1 93.12 µs 101.43 µs +8.9%
✅ gather/small_f16_threads=1-internal/4096 477.5 ns 517.1 ns +8.3%
✅ gather/small_f32_threads=1-internal/4096 678.1 ns 729.3 ns +7.5%
✅ reduce_mean/small_f32_threads=1-internal/4096 16.53 µs 17.77 µs +7.5%
✅ matmul/small_generic_bf16_threads=1/1x256x256 36.65 µs 39.03 µs +6.5%
✅ reduce_mean/large_f32_threads=1-internal/262144 1.02 ms 1.09 ms +6.4%
✅ block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1 458.56 µs 483.37 µs +5.4%
✅ matmul/large_generic_f32_threads=1/32x1024x1024 11.00 ms 11.55 ms +5.0%
✅ matmul/small_generic_bf16_threads=8/1x256x256 39.13 µs 40.86 µs +4.4%
✅ matmul/large_generic_bf16_threads=1/32x1024x1024 2.45 ms 2.55 ms +3.9%
✅ add/small_bf16_threads=1-internal/1024 501.2 ns 511.1 ns +2.0%
✅ matmul/large_generic_f16_threads=1/32x1024x1024 91.19 µs 91.85 µs +0.7%
✅ add/medium_f32_threads=1-internal/262144 28.06 µs 28.17 µs +0.4%
✅ sampling_latency/top_p_per_token 443.53 µs 441.33 µs -0.5%
✅ add/medium_bf16_threads=1-internal/262144 121.13 µs 119.07 µs -1.7%
✅ matmul/medium_generic_bf16_threads=8/32x512x512 687.05 µs 660.29 µs -3.9%
✅ qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline 3.90 ms 3.68 ms -5.7%
✅ matmul/medium_generic_f16_threads=8/32x512x512 51.92 µs 48.91 µs -5.8%
✅ gather/large_f32_threads=1-internal/131072 34.82 µs 32.79 µs -5.8%
✅ sampling_latency/top_k_per_token 62.01 µs 57.86 µs -6.7%
✅ tokenization/encode_tokens_per_second 460.18 µs 428.56 µs -6.9%
✅ add/medium_f16_threads=1-internal/262144 126.65 µs 117.93 µs -6.9%
✅ matmul/medium_generic_f16_threads=1/32x512x512 43.39 µs 40.24 µs -7.3%
✅ add/small_f16_threads=1-internal/1024 568.0 ns 518.4 ns -8.7%
✅ qwen3_sampling_processors/top_p_fast_after_top_k 588.76 µs 518.73 µs -11.9%
✅ sampling_latency/greedy_per_token 3.84 µs 3.39 µs -11.9%
✅ tokenization/decode_tokens_per_second 7.26 ms 6.34 ms -12.6%
✅ qwen3_sampling_processors/top_k_top_p_fast 758.40 µs 657.59 µs -13.3%
✅ matmul/medium_generic_bf16_threads=1/32x512x512 776.99 µs 666.11 µs -14.3%
✅ qwen3_sampling_processors/top_k_full_sort_baseline 2.80 ms 2.39 ms -14.6%
✅ sampling_latency/min_p_per_token 274.18 µs 233.31 µs -14.9%
🟢 logit_processing/seven_processor_chain_per_step 380.56 µs 321.59 µs -15.5%
🟢 kv_cache/alloc_dealloc_pages 45.51 µs 38.38 µs -15.7%
🟢 add/small_f32_threads=1-internal/1024 297.5 ns 250.2 ns -15.9%
🟢 qwen3_sampling_processors/top_k_top_p_full_sort_baseline 6.72 ms 5.64 ms -16.1%
🟢 grammar_masking/llguidance_compute_mask/32 91.78 µs 76.47 µs -16.7%
🟢 qwen3_sampling_processors/top_k_partial_selection 213.28 µs 176.16 µs -17.4%
🟢 block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024 94.96 µs 75.68 µs -20.3%
🟢 matmul/large_generic_f16_threads=8/32x1024x1024 115.64 µs 91.91 µs -20.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: { 2.61 3.02 5.61 }
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)

justinchuby pushed a commit that referenced this pull request Aug 19, 2026
`Relu` had only a 1 Mi bench case. At that size the op moves 8 MiB and
runs at ~63 GB/s on one core — it is bandwidth-bound, so it reports
almost nothing about the arithmetic. 4 Ki fits in L1, is the size a
decode step actually sees, and is where every other activation in the
grid is already measured.

The first thing it says is worth knowing: on `origin/main` Relu at 4 Ki
is 1.64x ORT, the **worst** ratio in the 4 Ki grid, and the AVX2 kernel
in the previous commit does not move it (1.642 -> 1.638, 2/5 rounds).
Relu is the cheapest op in the set, so ORT's own time is smallest and
our fixed per-`Run` cost is the largest share of it. That is a dispatch
problem, not a kernel one, and it is the number to watch as #1244/#1246
and their successors land.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@codecov

codecov Bot commented Aug 19, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 85.34483% with 17 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.15%. Comparing base (06c62e0) to head (e06208c).
⚠️ Report is 79 commits behind head on main.

Files with missing lines Patch % Lines
crates/onnx-runtime-ep-plugin/src/compute.rs 75.71% 16 Missing and 1 partial ⚠️
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #1244      +/-   ##
==========================================
- Coverage   80.88%   80.15%   -0.73%     
==========================================
  Files         364      364              
  Lines      160729   161028     +299     
  Branches   160729   161028     +299     
==========================================
- Hits       130005   129077     -928     
- Misses      26069    27294    +1225     
- Partials     4655     4657       +2     
Flag Coverage Δ
mlas 85.09% <ø> (-0.14%) ⬇️
offline 80.06% <85.34%> (-0.74%) ⬇️

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

Files with missing lines Coverage Δ
crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs 89.03% <100.00%> (+1.49%) ⬆️
crates/onnx-runtime-ep-plugin/src/compute.rs 80.38% <75.71%> (+0.45%) ⬆️

... and 10 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@justinchuby
justinchuby merged commit 4a9f4ec into main Aug 19, 2026
14 of 21 checks passed
@justinchuby
justinchuby deleted the squad/resch-small-node-dispatch branch August 19, 2026 13:41
justinchuby added a commit that referenced this pull request Aug 19, 2026
…d (2.31 -> 1.71 at depth 100) (#1387)

> **Scope note.** This PR now carries the whole #1077 dispatch-overhead
stack. Enabling auto-merge across the stacked PRs caused each to merge
into its *parent feature branch* — those branches are unprotected; only
`main` requires checks — and GitHub retargeted the children, collapsing
the stack into this one. **No `main` protection was bypassed and no work
was lost.** This PR still waits on required CI before it can reach
`main`.

**What is in here, in dependency order — each was independently reviewed
as its own PR before it merged:**

| was | what | evidence |
|---|---|---|
| #1387 | deterministic phase/counter instrumentation for the dispatch
path | Opus 4.8 REQUEST CHANGES → all findings fixed |
| #1392 | stop building debug strings production throws away | 5 → 4
allocations per input per `Run` |
| #1394 | the #1077 benchmark grid (depth 1/10/100, static + dynamic) |
produced the first real decomposition |
| #1397 | stop asking ORT three times where the inputs live | **7 → 4**
FFI calls per `Run`, Opus 4.8 APPROVE WITH NITS → fixed |
| #1401 | routed-loop instrumentation; O(depth²) retirement scan; lazy
absent strides; 8-element dispatch-isolating cases | **2.310 → 1.713**
at depth 100 (−24%) |

**Headline result** — 8-wide Relu chain, depth 100, 6000 interleaved
iterations, pinned cores, ORT arm as drift control (moved 1.7%):

| | baseline | after |
|---|---|---|
| ours p50 | 54.2 µs | **41.0 µs** |
| ratio vs ORT | 2.310 | **1.713** |
| per-node | 0.52 µs | **0.39 µs** |

Two of my own earlier conclusions were falsified along the way and are
documented in #1077: a 10-node chain is **one** fused dispatch (so the
per-node gap contains no FFI at all), and the original grid was
measuring **memory bandwidth**, not dispatch.

---

## Why

Issue #1077 says our per-`Run` dispatch overhead is higher than ORT's.
Every time I have had to answer *where* that overhead is, the method has
been the same: hand-edit `Instant` probes into the hot path, run the
A/B, read the numbers, tear the probes out. That is slow, it is
unreviewable, and it leaves nothing behind — the next person starts from
zero, and nothing catches a regression that quietly adds a round trip
back.

This makes the measurement a permanent, reviewable part of the crate.

## What

`dispatch_probe`, gated behind a **non-default** Cargo feature, wired
through the CPU dispatch path:

| Phase | Where |
|---|---|
| `CallbackEntry` | `compute_execute` entry, up to the first real work |
| `MetadataQuery` | `read_inputs` |
| `TensorBind` | single-node input view construction |
| `Allocate` | `allocate_output` |
| `DispatchLookup` | `infer_shapes`, `prepare_workspace` |
| `KernelInvoke` | `execute_with_workspace` |
| `StatusCrossing` | `status_with_code` |

Plus event tallies for the three things actually worth removing:
`OrtFfiCall`, `DispatchAlloc`, `StatusCreated` (and `ComputeExecute` /
`NodeExecuted` as denominators).

## Design notes

**The counters are kept twice, deliberately.** The thread-local copy is
the precise one — a dispatch is a single-threaded story, and only
per-thread accumulation lets one measurement be read without a
concurrently-running sibling test bleeding into it. That isolation is
what makes the exact-count assertions below possible at all; I tried
global-only first and the counts were nondeterministic under the default
parallel test runner.

The global mirror exists because ORT chooses the thread that runs
`Compute`, and it need not be the thread that called `Run`. The e2e
harness loads this EP as a cdylib through
`RegisterExecutionProviderLibrary`, so it reads the counters back
through an exported `nxrt_dispatch_probe_snapshot` symbol — and it has
no way to be on ORT's worker thread. A thread-local-only probe would
report zero there, which reads as "we made no FFI calls" rather than
"you are looking at the wrong thread". Writing both costs one `Cell`
store and one relaxed `fetch_add`.

**Production pays nothing.** Without the feature, `PhaseGuard` is a
zero-sized type with no `Drop` impl and every entry point is an
`#[inline(always)]` empty function. Call sites use `guard.end()` rather
than `drop(guard)` so the production build does not read as dropping a
`Drop`-less ZST.

**Timing is gated a second time**, on `ONNX_GENAI_PROFILE_DISPATCH=1`. A
test asserting FFI call counts wants the counters but emphatically does
not want two `Instant::now()` calls added to every phase it is
measuring.

## The counts are pinned, not documented

The point of the module is the regression guard, so the numbers are
asserted rather than written in a comment. `kernel_ctx` gains a
hand-built `OrtApi` — zeroed, then filled in with only the entry points
`read_inputs` actually uses, so any call it makes that the test did not
anticipate faults loudly instead of silently succeeding — which lets
`read_inputs` run with no live ONNX Runtime at all.

Today, per `Run`:

- **8 FFI round trips** for one input (within `read_inputs`) — one
shared `GetInputCount`, then 7 per input
- **5 heap allocations** — the `Vec<OwnedInput>` once, then `dims`, the
`format!("input {i}")` label, `shape` and `strides` per input
- three inputs cost `1 + 3×7` and `1 + 3×4`, so this is a linear model
rather than a single data point
- an absent optional input short-circuits at 2 calls, which pins the
optional-slot path

Any change that adds a round trip or an allocation to the per-`Run` path
now has to come here and edit a number, in a diff a reviewer can see.

## Verification

Both feature configurations: `cargo fmt`, `cargo clippy --all-targets`
(zero warnings), and the full `--lib` suite — 251 tests without the
feature, 257 with it. The real ORT-driven suite passes unchanged:
**55/55 `plugin_ort_e2e`** with `NXRT_REQUIRE_ORT_TESTS=1`, including
`every_assigned_node_is_also_executed_by_this_ep`.

**Every test here was verified by mutation**, because a test that passes
against the known-broken version is not a test:

- Injecting one extra `get_type_shape` call and one extra `Vec`
allocation into `read_inputs` failed exactly the three count assertions
and correctly left the other two green.
- The zero-cost assertion **was vacuous when I first wrote it** — it
compared `snapshot()` after counting against `snapshot()` before, which
are two calls to the same constant function in a disabled build. It
passed unchanged against a mutant whose "disabled" `count_n` recorded
into a static and whose `snapshot` returned non-zero. It now compares
against `Counters::default()` — an absolute claim rather than a
self-consistency one — and fails against that mutant. It also asserts
`!needs_drop::<PhaseGuard>()`, closing the one way the production guard
could regain a cost while staying zero-sized.

## Review round 1 — Opus 4.8, `REQUEST CHANGES`, all addressed

The review found a real defect, and it was the module's central claim.

**Blocker: `OrtFfiCall` was documented as a per-`Run` total but only
`read_inputs` and `allocate_output` were instrumented.** On the happy
path of a single-node `Run`, `device_mem_info`, `ort_input_mem_info`,
`mem_info_is_device`, `alloc_scratch`, `host_pool::install` and
`CreateStatus` all reach into ORT and none were counted. A probe that
under-counts while calling itself exact is worse than no probe — it
invites exactly the wrong conclusion about where the overhead is.

Fixed by instrumenting every leaf `OrtApi` call in the crate, and then
by making that property **self-enforcing** rather than a promise. The
new `ffi_coverage` tests scan the source of each file that touches
`OrtApi` and extract the API members it names — ORT's C API is
`CamelCase` and our own fields are `snake_case`, so `.CamelCase`
isolates them cleanly — then fail if that set or the instrumentation
count moves. Verified by mutation: adding a stray
`GetTensorShapeElementCount` reference fails with

> `compute.rs now names 10 ORT API members ([…
"GetTensorShapeElementCount" …]), not 9. If you added an FFI call, add
an ort_call() beside it and update this table.`

A second test pins that the scan finds real members, so it cannot rot
into a heuristic that matches nothing and passes forever.

**`DispatchAlloc` got the opposite treatment: it is now documented as a
lower bound.** The reviewer was right that hand-placed allocation counts
cannot be exhaustive — `Vec::new()` does not allocate,
`Vec::with_capacity(0)` does not allocate, and whether a `collect`
allocates once or twice depends on `size_hint`, not on anything visible
at the call site. Claiming exactness there would repeat the same mistake
in a place where it cannot be fixed. It is exact for `read_inputs`,
where the sites are few and each is pinned; callers needing the true
whole-`Run` figure get `CountingAllocator`, a `GlobalAlloc` wrapper to
install in a test or bench. This crate deliberately does not install it,
since a library defining a `#[global_allocator]` takes that choice away
from every dependent.

Also fixed:
- A node with **no** inputs was charged one allocation for a
`Vec::with_capacity(0)` that never happens. Now conditional, pinned by
`a_node_with_no_inputs_allocates_nothing`.
- `CreateStatus` was counted as a `StatusCreated` but not as an FFI
call, undercounting error paths.
- **Phases are documented as not being a partition.** On success they
are non-exhaustive; on error they *nest*, because a guard closes at
scope exit rather than at the early `return`, so `StatusCrossing` opens
inside whichever phase was live. `DispatchLookup` is entered twice per
node and reports the total. Summing `phase_ns` gives neither wall time
nor a pie chart, and a reader who assumed otherwise would be misled.
- `probe_is_compiled_out_in_production` no longer implies it has proven
`count` is side-effect-free — a disabled build has no storage to observe
one. It proves what it can: the guard is a ZST that does not need
dropping.

Confirmed by the reviewer independently: `mem::zeroed::<OrtApi>()` is
sound (all 422 fields are `Option<unsafe extern "C" fn>`, so NPO makes
all-zero a valid `None` everywhere); the `without_provenance` fake
pointers are never dereferenced; and all four `read_inputs` counts are
correct, checked against a counting global allocator.

Re-validated after the fixes: clippy clean and `--lib` green in both
configurations (253 / 260), and **55/55 `plugin_ort_e2e`** against real
ORT.

## What this sets up

Two allocations are already visibly waste and are removed in a follow-up
that this PR exists to measure:

1. `compute_execute` calls `staging_log(&format!(…))`
**unconditionally** — `staging_log` early-returns when disabled, but the
`String` is built and three values formatted on every dispatch
regardless.
2. `read_inputs` builds `format!("input {i}")` eagerly per input, on the
success path, purely to label an error that is usually not raised.

Relates to #1077. Independent of #1244/#1246 — additive and
feature-gated, so it can land in any order.

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby added a commit that referenced this pull request Aug 19, 2026
)

## What this is

`read_inputs` runs **once per input per `Run`**, and it asked ORT for
the element
type and the dimensions through the classic five-call sequence:

```
GetTensorTypeAndShape  →  GetTensorElementType
                       →  GetDimensionsCount
                       →  GetDimensions
                       →  ReleaseTensorTypeAndShapeInfo
```

The first of those makes ORT **heap-allocate** an
`OrtTensorTypeAndShapeInfo` and
copy the shape into it; the last frees it. Five FFI crossings and one
allocation
on ORT's side of the boundary, per input, per `Run`, to read data the
`OrtValue`
already owns.

`GetTensorElementTypeAndShapeDataReference` (ORT C API, since 1.24)
returns the
element type **and a reference to the value's own shape array** in one
call,
allocating nothing. This PR routes `read_inputs` through it.

Stacked on #1244 (since merged; this branch now targets `main`
directly).
No kernel is touched, no numerics change, and node assignment
and execution ownership are unchanged: the CPU EP still executes every
node it
claims, locally, with ORT CPU fallback off.

## The three cuts

**1. Five ORT calls and one ORT-side allocation per input → one call, no
allocation.**

The plugin already fails closed below API 27, so the hook is always
present in
practice. The five-call sequence is kept as a fallback for a host that
leaves it
null, and a host offering **neither** now fails closed with a message
naming
both — it does not silently read garbage shapes.

The borrowed pointer is used only to build the owned `shape`/`strides`
of the
`OwnedInput`; nothing retains it. ORT spells a scalar as a **null**
pointer with
count 0, and `slice::from_raw_parts` is UB on null, so that case is
guarded
explicitly — see *Correctness* below for how that guard is actually
enforced.

**2. `validate_dims` takes a lazily-formatted label, not `&str`.**

Every call site passed `&format!("input {i}")`, i.e. a heap `String`
built on the
success path purely to label an error that almost never happens. Callers
now pass
`format_args!(..)` and the string is only materialised inside the error
branches.
One `String` per input per `Run` gone. All call sites converted,
including
`transfer.rs` and the unit tests; no message text changed.

(The merge with `main` took main's `impl std::fmt::Display` signature
here, which
is a superset of this branch's original `Arguments<'_>` and keeps both
callers
working unchanged.)

**3. `allocate_output` converts the output shape to `i64` on the
stack.**

`let dims: Vec<i64> = shape.iter().map(|&d| d as i64).collect()` was a
heap
allocation per output per `Run` whose entire lifetime was the
`KernelContext_GetOutput`
call below it. Rank ≤ 8 now goes to a stack array; a taller tensor still
works,
through the `Vec`.

For a 1-in/1-out elementwise node that is **5 fewer ORT FFI calls**, one
fewer
ORT-side allocation, and **2 fewer heap allocations** per `Run`, on top
of #1244.

## A/B, one thread

`taskset -c 8-15`, this branch vs its merge parent (`origin/main` +
#1244),
**interleaved** (A,B,A,B,… so drift hits both arms equally), 400
iterations after
50 warmup, **five rounds**. Ratio is **ours/ORT, lower is better**; each
column is
the median of the five rounds. "ORT drift" is the control: ORT's own p50
measured
in the two arms, which must not move for the comparison to mean
anything.

| case | before p50 | after p50 | before p90 | after p90 | ORT drift |
rounds won |
|---|---|---|---|---|---|---|
| `sqrt_f32_4k` | 1.023 | **0.952** | 1.031 | **0.964** | +2.9 % | 4/5 |
| `leakyrelu_f32_4k` | 1.202 | **1.097** | 1.210 | **1.104** | 0.0 % |
5/5 |
| `hardsigmoid_f32_4k` | 1.213 | **1.144** | 1.226 | **1.149** | 0.0 % |
5/5 |
| `thresholdedrelu_f32_4k` | 1.316 | **1.209** | 1.314 | **1.220** | 0.0
% | 5/5 |
| `sigmoid_f32_4k` | 1.328 | **1.252** | 1.333 | **1.258** | 0.0 % | 5/5
|
| `tanh_f32_4k` | 1.368 | **1.287** | 1.369 | **1.291** | 0.0 % | 5/5 |
| `erf_f32_4k` | 1.383 | **1.327** | 1.387 | **1.329** | 0.0 % | 5/5 |
| `celu_f32_4k` | 0.430 | **0.411** | 0.435 | **0.418** | 0.0 % | 4/5 |
| `elu_f32_4k` | 0.437 | **0.416** | 0.443 | **0.422** | −1.8 % | 4/5 |
| `selu_f32_4k` | 0.467 | **0.441** | 0.471 | **0.446** | 0.0 % | 5/5 |
| `log_f32_4k` | 0.727 | **0.700** | 0.728 | **0.706** | 0.0 % | 5/5 |
| `mish_f32_4k` | 0.268 | **0.264** | 0.271 | **0.266** | 0.0 % | 4/5 |

12 of 12 cases improve at p50 and p90; 56 of 60 individual arm-pairs
favour the
change, and no case regressed in its median. In absolute terms it is
**0.2–0.3 µs
off every `Run`**: `tanh_f32_4k` 4.20 µs → 3.90 µs against ORT's 3.10
µs,
`thresholdedrelu_f32_4k` 3.10 µs → 2.80 µs against ORT's 2.30 µs. That
is the
shape of a fixed per-call cost being removed, which is what it is.

`sqrt_f32_4k` crosses below 1.00 — the plugin path now dispatches that
op faster
than plain ORT does.

## A/B, four threads

Same protocol, `NXRT_MM_BENCH_THREADS` =
`ONNX_GENAI_MLAS_THREADPOOL_THREADS` =
`RAYON_NUM_THREADS` = 4, three rounds.

| case | before p50 | after p50 | before p90 | after p90 |
|---|---|---|---|---|
| `leakyrelu_f32_4k` | 1.069 | **0.966** | 1.075 | **0.976** |
| `hardsigmoid_f32_4k` | 1.152 | **1.050** | 1.164 | **1.060** |
| `thresholdedrelu_f32_4k` | 1.212 | **1.091** | 1.228 | **1.106** |
| `sqrt_f32_4k` | 1.121 | **1.036** | 1.130 | **1.045** |
| `sigmoid_f32_4k` | 1.346 | **1.264** | 1.359 | **1.263** |
| `tanh_f32_4k` | 1.434 | **1.319** | 1.453 | **1.325** |
| `erf_f32_4k` | 1.581 | **1.502** | 1.588 | **1.510** |
| `celu_f32_4k` | 0.460 | **0.437** | 0.476 | **0.445** |
| `elu_f32_4k` | 0.481 | **0.448** | 0.486 | **0.468** |
| `selu_f32_4k` | 0.500 | **0.458** | 0.482 | **0.464** |
| `log_f32_4k` | 0.766 | **0.731** | 0.781 | **0.734** |
| `mish_f32_4k` | 0.332 | **0.322** | 0.345 | **0.339** |

36 of 36 arm-pairs favour the change. The cost is per call, not per
element or
per worker, so the gain is the same absolute microseconds at every
thread count.

## Drift control: the 1 Mi grid does not move

Same protocol, `_f32_1m`, 200 iterations after 30 warmup, five
interleaved rounds,
one thread. A 1 Mi tensor amortises a fixed per-call cost away, so
**these must
not move** — and they don't:

| case | before | after | | case | before | after |
|---|---|---|---|---|---|---|
| `relu_f32_1m` | 1.024 | 1.028 | | `sqrt_f32_1m` | 0.313 | 0.309 |
| `exp_f32_1m` | 0.580 | 0.576 | | `log_f32_1m` | 0.264 | 0.258 |
| `sigmoid_f32_1m` | 0.588 | 0.584 | | `mish_f32_1m` | 0.100 | 0.099 |
| `tanh_f32_1m` | 0.618 | 0.620 | | `celu_f32_1m` | 0.127 | 0.129 |
| `erf_f32_1m` | 0.651 | 0.668 | | `elu_f32_1m` | 0.129 | 0.129 |
| `gelu_exact_f32_1m` | 0.585 | 0.584 | | `selu_f32_1m` | 0.136 | 0.138
|
| `gelu_tanh_f32_1m` | 0.602 | 0.646 | | `hardsigmoid_f32_1m` | 0.455 |
0.453 |
| `fastgelu_f32_1m` | 0.608 | 0.639 | | `leakyrelu_f32_1m` | 0.409 |
0.417 |
| `quickgelu_f32_1m` | 0.429 | 0.438 | | `thresholdedrelu_f32_1m` |
0.567 | 0.573 |

The per-case round-win counts here are 0/5–4/5 with a median of 2/5,
i.e. a coin
flip — the signature of noise, not of an effect. The two largest movers,
`gelu_tanh` (+7 %) and `fastgelu` (+5 %), are the two cases whose
**ORT** side also
drifted most in the same rounds (−4.4 % and −2.8 %), so the ratio moved
because
the denominator did. This host's floor is roughly ±5 % on the 1 Mi grid
and I am
not claiming anything below it.

## Correctness

**Five new tests, three with a verified falsifier**, driving
`read_inputs` and
`allocate_output` against a hand-built `OrtApi` whose hooks count their
calls.

| test | what it pins | falsifier |
|---|---|---|
| `input_shapes_come_from_one_call_when_ort_offers_the_reference_hook` |
1 reference call, **0** legacy calls, and the shape/strides/dtype that
come out | forcing the legacy route fails it with `left: 0, right: 1`
(run) |
| `the_five_call_fallback_produces_the_same_input_as_the_reference_hook`
| the fallback is exercised and its `OwnedInput` matches the reference
path field for field | — (it *is* the parity check) |
| `a_borrowed_scalar_shape_never_dereferences_null` | ORT's documented
scalar spelling — null pointer, count 0 — takes the borrowed route and
yields rank 0 | deleting the null guard makes **Miri** report
`out-of-bounds pointer use: null pointer is a dangling pointer` at
`kernel_ctx.rs:280` (run) |
| `read_inputs_fails_closed_when_no_shape_route_exists` | the error
names **both** routes; no silent garbage shapes | — |
| `output_dims_are_identical_on_the_inline_and_heap_ranks` | ranks
0/1/8/9/12 all reach ORT with every dimension intact | dropping the heap
arm delivers rank 9 truncated to 8 dims (run) |

**The null guard is now actually enforced, not just asserted.**
Natively,
`from_raw_parts(null, 0)` returns an empty slice, so a test asserting
the
resulting shape passes with or without the guard — the guarantee only
exists if
Miri sees it. `onnx-runtime-ep-plugin` was not in the Miri matrix, so
this PR adds
`kernel_ctx::` as a lane in `.github/workflows/miri.yml` (the crate's
other
modules dlopen ORT and are not Miri-tractable, which is why the lane is
scoped to
the module rather than the crate). 25 tests pass under Miri in 4.15 s;
with the
guard deleted the lane fails. That pairing is what makes the test
non-vacuous.

**Suites.** `-p onnx-runtime-ep-plugin`: **252** unit tests pass (247
before,
+5 new). `-p onnx-runtime-ep-cpu-plugin` with
`NXRT_REQUIRE_ORT_TESTS=1`: every
suite green, including all **55** `plugin_ort_e2e` cases —
`every_assigned_node_is_also_executed_by_this_ep`,
`no_supported_node_is_ever_left_to_the_ort_cpu_ep`,
`no_matmul_family_node_escapes_to_the_ort_cpu_ep` and
`every_fixture_loads_with_cpu_fallback_disabled` all pass, so **assigned
still
equals executed** with ORT CPU fallback disabled. `cargo clippy
--release
--all-targets -p onnx-runtime-ep-plugin` and `cargo fmt` are clean.

**Build identity.** Pure native CPU EP: no MLAS at runtime, no ORT CPU
EP
fallback, no new dependency. AVX2/FMA host (`avx2 fma f16c`, no
AVX-512), so ORT
and we are on the same instruction footing.

## Independent review

Reviewed by **Claude Opus 4.8**, read-only, briefed with the exact ORT
contract
for the borrowed pointer and asked specifically to hunt UB and vacuous
tests.
Verdict **APPROVE**, no blockers. It confirmed the null/scalar guard is
unreachable-by-construction for `from_raw_parts`, that the borrow cannot
outlive
the `OrtValue` (its only consumer copies), that
`ReleaseTensorTypeAndShapeInfo`
still runs on every legacy error path — and noted that moving
`DataType::from_onnx`
after the match incidentally closes a **pre-existing** leak of
`type_shape` on the
unsupported-dtype path.

It raised two test-quality defects, both real and both **fixed** in
`6027f6e16`:

1. The three counting tests shared two process-wide `AtomicUsize`es
while
asserting exact counts, and libtest runs them in parallel — one test's
reset
   could land inside another's assertion window. They now take a shared
   `SHAPE_COUNTER_LOCK` that resets both counters under the guard.
2. `a_borrowed_scalar_shape_never_dereferences_null` was **vacuous**
with respect
to its name, for exactly the reason above, and the crate was not under
Miri.
Hence the Miri lane, and the test now also asserts the scalar went
through the
   borrowed route so it cannot pass on the fallback.

It also flagged that `OrtStatus` is not released on error paths in this
file.
That is pre-existing and repo-wide in `kernel_ctx.rs` — the old
five-call code
leaked identically — so it is not touched here; it belongs in its own
change.

## What is left

After #1244 and this PR, a 1-in/1-out elementwise `Run` is down to
roughly
`KernelContext_GetInputCount` + `GetInput` + the one shape reference +
`GetTensorData` + `GetOutput` + `GetTensorMutableData`. What remains:

* **`OwnedInput`'s `shape` and `strides` `Vec`s**, `kernel_inputs`,
`infer_shapes`'s `Vec<Vec<usize>>`, `slot_map`, `output_views`, and the
`Box<HostPool>` in `host_pool::install` — each worth ~25–35 ns. Removing
them
needs an inline-capacity vector type or per-session caching of the parts
that
cannot change between `Run`s. Both are worth doing; neither belongs
here.
* **The `HostPool` box** specifically is @sebastian's 16-thread
scheduling file
  and should come from him or after his PRs land.

---

## Refreshed against `main` @ `6a855d5e0`, and measured as a stack

`origin/main` moved a long way while this sat in the CI queue (#1346,
#1352 and
#1361 on the quality lane; #1154, #1232, #1238 on the CPU side). Merged
in
normally — no rebase — and re-measured from scratch against the new
baseline.

Production pure-native A/B, plain ORT as the control arm. No MLAS, no
ORT CPU
fallback, no deferral. `taskset -c 8-15`, one thread, 400 iterations,
five
interleaved rounds out of two worktrees, started only once cores 8-15
were
>=93% idle. **Ratio is ours/ORT, lower is better.** `before` is `main`
at
`6a855d5e0`; `after` is #1244 + #1246 together, since #1246 is stacked
on #1244
and the pair is what a user gets.

| case | ratio p50 main | ratio p50 stack | Δ | ratio p90 main | ratio
p90 stack | ours us | ORT drift | rounds won |
|---|---|---|---|---|---|---|---|---|
| `thresholdedrelu_f32_4k` | 1.512 | **1.216** | -19.6% | 1.519 |
**1.216** | 3.6 → **2.8** | -4.2% | 5/5 |
| `tanh_f32_4k` | 1.511 | **1.279** | -15.4% | 1.513 | **1.284** | 4.6 →
**3.9** | +0.0% | 5/5 |
| `sigmoid_f32_4k` | 1.475 | **1.248** | -15.4% | 1.484 | **1.256** |
4.7 → **4.0** | +0.0% | 5/5 |
| `erf_f32_4k` | 1.470 | **1.363** | -7.3% | 1.489 | **1.364** | 7.5 →
**7.0** | +0.0% | 5/5 |
| `hardsigmoid_f32_4k` | 1.416 | **1.127** | -20.4% | 1.426 | **1.133**
| 3.5 → **2.8** | +0.0% | 5/5 |
| `leakyrelu_f32_4k` | 1.361 | **1.094** | -19.6% | 1.371 | **1.104** |
3.5 → **2.8** | +0.0% | 5/5 |
| `sqrt_f32_4k` | 1.141 | **0.946** | -17.1% | 1.150 | **0.955** | 4.0 →
**3.4** | -2.8% | 5/5 |
| `log_f32_4k` | 0.767 | **0.698** | -9.0% | 0.776 | **0.703** | 7.9 →
**7.2** | +0.0% | 5/5 |
| `selu_f32_4k` | 0.505 | **0.440** | -12.9% | 0.512 | **0.444** | 5.6 →
**4.9** | +0.0% | 5/5 |
| `elu_f32_4k` | 0.480 | **0.416** | -13.3% | 0.488 | **0.421** | 5.3 →
**4.6** | +0.0% | 5/5 |
| `celu_f32_4k` | 0.470 | **0.411** | -12.6% | 0.478 | **0.418** | 5.7 →
**5.0** | +0.0% | 5/5 |
| `mish_f32_4k` | 0.276 | **0.264** | -4.3% | 0.278 | **0.268** | 17.3 →
**16.6** | +0.0% | 5/5 |

**Every case, every round.** The two rows with a moving control (`sqrt`
-2.8%,
`thresholdedrelu` -4.2%) are reported rather than dropped; both won 5/5
anyway
and their absolute time fell by the same ~0.7 us as everything else.

That constant ~0.7 us is the point. It is not proportional to tensor
size — the
same absolute amount comes off `hardsigmoid` (3.5 -> 2.8 us) as off
`mish`
(17.3 -> 16.6 us) — which is what a fixed per-`Run` cost looks like when
you
remove some of it. It moves the cheap ops the most because they had the
least
to hide it behind, and `sqrt` crosses from 1.141 to **0.946**, from a
loss to a
win.

### Where the remaining time goes

Measured directly, by instrumenting `compute_execute` segment by segment
on top
of this stack (temporary probe, not committed; `perf` is unavailable on
this
host — `perf_event_paranoid=4`). Per `Run`, one-in/one-out elementwise
node,
4096 `f32`, microseconds:

| segment | us | note |
|---|---|---|
| `KernelContext_GetOutput` | 0.35 | ORT's own API — ours to call, not
to optimise |
| `read_inputs` | 0.15 | 4 ORT FFI calls, already one shape call after
#1246 |
| rest of `allocate_output` | 0.13 | `GetTensorMutableData` + strides |
| `prepare_workspace` | 0.09 | metadata vector + plan-cache lookup, for
a kernel needing 0 bytes |
| `host_pool::install` | 0.05 | @sebastian's, not touched |
| `infer_shapes` | 0.05 | |
| `kernel_inputs` | 0.04 | |
| `output_views` | 0.04 | |

Non-kernel node cost is **~1.25 us and near-constant across all twelve
operators** (0.28 to 15.2 us of kernel time), which is the direct
confirmation
that small-node ratios on this EP are dispatch-bound rather than
kernel-bound.
There is no single large item left — the biggest,
`KernelContext_GetOutput`, is
ORT's. The rest is a long tail of 0.04-0.15 us items, which is what
#1358
(`InlineVec`) starts on.

### And nothing breaks at 1 Mi

Same harness, 1048576 elements, 120 iterations, 3 rounds. A fixed
per-`Run`
cost should be invisible here, and it is:

| case | ratio p50 main | ratio p50 stack | ours us | ORT drift |
|---|---|---|---|---|
| `celu_f32_1m` | 0.141 | 0.140 | 382.1 → 377.9 | -0.0% |
| `elu_f32_1m` | 0.138 | 0.136 | 347.7 → 343.6 | +0.1% |
| `erf_f32_1m` | 0.671 | 0.670 | 595.2 → 594.6 | +0.1% |
| `exp_f32_1m` | 0.606 | 0.590 | 247.0 → 240.5 | +0.1% |
| `fastgelu_f32_1m` | 0.643 | 0.648 | 415.3 → 413.2 | -1.7% |
| `gelu_exact_f32_1m` | 0.572 | 0.592 | 706.1 → 710.6 | -2.8% ⚠ |
| `gelu_tanh_f32_1m` | 0.657 | 0.651 | 414.8 → 410.5 | +0.3% |
| `hardsigmoid_f32_1m` | 0.376 | 0.372 | 88.0 → 69.6 | -0.9% |
| `leakyrelu_f32_1m` | 0.421 | 0.430 | 90.1 → 87.3 | -2.2% ⚠ |
| `log_f32_1m` | 0.270 | 0.274 | 616.4 → 612.4 | +1.7% |
| `mish_f32_1m` | 0.105 | 0.105 | 1626.0 → 1626.0 | -0.2% |
| `quickgelu_f32_1m` | 0.455 | 0.453 | 321.5 → 320.0 | -0.5% |
| `relu_f32_1m` | 1.034 | 1.022 | 131.7 → 130.2 | +0.0% |
| `selu_f32_1m` | 0.147 | 0.148 | 374.1 → 371.8 | -1.5% |
| `sigmoid_f32_1m` | 0.471 | 0.606 | 231.6 → 229.0 | -23.4% ⚠ |
| `sqrt_f32_1m` | 0.314 | 0.302 | 148.1 → 143.9 | +1.1% |
| `tanh_f32_1m` | 0.644 | 0.631 | 226.2 → 221.9 | -0.2% |
| `thresholdedrelu_f32_1m` | 0.500 | 0.486 | 70.9 → 68.6 | -0.3% |

Flat, as predicted — 0.7 us against 70-1626 us of work. Absolute time is
equal
or better in 16 of 18 cases. The two ⚠ rows had the control move more
than the
effect: `sigmoid` is unusable (ORT itself moved -23.4%; our own absolute
went
231.6 -> 229.0 us), and `gelu_exact`'s +0.6% absolute sits inside its
-2.8%
control. Reported rather than dropped.

This is the coverage claim for the change: it buys ~0.7 us at every
size, which
is 20% of a small node and nothing at all of a large one, and it costs
nothing
anywhere.

---

## Status after merging latest `main` (2026-08-19)

Validated on latest `main` after the six-PR #1077 stack landed
(#1387, #1409, #1412, #1430, #1433, #1472). Full evidence in the PR
comment
below; the measured summary:

**Deterministic counters, `relu_1_tiny`** — `OrtFfiCall` **10 → 6 per
`Run`**,
`DispatchAlloc` **24 → 20**. Four fewer round trips for one input, i.e.
the
predicted 7 → 3 per input.

**Timing A/B**, production build, 3 clean reps (a 4th discarded for
contention —
it read 0.973, which would have flattered us):

| case | main | this PR |
|---|---|---|
| `relu_1_tiny` | 1.397 | **1.271** |
| `relu_10_tiny` | 1.250 | **1.173** |
| `relu_100_tiny` | 1.163 | **1.110** |

The gain shrinks as depth grows — the signature of a fixed-cost fix,
since this
removes FFI calls per `Run`, not per node. Fixed per-`Run` overhead
**1.38 →
1.23** against ORT; per-node slope unchanged.

**Test coverage gap this merge exposed and closed:** the pinned costs
(`1 + 7` calls, `1 + 3` allocations) run against a zeroed `fake_api()`,
where the
reference hook is null — so they describe the *legacy fallback*, not the
fast
path this PR exists to add. Added
`the_reference_hook_path_costs_exactly_three_ort_calls_per_input` (pins
`1 + 3`)
and `the_reference_hook_path_allocates_less_than_the_legacy_path`.
`dispatch_probe`'s FFI-coverage table updated to 12 members / 12
`ort_call()`
sites, which is the guard that flagged the new API member in the first
place.

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Resch <resch@squad.local>
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