fix(runtime): infer Compress output shapes for large conditions (replaces #1195)

[justinchuby]

What this is

This is #1195 (fix(runtime): infer Compress output shapes for large conditions) rebased onto current main (which now includes #1251, the replacement for its stacked base #1159), plus a follow-up commit that adds the falsifying test the original was missing. Opened as a replacement branch because the original fix/compress-condition-shapes was stacked on a split/* branch, which makes gh pr merge 403. Closes #1195.

The bug

A dynamic Compress output shape is the count of selected entries, so the executor must read the boolean condition tensor to size the output. The dynamic-shape reader (resolve_node_outputs) admitted a shape input only if it passed bounded_shape_input, whose cap is MAX_SHAPE_DATA_ELEMS = 1024 — sized for scalar/vector shape metadata. A vision graph whose Compress condition spans a full image is longer than 1024, so it was rejected, dynamic_output_shapes received a None condition, and the run failed with UnresolvedShape { op: "Compress" } (a 500).

The fix

• geometry.rs: add bounded_compress_condition — a 1-D Bool tensor up to MAX_COMPRESS_CONDITION_ELEMS = 1 << 20.
• dispatch.rs: route the Compress condition input (default domain, op_type == "Compress", input index 1) through a new compress_condition_i64 that uses the larger bound, while every other shape input still goes through the ordinary shape_input_i64. So the one operand Compress genuinely needs is materialized, and an unrelated large tensor still is not.

Test-discipline gap I closed

The PR shipped compress_condition_allows_image_sized_boolean_vectors, which only asserts the bounded_compress_condition predicate. Reverting the dispatch.rs routing (the behavioural half of the fix) leaves that test green — it cannot catch a regression of the fix. I added:

1. compress_runs_with_image_sized_condition_including_empty_and_full — builds a Compress graph whose Bool condition initializer is MAX_SHAPE_DATA_ELEMS + 500 (1524) elements and runs it end to end on the CPU EP, asserting the output shape resolves and exactly the selected rows are gathered. Covers the all-false (empty output) and all-true (full output) boundaries through the real resolve + kernel path.
2. dynamic_output_shapes_compress_boundary_counts — the selected-count sizer at its boundaries: condition longer than the axis (clamped), shorter than the axis, all-false, all-true.

Falsification (revert -> RED -> restore -> GREEN)

Reverting only the dispatch.rs routing back to shape_input_i64 (keeping the helper + method so it stays a clean behavioural RED, not a compile error) makes compress_runs_with_image_sized_condition_including_empty_and_full fail with exactly:

UnresolvedShape { value: "y", op: "Compress" }

while the shipped predicate-only test and the counting test stay green (confirming they do not cover the routing). Restoring the routing returns all 4 Compress tests to green.

Validation

• cargo test -p onnx-runtime-session --lib: 185 passed, 0 failed on current main.
• Compress subset: 4 passed (shipped predicate test + shipped counting test + my 2 new tests).

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

[codecov]

Codecov Report

❌ Patch coverage is 71.87500% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.93%. Comparing base (5417d04) to head (d1686f0).
⚠️ Report is 12 commits behind head on main.

Files with missing lines	Patch %	Lines
...ates/onnx-runtime-session/src/executor/dispatch.rs	68.96%	6 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1253      +/-   ##
==========================================
- Coverage   79.95%   79.93%   -0.02%     
==========================================
  Files         360      360              
  Lines      158518   158790     +272     
  Branches   158518   158790     +272     
==========================================
+ Hits       126741   126936     +195     
- Misses      27160    27231      +71     
- Partials     4617     4623       +6     

Flag	Coverage Δ
mlas	85.19% <ø> (+0.09%)	⬆️
offline	79.83% <71.87%> (-0.02%)	⬇️

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

Files with missing lines	Coverage Δ
...ates/onnx-runtime-session/src/executor/geometry.rs	92.87% <100.00%> (+0.06%)	⬆️
...ates/onnx-runtime-session/src/executor/dispatch.rs	69.13% <68.96%> (-0.04%)	⬇️

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[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
🔴	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	53.65 µs	123.25 µs	+129.7%
🔴	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	62.63 µs	127.82 µs	+104.1%
🔴	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	61.92 µs	114.01 µs	+84.1%
🔴	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	735.93 µs	1.20 ms	+63.6%
🔴	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	392.70 µs	592.57 µs	+50.9%
🔴	matmul/medium_generic_bf16_threads=8/32x512x512	453.10 µs	648.31 µs	+43.1%
🔴	grammar_masking/llguidance_compute_mask/32	81.10 µs	110.56 µs	+36.3%
🔴	gather/large_f32_threads=1-internal/131072	29.53 µs	38.85 µs	+31.6%
⚠️	matmul/large_generic_bf16_threads=8/32x1024x1024	1.88 ms	2.22 ms	+18.3%
✅	sampling_latency/top_p_per_token	384.31 µs	435.88 µs	+13.4%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	96.23 µs	108.45 µs	+12.7%
✅	gather/small_bf16_threads=1-internal/4096	498.1 ns	556.4 ns	+11.7%
✅	logit_processing/seven_processor_chain_per_step	342.87 µs	379.85 µs	+10.8%
✅	matmul/large_generic_f32_threads=8/32x1024x1024	5.75 ms	6.36 ms	+10.7%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	88.93 µs	97.71 µs	+9.9%
✅	reduce_mean/medium_f32_threads=1-internal/65536	296.09 µs	318.92 µs	+7.7%
✅	reduce_mean/small_f32_threads=1-internal/4096	17.89 µs	19.15 µs	+7.1%
✅	sampling_latency/top_k_per_token	55.40 µs	57.86 µs	+4.4%
✅	sampling_latency/min_p_per_token	226.56 µs	234.54 µs	+3.5%
✅	matmul/small_generic_f16_threads=1/1x256x256	32.19 µs	32.78 µs	+1.8%
✅	sampling_latency/greedy_per_token	3.60 µs	3.62 µs	+0.7%
✅	qwen3_sampling_processors/top_k_top_p_fast	729.99 µs	726.04 µs	-0.5%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	2.38 ms	2.37 ms	-0.6%
✅	reduce_mean/large_f32_threads=1-internal/262144	1.13 ms	1.10 ms	-2.5%
✅	gather/small_f16_threads=1-internal/4096	569.0 ns	553.8 ns	-2.7%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	600.13 µs	576.94 µs	-3.9%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	6.49 ms	6.19 ms	-4.6%
✅	gather/small_f32_threads=1-internal/4096	786.0 ns	748.5 ns	-4.8%
✅	gather/medium_f32_threads=1-internal/32768	4.23 µs	4.00 µs	-5.6%
✅	kv_cache/alloc_dealloc_pages	44.57 µs	40.99 µs	-8.0%
✅	gather/large_f16_threads=1-internal/131072	15.93 µs	14.40 µs	-9.6%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	660.42 µs	569.75 µs	-13.7%
🟢	matmul/small_generic_f32_threads=8/1x256x256	54.24 µs	45.56 µs	-16.0%
🟢	tokenization/encode_tokens_per_second	518.32 µs	433.09 µs	-16.4%
🟢	tokenization/decode_tokens_per_second	8.69 ms	7.17 ms	-17.5%
🟢	matmul/large_generic_f32_threads=1/32x1024x1024	12.22 ms	10.03 ms	-17.9%
🟢	matmul/small_generic_f32_threads=1/1x256x256	46.22 µs	37.47 µs	-18.9%
🟢	matmul/medium_generic_f32_threads=8/32x512x512	1.20 ms	956.80 µs	-20.1%
🟢	qwen3_sampling_processors/top_k_full_sort_baseline	3.02 ms	2.38 ms	-21.0%
🟢	add/large_bf16_threads=1-internal/4194304	2.73 ms	2.15 ms	-21.2%
🟢	add/medium_bf16_threads=1-internal/262144	157.71 µs	122.86 µs	-22.1%
🟢	gather/medium_bf16_threads=1-internal/32768	3.54 µs	2.74 µs	-22.6%
🟢	add/small_f32_threads=1-internal/1024	287.0 ns	221.2 ns	-22.9%
🟢	matmul/medium_generic_f16_threads=1/32x512x512	39.39 µs	30.33 µs	-23.0%
🟢	matmul/medium_generic_f16_threads=8/32x512x512	40.00 µs	30.26 µs	-24.3%
🟢	matmul/medium_generic_f32_threads=1/32x512x512	3.12 ms	2.35 ms	-24.9%
🟢	qwen3_sampling_processors/top_k_partial_selection	218.22 µs	162.11 µs	-25.7%
🟢	gather/large_bf16_threads=1-internal/131072	20.84 µs	15.17 µs	-27.2%
🟢	matmul/small_generic_f16_threads=8/1x256x256	42.91 µs	30.80 µs	-28.2%
🟢	matmul/small_generic_bf16_threads=1/1x256x256	44.97 µs	31.76 µs	-29.4%
🟢	add/small_bf16_threads=1-internal/1024	752.6 ns	510.4 ns	-32.2%
🟢	add/large_f32_threads=1-internal/4194304	1.13 ms	756.52 µs	-32.8%
🟢	gather/medium_f16_threads=1-internal/32768	4.05 µs	2.70 µs	-33.1%
🟢	add/medium_f32_threads=1-internal/262144	43.75 µs	29.13 µs	-33.4%
🟢	add/small_f16_threads=1-internal/1024	760.5 ns	494.1 ns	-35.0%
🟢	matmul/small_generic_bf16_threads=8/1x256x256	48.96 µs	31.37 µs	-35.9%
🟢	add/large_f16_threads=1-internal/4194304	2.89 ms	1.71 ms	-40.7%
🟢	add/medium_f16_threads=1-internal/262144	204.95 µs	120.35 µs	-41.3%
🟢	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	6.67 ms	3.87 ms	-42.0%

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.22 3.49 6.16 }

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)