Container-aware control flow: If/Loop sequence threading (#449 inc3a)

[justinchuby]

Follows #486 (inc2 sequence mutations + conversions, now merged). Rebased onto main since #486's branch was deleted on merge. refs #449. CPU-only, pure shape-inference crate.

Summary

Increment 3a of the #449 container-type roadmap: ValueType (Sequence/Optional/Map) container types now thread through If branches and Loop carried dependencies — the two headline control-flow cases where sequences are actually used in real models (an If returning a sequence from both branches; a Loop carrying a sequence accumulator).

What's implemented

• If: output container type = unify(then_out, else_out) via a new recursive unify_value_type (Tensor⨝Tensor at the leaves reusing the foundation's unify_tensor_type; Sequence/Optional recurse; Map requires an equal key dtype). Element-extent disagreement → fresh symbol; element dtype disagreement → error; one branch container + other tensor → error.
• Loop: a container carried operand seeds the body formal input (body_container_seeds); the body threads/mutates it (SequenceInsert/Identity/passthrough); the carried body-output container type flows to the Loop output, remapped to the parent symbol space. Loop scan outputs stack tensors → never containers.
• InferenceReport.containers is the observation channel for tests and the body→parent threading channel. infer_graph_scoped gains a seed_containers param and raises the child interner floor above seeded container symbols. New CfOutput::Container variant.

DRY

Promoted the foundation's merge_tensor/merge_shape from handlers/container.rs to context.rs as unify_tensor_type/merge_element_shape (shared by the container handlers and infer.rs). Container symbol remap mirrors the exact per-dim rule the tensor CF path uses.

Byte-identical tensor path

All container work is gated on a non-empty container map — empty for pure-tensor graphs, no extra fresh_dim minted → symbol numbering unchanged. tensor_only_path_is_byte_identical_after_container_type_model stays GREEN, as do all existing If/Loop/Scan tensor tests.

Honest limitation

A seeded symbolic element dim crossing a Loop body degrades to a fresh parent symbol (container seeds never touch a body IR Value, so it's absent from child_to_parent) — SOUND, dtype/structure/concrete extents preserved. Covered by loop_passthrough_preserves_seeded_sequence_dtype.

Scoped out

• inc3b: Scan container state vars (mechanically identical to Loop seeding, rare, no scan-output payoff).
• inc4: SequenceMap + cross-subgraph container capture (remap_node_io sets value_type: None today).

Verification

• cargo test -p onnx-runtime-shape-inference: graph_inference 27 → 33 (+6 container CF tests); op_rules 275, container 16, doctest 1 — all GREEN; byte-identical regression GREEN.
• cargo fmt --all --check clean; cargo clippy -p onnx-runtime-shape-inference --all-targets -D warnings clean.
• cargo build -p onnx-runtime-session -p onnx-runtime-eager GREEN.
• Catalog counts unchanged at 217 ops / 262 entries (inc3a registers no new operators).

Design + status notes: .squad/decisions/inbox/harry-449-inc3-design.md, .squad/decisions/inbox/harry-449-inc3.md.

Do NOT merge.

[codecov]

Codecov Report

❌ Patch coverage is 88.78924% with 25 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.59%. Comparing base (2b0b2f9) to head (684a60d).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/onnx-runtime-shape-inference/src/infer.rs	88.73%	13 Missing and 3 partials ⚠️
crates/onnx-runtime-shape-inference/src/context.rs	87.32%	9 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff            @@
##             main     #527    +/-   ##
========================================
  Coverage   80.58%   80.59%            
========================================
  Files         315      315            
  Lines      122900   123077   +177     
  Branches   122900   123077   +177     
========================================
+ Hits        99037    99191   +154     
- Misses      19838    19859    +21     
- Partials     4025     4027     +2     

Flag	Coverage Δ
cli-ort-linux	83.27% <ø> (ø)
cli-ort-windows	82.78% <ø> (ø)
mlas	77.91% <ø> (ø)
offline	80.51% <88.78%> (+<0.01%)	⬆️

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

Files with missing lines	Coverage Δ
...rates/onnx-runtime-shape-inference/src/dim_expr.rs	97.10% <100.00%> (+0.07%)	⬆️
...-runtime-shape-inference/src/handlers/container.rs	86.33% <100.00%> (-2.50%)	⬇️
crates/onnx-runtime-shape-inference/src/report.rs	100.00% <ø> (ø)
crates/onnx-runtime-shape-inference/src/context.rs	89.91% <87.32%> (-0.65%)	⬇️
crates/onnx-runtime-shape-inference/src/infer.rs	92.03% <88.73%> (-0.40%)	⬇️

🚀 New features to boost your workflow:

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

[justinchuby]

VERDICT: APPROVE

Independent review by Lori (reviewer). Author Harry locked out. Worktree origin/squad/449-container-inc3 @ 684a60d, base main. CPU-only shape-inference crate. Reviewed the isolated commit diff (HEAD~1..HEAD), not the noisy main...HEAD range.

1. Byte-identical tensor control-flow [CRITICAL] — PASS, gate is real

The container path is gated on a genuinely-empty container map on tensor graphs:

• seed_containers defaults to &HashMap::new() at the top level; containers = seed_containers.clone() (infer.rs:87) is empty for pure-tensor graphs and only ever grows via CfOutput::Container, which requires a branch/body report.containers hit (infer.rs:494-513, :955-963) — impossible without a Sequence/Optional/Map producer.
• Interner floor: seed_next_symbol(graph).max(next_container_symbol(seed_containers)) (infer.rs:70-72); next_container_symbol returns 0 for an empty map (infer.rs:640-648), so the floor and thus symbol numbering are unchanged on tensor graphs — no extra fresh_dim minted.
• body_container_seeds early-returns an empty map unless the node is a default-domain Loop with a container operand (infer.rs:625-637).
• tensor_only_path_is_byte_identical_after_container_type_model GREEN in isolation, and all existing If/Loop/Scan tensor tests remain in the 33-test graph_inference suite, all GREEN.

2. If unify semantics + disagreement — PASS, honest

unify_value_type (context.rs, new): Tensor leaves delegate to the foundation unify_tensor_type (genuine reuse, not a reimplementation); Sequence/Optional recurse structurally; Map requires equal key dtype else error. Disagreement is honest:

• dtype mismatch → Invalid error (via unify_tensor_type), attributed to "If" — asserted by if_sequence_branch_dtype_disagreement_errors.
• extent disagreement → fresh symbol (matching dim preserved) — asserted by if_sequence_branch_extent_disagreement_degrades_to_symbol.
• container-vs-tensor branch → explicit Invalid error ("one branch produces a container output while the other produces a tensor", infer.rs:504-511) — asserted by if_container_versus_tensor_branch_disagreement_errors. Consistent with the tensor-branch merge semantics of the foundation.

3. Loop carried-container threading + remap — PASS, sound

• Carried operand 2+i seeds body carried input 2+i via body_container_seeds (infer.rs:625-637); body threads/mutates it (SequenceInsert/Identity/passthrough); carried body-output container flows to the Loop output gated on slot < carried (infer.rs:955-963) — scan outputs (slot >= carried) stay tensors, never containers.
• map_container_to_parent/map_container_dim (infer.rs:673-729) mirror the tensor CF per-dim rule exactly: const→const, parent-origin symbol (in child_to_parent)→pass through, else→fresh parent symbol. dtype and structure recurse unchanged.
• Documented degradation is SOUND: a seeded symbolic element dim never lands on a body Value, so it is absent from child_to_parent and degrades to a fresh (maximally-permissive) parent symbol — information loss only on the symbolic extent; dtype/structure/concrete extents preserved. Aliasing is prevented by raising the child floor above seeded container symbols (next_container_symbol + DimExpr::symbol_ids). Exactly exhibited and asserted by loop_passthrough_preserves_seeded_sequence_dtype (dtype Float16 kept, concrete dim 3 kept, symbolic batch → fresh symbol).

4. Tests meaningful (assert types, not non-crash) — PASS

graph_inference 27→33 (+6). All six assert real container types: matching-branch unify keeps concrete [2,3]; extent-disagreement asserts a symbol in slot 1 with slot 0 concrete; two disagreement tests assert error attribution ("If"/"container"); Loop accumulator asserts element shape [2,3] preserved across SequenceInsert; passthrough asserts dtype + concrete-dim preservation + symbolic degradation. None are non-panic-only.

5. DRY promotion genuine — PASS

Old merge_tensor/merge_shape are DELETED from handlers/container.rs and replaced by calls to context::unify_tensor_type/merge_element_shape (taking &mut SymbolInterner via ctx.interner_mut()); infer.rs's unify_value_type reuses unify_tensor_type at leaves. One implementation, shared by handlers + infer.rs — not a copy. Catalog UNCHANGED: operator_count()==217, entry_count()==262; expanded_registry_catalog_count_is_pinned GREEN (no new ops — threading only).

6. Verification — ALL GREEN

• cargo test -p onnx-runtime-shape-inference: lib 16, graph_inference 33, op_rules 275, doctest 1 — 0 failed. Byte-identical regression GREEN; catalog pin GREEN.
• cargo fmt --all --check: clean.
• cargo clippy -p onnx-runtime-shape-inference --all-targets -- -D warnings: clean (no warnings/errors).
• cargo build -p onnx-runtime-session -p onnx-runtime-eager: Finished OK — container ValueType default-None does not break session/eager (blast radius clear).

Scoped-out items (inc3b Scan container state vars; inc4 SequenceMap + cross-subgraph capture with remap_node_io value_type:None) are documented and consistent with the code. No blocking issues found.

[github-actions]

✅ Benchmarks — No Regression

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
✅	grammar_masking/llguidance_compute_mask/32	74.15 µs	74.98 µs	+1.1%
✅	logit_processing/seven_processor_chain_per_step	1.16 ms	1.16 ms	+0.5%
✅	kv_cache/alloc_dealloc_pages	37.51 µs	36.38 µs	-3.0%
✅	gather/large_bf16_threads=1-internal/131072	14.17 µs	13.46 µs	-5.1%
✅	sampling_latency/top_k_per_token	524.20 µs	463.78 µs	-11.5%
✅	sampling_latency/greedy_per_token	3.61 µs	3.19 µs	-11.7%
✅	sampling_latency/top_p_per_token	1.09 ms	963.77 µs	-11.8%
🟢	tokenization/encode_tokens_per_second	449.02 µs	372.81 µs	-17.0%
🟢	matmul/small_generic_bf16_threads=8/1x256x256	38.54 µs	31.84 µs	-17.4%
🟢	tokenization/decode_tokens_per_second	7.35 ms	6.05 ms	-17.8%
🟢	matmul/small_generic_bf16_threads=1/1x256x256	39.82 µs	32.47 µs	-18.5%
🟢	sampling_latency/min_p_per_token	436.27 µs	346.07 µs	-20.7%
🟢	matmul/large_generic_bf16_threads=8/32x1024x1024	1.65 ms	1.30 ms	-21.3%
🟢	matmul/medium_generic_f32_threads=1/32x512x512	2.93 ms	2.30 ms	-21.5%
🟢	matmul/medium_generic_f32_threads=8/32x512x512	1.24 ms	963.03 µs	-22.5%
🟢	matmul/large_generic_f32_threads=1/32x1024x1024	11.96 ms	9.25 ms	-22.7%
🟢	add/large_f32_threads=1-internal/4194304	54.02 ms	39.92 ms	-26.1%
🟢	matmul/large_generic_bf16_threads=1/32x1024x1024	2.62 ms	1.92 ms	-26.5%
🟢	reduce_mean/large_f32_threads=1-internal/262144	1.35 ms	959.62 µs	-29.1%
🟢	matmul/small_generic_f32_threads=8/1x256x256	59.43 µs	41.95 µs	-29.4%
🟢	matmul/small_generic_f32_threads=1/1x256x256	51.49 µs	35.86 µs	-30.4%
🟢	matmul/large_generic_f16_threads=1/32x1024x1024	114.20 µs	79.50 µs	-30.4%
🟢	add/small_f16_threads=1-internal/1024	19.08 µs	13.17 µs	-31.0%
🟢	add/medium_bf16_threads=1-internal/262144	3.80 ms	2.62 ms	-31.0%
🟢	add/small_bf16_threads=1-internal/1024	18.50 µs	12.74 µs	-31.1%
🟢	add/medium_f16_threads=1-internal/262144	3.70 ms	2.54 ms	-31.3%
🟢	add/large_bf16_threads=1-internal/4194304	61.28 ms	42.03 ms	-31.4%
🟢	gather/small_f32_threads=1-internal/4096	970.4 ns	664.0 ns	-31.6%
🟢	reduce_mean/small_f32_threads=1-internal/4096	21.50 µs	14.64 µs	-31.9%
🟢	add/large_f16_threads=1-internal/4194304	58.93 ms	40.11 ms	-31.9%
🟢	reduce_mean/medium_f32_threads=1-internal/65536	354.47 µs	237.40 µs	-33.0%
🟢	add/medium_f32_threads=1-internal/262144	3.82 ms	2.50 ms	-34.6%
🟢	matmul/small_generic_f16_threads=1/1x256x256	47.78 µs	31.04 µs	-35.0%
🟢	gather/small_bf16_threads=1-internal/4096	736.9 ns	469.0 ns	-36.4%
🟢	gather/small_f16_threads=1-internal/4096	733.6 ns	458.3 ns	-37.5%
🟢	matmul/large_generic_f16_threads=8/32x1024x1024	141.81 µs	87.67 µs	-38.2%
🟢	matmul/medium_generic_bf16_threads=8/32x512x512	618.10 µs	379.01 µs	-38.7%
🟢	matmul/small_generic_f16_threads=8/1x256x256	51.39 µs	31.41 µs	-38.9%
🟢	gather/medium_f16_threads=1-internal/32768	4.13 µs	2.50 µs	-39.4%
🟢	matmul/medium_generic_bf16_threads=1/32x512x512	879.41 µs	514.76 µs	-41.5%
🟢	gather/medium_f32_threads=1-internal/32768	6.79 µs	3.92 µs	-42.4%
🟢	matmul/large_generic_f32_threads=8/32x1024x1024	6.62 ms	3.74 ms	-43.4%
🟢	add/small_f32_threads=1-internal/1024	340.0 ns	189.4 ns	-44.3%
🟢	gather/medium_bf16_threads=1-internal/32768	4.43 µs	2.39 µs	-46.0%
🟢	matmul/medium_generic_f16_threads=8/32x512x512	59.62 µs	30.72 µs	-48.5%
🟢	gather/large_f16_threads=1-internal/131072	23.72 µs	12.02 µs	-49.4%
🟢	matmul/medium_generic_f16_threads=1/32x512x512	61.78 µs	31.26 µs	-49.4%
🟢	gather/large_f32_threads=1-internal/131072	55.98 µs	28.12 µs	-49.8%

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.4.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 3.90 4.68 8.41 }

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)