From bba27cef9815c7d1233260f9f34f2f71d248c95e Mon Sep 17 00:00:00 2001 From: human9000 Date: Wed, 5 Aug 2026 12:13:25 +0500 Subject: [PATCH] onnx: fix two graph-pass deletion bugs, prune the dead runner MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stage 0 of the three-layer optimizer arc: correctness and dead-code prerequisites, so the layers that follow are measured against a clean baseline. No perf win is claimed. remove_dropout_nodes deleted nodes by NodeProto.name, which ONNX makes optional. A single unnamed Dropout put the empty string into the delete set, and `retain(|n| !dropout_names.contains(&n.name))` then removed every unnamed node in the graph — the new fixture remove_dropout_keeps_unnamed_siblings drops from three nodes to zero against the old code. The same block also deleted malformed Dropout nodes that the rewiring loop above had skipped, orphaning their consumers. It now records the indices it actually rewired and removes those, matching the reverse-removal idiom the other passes use. eliminate_squeeze_unsqueeze_pairs accepted overlapping inverse pairs. A Squeeze -> Unsqueeze -> Squeeze chain with matching axes satisfies the predicate at both i and i+1, and the reverse-removal loop then deleted already-shifted indices — in the new eliminate_squeeze_unsqueeze_handles_overlapping_chain fixture the downstream Relu disappears entirely. A next_free cursor now skips matches that overlap an accepted pair. Nine of the twelve passes called rebuild_runtime_index() on exit, so one optimize_onnx_graph re-ran execution-plan construction and weight prepacking about ten times over. Each pass now exposes a `run` that mutates without rebuilding and reports whether it changed anything; the driver calls those and rebuilds once at the end. The public per-pass entry points keep rebuilding, because a stale plan holds node indices and yscv-quantize-cli and llm-bench call several passes standalone. fold_conv_mul and fold_conv_add_const lose their public wrappers — the driver was their only caller and they were never re-exported from lib.rs. run_onnx_model_sequential is gone: ~600 lines re-deriving Conv+BN+Relu, Conv+SiLU and Conv+Add patterns on every inference. It was reachable only when the load-time execution plan was empty, which no graph with nodes produces; confirmed by panicking in that branch and running the workspace suite (53 binaries, zero hits). Its orphaned helpers went with it — find_relu_after_identity_chain, mark_skip_indices, exec_reshape_zerocopy, and the use_counts-taking try_reshape_nhwc_passthrough. The plan path's _inner variant and the NHWC-passthrough optimization are untouched. TensorEnv::get_mut survives behind #[cfg(feature = "gpu")], its only remaining callers. Validation: cargo fmt, clippy --workspace --all-targets --all-features -D warnings, clippy with `gpu` and with `gpu rknn native-camera`, cargo test --workspace (2305 passed), check-doc-counts.sh (all 14 rows match; onnx-cpu-ops still 122). The Windows cross-check could not run locally — x86_64-pc-windows-msvc std is not installed; nothing here is platform-specific and CI covers it. --- CHANGELOG.md | 34 + .../eliminate_squeeze_unsqueeze_pairs.rs | 10 +- .../yscv-onnx/src/optimizer/fold_constants.rs | 14 +- .../src/optimizer/fold_conv_add_const.rs | 8 +- .../yscv-onnx/src/optimizer/fold_conv_bn.rs | 12 +- .../yscv-onnx/src/optimizer/fold_conv_mul.rs | 8 +- .../yscv-onnx/src/optimizer/fuse_bn_relu.rs | 10 +- .../yscv-onnx/src/optimizer/fuse_conv_relu.rs | 10 +- crates/yscv-onnx/src/optimizer/mod.rs | 26 +- .../src/optimizer/remove_dropout_nodes.rs | 21 +- .../strip_qdq_within_fusion_chains.rs | 12 +- crates/yscv-onnx/src/runner/execute.rs | 603 ------------------ crates/yscv-onnx/src/runner/layout.rs | 115 +--- crates/yscv-onnx/src/runner/mod.rs | 49 +- crates/yscv-onnx/src/runner/reshape.rs | 56 +- crates/yscv-onnx/src/runner/tensor_env.rs | 4 + crates/yscv-onnx/src/tests/optimizer.rs | 108 ++++ 17 files changed, 270 insertions(+), 830 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0bae977..bbaf491 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,40 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed + +- `remove_dropout_nodes` deleted nodes by `NodeProto.name`, which ONNX makes + optional. A single unnamed `Dropout` put the empty string into the delete set + and took every other unnamed node in the graph with it — on a fully unnamed + three-node fixture the pass emptied the graph. It also deleted malformed + `Dropout` nodes it had skipped rewiring. Now removes only the nodes it + actually rewired, by index. +- `eliminate_squeeze_unsqueeze_pairs` could accept overlapping inverse pairs + `(i, i+1)` and `(i+1, i+2)` from a chain such as `Squeeze -> Unsqueeze -> + Squeeze`. The reverse-removal loop then deleted already-shifted indices, + taking unrelated nodes with them. Overlapping matches are now skipped. + +### Changed + +- Optimizer passes no longer rebuild the runtime index individually. Nine of + the twelve passes called `rebuild_runtime_index()` on exit, so a single + `optimize_onnx_graph` re-ran execution-plan construction and weight + prepacking about ten times. The driver now rebuilds once after the whole + pipeline; the public per-pass entry points still rebuild for standalone + callers. + +### Removed + +- `run_onnx_model_sequential`, the ~600-line per-inference fusion scanner in + `runner/execute.rs`. It re-derived Conv+BN+Relu / Conv+SiLU / Conv+Add + patterns on every inference and was reachable only when the load-time + execution plan was empty, which no graph with nodes produces. Verified + unreachable across the workspace suite before removal. Its now-orphaned + helpers went with it: `find_relu_after_identity_chain`, `mark_skip_indices`, + `exec_reshape_zerocopy`, and the `use_counts`-taking + `try_reshape_nhwc_passthrough` (the plan path's `_inner` variant, and the + NHWC-passthrough optimization itself, are unaffected). + ## [0.1.11] — 2026-07-26 ### Added diff --git a/crates/yscv-onnx/src/optimizer/eliminate_squeeze_unsqueeze_pairs.rs b/crates/yscv-onnx/src/optimizer/eliminate_squeeze_unsqueeze_pairs.rs index dd6dc6f..545be70 100644 --- a/crates/yscv-onnx/src/optimizer/eliminate_squeeze_unsqueeze_pairs.rs +++ b/crates/yscv-onnx/src/optimizer/eliminate_squeeze_unsqueeze_pairs.rs @@ -6,8 +6,16 @@ use crate::loader::OnnxNode; /// Squeeze(Unsqueeze(x, axes=A), axes=A) and Unsqueeze(Squeeze(x, axes=A), axes=A). pub fn eliminate_squeeze_unsqueeze_pairs(model: &mut OnnxModel) { let mut remove_pairs: Vec<(usize, usize, String, String)> = Vec::new(); + // Index of the first node not already claimed by an accepted pair. A chain + // like Squeeze -> Unsqueeze -> Squeeze matches at both `i` and `i + 1`; + // accepting both would make the reverse-removal loop below delete shifted + // indices, taking unrelated nodes with them. + let mut next_free = 0usize; for i in 0..model.nodes.len().saturating_sub(1) { + if i < next_free { + continue; + } let first = &model.nodes[i]; let second = &model.nodes[i + 1]; let is_pair = (first.op_type == "Squeeze" && second.op_type == "Unsqueeze") @@ -38,6 +46,7 @@ pub fn eliminate_squeeze_unsqueeze_pairs(model: &mut OnnxModel) { continue; } remove_pairs.push((i, i + 1, first.inputs[0].clone(), second.outputs[0].clone())); + next_free = i + 2; } for &(first_idx, second_idx, ref producer_input, ref consumer_output) in @@ -58,7 +67,6 @@ pub fn eliminate_squeeze_unsqueeze_pairs(model: &mut OnnxModel) { model.nodes.remove(second_idx); model.nodes.remove(first_idx); } - model.rebuild_runtime_index(); } fn node_axes(model: &OnnxModel, node: &OnnxNode) -> Option> { diff --git a/crates/yscv-onnx/src/optimizer/fold_constants.rs b/crates/yscv-onnx/src/optimizer/fold_constants.rs index 81c7a97..a86a265 100644 --- a/crates/yscv-onnx/src/optimizer/fold_constants.rs +++ b/crates/yscv-onnx/src/optimizer/fold_constants.rs @@ -10,6 +10,17 @@ use crate::runner::run_onnx_model; /// /// Iterates until a fixed point is reached (no more foldable nodes). pub fn fold_constants(model: &mut OnnxModel) { + if run(model) { + model.rebuild_runtime_index(); + } +} + +/// Folds without rebuilding the caller's runtime index; returns whether the +/// graph changed. The driver calls this and rebuilds once for the whole +/// pipeline. The throwaway `mini_model` built per folded node still rebuilds +/// its own index — it has to, since it is executed to produce the constant. +pub(super) fn run(model: &mut OnnxModel) -> bool { + let mut changed = false; loop { let foldable = model.nodes.iter().enumerate().find(|(_, node)| { !node.inputs.is_empty() @@ -85,11 +96,12 @@ pub fn fold_constants(model: &mut OnnxModel) { model.initializers.insert(name, tensor); } model.nodes.remove(idx); + changed = true; } Err(_) => { break; } } } - model.rebuild_runtime_index(); + changed } diff --git a/crates/yscv-onnx/src/optimizer/fold_conv_add_const.rs b/crates/yscv-onnx/src/optimizer/fold_conv_add_const.rs index 3fe5746..4af44c0 100644 --- a/crates/yscv-onnx/src/optimizer/fold_conv_add_const.rs +++ b/crates/yscv-onnx/src/optimizer/fold_conv_add_const.rs @@ -10,7 +10,10 @@ use super::fold_conv_mul::broadcast_scale_to_oc; /// Residual (non-constant) Adds are handled at runtime by `NodeAction::ConvAdd` /// and are NOT touched here — we only fold when the second Add operand is a /// graph initializer broadcastable to OC. -pub fn fold_conv_add_const(model: &mut OnnxModel) { +/// Folds without rebuilding the runtime index; returns whether the graph +/// changed. The driver rebuilds once for the whole pipeline. +pub(super) fn run(model: &mut OnnxModel) -> bool { + let mut changed = false; let mut fuse_pairs: Vec<(usize, usize, usize)> = Vec::new(); for i in 0..model.nodes.len().saturating_sub(1) { @@ -119,6 +122,7 @@ pub fn fold_conv_add_const(model: &mut OnnxModel) { model.nodes[conv_idx].outputs[0] = add_out; model.nodes.remove(add_idx); + changed = true; } - model.rebuild_runtime_index(); + changed } diff --git a/crates/yscv-onnx/src/optimizer/fold_conv_bn.rs b/crates/yscv-onnx/src/optimizer/fold_conv_bn.rs index 547f349..6e7b7ee 100644 --- a/crates/yscv-onnx/src/optimizer/fold_conv_bn.rs +++ b/crates/yscv-onnx/src/optimizer/fold_conv_bn.rs @@ -11,6 +11,15 @@ use crate::loader::OnnxModel; /// `b_fused[c] = (b[c] - mean_c) * scale_c + beta_c` /// The Conv initializers are replaced and the BN node is removed. pub fn fold_conv_bn(model: &mut OnnxModel) { + if run(model) { + model.rebuild_runtime_index(); + } +} + +/// Folds without rebuilding the runtime index; returns whether the graph +/// changed. The driver calls this and rebuilds once for the whole pipeline. +pub(super) fn run(model: &mut OnnxModel) -> bool { + let mut changed = false; let mut fuse_pairs: Vec<(usize, usize)> = Vec::new(); for i in 0..model.nodes.len().saturating_sub(1) { @@ -166,6 +175,7 @@ pub fn fold_conv_bn(model: &mut OnnxModel) { model.nodes[conv_idx].outputs[0] = bn_output; model.nodes.remove(bn_idx); + changed = true; } - model.rebuild_runtime_index(); + changed } diff --git a/crates/yscv-onnx/src/optimizer/fold_conv_mul.rs b/crates/yscv-onnx/src/optimizer/fold_conv_mul.rs index b7917ff..e0705ed 100644 --- a/crates/yscv-onnx/src/optimizer/fold_conv_mul.rs +++ b/crates/yscv-onnx/src/optimizer/fold_conv_mul.rs @@ -64,7 +64,10 @@ fn scale_weight_inplace( /// Conv output is only consumed by the Mul. After fusion: /// W'[c] = W[c] · s[c] /// b'[c] = b[c] · s[c] -pub fn fold_conv_mul(model: &mut OnnxModel) { +/// Folds without rebuilding the runtime index; returns whether the graph +/// changed. The driver rebuilds once for the whole pipeline. +pub(super) fn run(model: &mut OnnxModel) -> bool { + let mut changed = false; let mut fuse_pairs: Vec<(usize, usize, usize)> = Vec::new(); for i in 0..model.nodes.len().saturating_sub(1) { @@ -171,6 +174,7 @@ pub fn fold_conv_mul(model: &mut OnnxModel) { model.nodes[conv_idx].outputs[0] = mul_out; model.nodes.remove(mul_idx); + changed = true; } - model.rebuild_runtime_index(); + changed } diff --git a/crates/yscv-onnx/src/optimizer/fuse_bn_relu.rs b/crates/yscv-onnx/src/optimizer/fuse_bn_relu.rs index c38b496..0273ad0 100644 --- a/crates/yscv-onnx/src/optimizer/fuse_bn_relu.rs +++ b/crates/yscv-onnx/src/optimizer/fuse_bn_relu.rs @@ -4,6 +4,14 @@ use crate::loader::OnnxModel; /// (annotation-only; execution still handles them separately, but this /// reduces graph traversal overhead for large models). pub fn fuse_bn_relu(model: &mut OnnxModel) { + if run(model) { + model.rebuild_runtime_index(); + } +} + +/// Fuses without rebuilding the runtime index; returns whether the graph +/// changed. The driver calls this and rebuilds once for the whole pipeline. +pub(super) fn run(model: &mut OnnxModel) -> bool { let mut fuse_pairs: Vec<(usize, usize)> = Vec::new(); for i in 0..model.nodes.len().saturating_sub(1) { if model.nodes[i].op_type == "BatchNormalization" @@ -31,5 +39,5 @@ pub fn fuse_bn_relu(model: &mut OnnxModel) { model.nodes[bn_idx].op_type = "BatchNormalization_Relu".to_string(); model.nodes.remove(relu_idx); } - model.rebuild_runtime_index(); + !fuse_pairs.is_empty() } diff --git a/crates/yscv-onnx/src/optimizer/fuse_conv_relu.rs b/crates/yscv-onnx/src/optimizer/fuse_conv_relu.rs index 50e0785..0498bd6 100644 --- a/crates/yscv-onnx/src/optimizer/fuse_conv_relu.rs +++ b/crates/yscv-onnx/src/optimizer/fuse_conv_relu.rs @@ -2,6 +2,14 @@ use crate::loader::OnnxModel; /// Fuse `Conv` -> `Relu` node pairs into a single fused-activation Conv. pub fn fuse_conv_relu(model: &mut OnnxModel) { + if run(model) { + model.rebuild_runtime_index(); + } +} + +/// Fuses without rebuilding the runtime index; returns whether the graph +/// changed. The driver calls this and rebuilds once for the whole pipeline. +pub(super) fn run(model: &mut OnnxModel) -> bool { let mut fuse_pairs: Vec<(usize, usize)> = Vec::new(); for i in 0..model.nodes.len().saturating_sub(1) { if model.nodes[i].op_type == "Conv" @@ -29,5 +37,5 @@ pub fn fuse_conv_relu(model: &mut OnnxModel) { model.nodes[conv_idx].op_type = "Conv_Relu".to_string(); model.nodes.remove(relu_idx); } - model.rebuild_runtime_index(); + !fuse_pairs.is_empty() } diff --git a/crates/yscv-onnx/src/optimizer/mod.rs b/crates/yscv-onnx/src/optimizer/mod.rs index daee383..1aba524 100644 --- a/crates/yscv-onnx/src/optimizer/mod.rs +++ b/crates/yscv-onnx/src/optimizer/mod.rs @@ -26,9 +26,7 @@ use crate::{ pub use analyze_nchwc::analyze_nchwc; pub use fold_constants::fold_constants; -pub use fold_conv_add_const::fold_conv_add_const; pub use fold_conv_bn::fold_conv_bn; -pub use fold_conv_mul::fold_conv_mul; pub use fuse_bn_relu::fuse_bn_relu; pub use fuse_conv_relu::fuse_conv_relu; pub use graph_cost::{ @@ -41,6 +39,9 @@ pub use strip_qdq_within_fusion_chains::strip_qdq_within_fusion_chains; /// Optimizes an ONNX model graph in-place for inference. /// /// Applies load-time passes modeled after ORT's Level-1 optimizer: +/// - Node reordering (topological re-sort so producers sit immediately before +/// consumers; the fusion passes below match on adjacency, so this is a +/// correctness prerequisite for them, not a nicety) /// - Dropout removal (inference-only, rewire consumers to Dropout input) /// - ConvTranspose(k==s) rewrite to Conv1x1 + DepthToSpace (GEMM-backed path, /// also unlocks backends without a ConvTranspose kernel) @@ -49,25 +50,28 @@ pub use strip_qdq_within_fusion_chains::strip_qdq_within_fusion_chains; /// - Conv-Add(const) bias absorption (absorb per-channel Add into Conv bias) /// - Constant folding (execute nodes with all-initializer inputs at load) /// - Squeeze/Unsqueeze pair elimination (drop inverse pairs left by PyTorch export) -/// - Conv+Clip(0,max) fusion (ReLU6-style clamped activation) /// - Conv+Relu / BN+Relu fusion (annotation-only; kernel dispatches on op_type) /// - Dead code elimination (iterate to fixpoint) /// /// Order matters: `fold_conv_bn` runs before Conv-Mul/Conv-Add because BN /// usually absorbs into Conv already; only stray scale/bias left over fall -/// through. Constant folding runs before Relu/Clip fusions because folded -/// Relus may turn Clip-style patterns into plain activations. +/// through. Constant folding runs before the activation fusions because a +/// folded constant can turn a data-dependent activation into a plain one. +/// +/// The runtime index is rebuilt once here, after every pass has run. Individual +/// passes must not rebuild it themselves — doing so re-runs plan construction +/// and weight prepacking once per pass. pub fn optimize_onnx_graph(model: &mut OnnxModel) { reorder_nodes_for_fusion(model); remove_dropout_nodes(model); rewrite_convtranspose_dts(model); - fold_conv_bn(model); - fold_conv_mul(model); - fold_conv_add_const(model); - fold_constants(model); + fold_conv_bn::run(model); + fold_conv_mul::run(model); + fold_conv_add_const::run(model); + fold_constants::run(model); eliminate_squeeze_unsqueeze_pairs(model); - fuse_conv_relu(model); - fuse_bn_relu(model); + fuse_conv_relu::run(model); + fuse_bn_relu::run(model); eliminate_dead_code(model); model.rebuild_runtime_index(); diff --git a/crates/yscv-onnx/src/optimizer/remove_dropout_nodes.rs b/crates/yscv-onnx/src/optimizer/remove_dropout_nodes.rs index 22b0dd3..b3b30ac 100644 --- a/crates/yscv-onnx/src/optimizer/remove_dropout_nodes.rs +++ b/crates/yscv-onnx/src/optimizer/remove_dropout_nodes.rs @@ -1,14 +1,19 @@ -use rustc_hash::FxHashSet; - use crate::loader::OnnxModel; /// Removes Dropout nodes by rewiring their consumers to the Dropout's input. pub fn remove_dropout_nodes(model: &mut OnnxModel) { let mut rewire: Vec<(String, String)> = Vec::new(); + // Indices of the Dropouts we actually rewired, ascending. Deleting by node + // name instead would be wrong twice over: `NodeProto.name` is optional in + // ONNX, so one unnamed Dropout would take every other unnamed node in the + // graph with it, and a malformed Dropout that was skipped above would be + // deleted without its consumers ever being rewired. + let mut drop_indices: Vec = Vec::new(); - for node in &model.nodes { + for (idx, node) in model.nodes.iter().enumerate() { if node.op_type == "Dropout" && !node.inputs.is_empty() && !node.outputs.is_empty() { rewire.push((node.outputs[0].clone(), node.inputs[0].clone())); + drop_indices.push(idx); } } @@ -31,11 +36,7 @@ pub fn remove_dropout_nodes(model: &mut OnnxModel) { } } - let dropout_names: FxHashSet = model - .nodes - .iter() - .filter(|n| n.op_type == "Dropout") - .map(|n| n.name.clone()) - .collect(); - model.nodes.retain(|n| !dropout_names.contains(&n.name)); + for &idx in drop_indices.iter().rev() { + model.nodes.remove(idx); + } } diff --git a/crates/yscv-onnx/src/optimizer/strip_qdq_within_fusion_chains.rs b/crates/yscv-onnx/src/optimizer/strip_qdq_within_fusion_chains.rs index b5463a3..488dbf0 100644 --- a/crates/yscv-onnx/src/optimizer/strip_qdq_within_fusion_chains.rs +++ b/crates/yscv-onnx/src/optimizer/strip_qdq_within_fusion_chains.rs @@ -18,6 +18,17 @@ use crate::loader::{OnnxModel, OnnxNode}; /// /// Returns the number of Q+DQ pairs removed. pub fn strip_qdq_within_fusion_chains(model: &mut OnnxModel) -> usize { + let removed = run(model); + if removed != 0 { + model.rebuild_runtime_index(); + } + removed +} + +/// Strips without rebuilding the runtime index; returns the number of pairs +/// removed. This pass is not in the driver pipeline — the quantization tools +/// call the public wrapper above — but it follows the same contract. +pub(super) fn run(model: &mut OnnxModel) -> usize { let conv_like = |op: &str| matches!(op, "Conv" | "Conv_Relu" | "MatMul" | "Gemm"); let mut to_remove: Vec = Vec::new(); @@ -96,7 +107,6 @@ pub fn strip_qdq_within_fusion_chains(model: &mut OnnxModel) -> usize { for &idx in to_remove.iter().rev() { model.nodes.remove(idx); } - model.rebuild_runtime_index(); removed } diff --git a/crates/yscv-onnx/src/runner/execute.rs b/crates/yscv-onnx/src/runner/execute.rs index bfa86c0..b540b64 100644 --- a/crates/yscv-onnx/src/runner/execute.rs +++ b/crates/yscv-onnx/src/runner/execute.rs @@ -183,606 +183,3 @@ pub(crate) fn run_onnx_model_jit( } Ok(result) } - -pub(crate) fn run_onnx_model_sequential( - model: &OnnxModel, - mut env: TensorEnv<'_, '_>, -) -> Result, OnnxError> { - // --- Operator fusion: scan for fusible patterns --- - // Build a set of node indices that should be skipped because they were - // fused into the preceding node. We also create synthetic "fused" nodes - // that carry a combined op_type (e.g. "Conv_Relu"). - let nodes = &model.nodes; - let node_kinds = &model.runtime_index.node_kinds; - let mut skip = vec![false; nodes.len()]; - - // Build reference counts: how many nodes consume each tensor as input. - // Used by SiLU fusions to decide in-place vs allocating path. - let use_counts = &model.runtime_index.use_counts; - let use_counts_by_id = &model.runtime_index.use_counts_by_id; - - // Mutable remaining-use counters for early tensor deallocation. - // When a tensor's remaining uses reach zero, free it to reduce working set. - let mut remaining_uses: Vec = use_counts_by_id.clone(); - let output_id_mask = build_output_id_mask(model, &env, use_counts_by_id.len()); - - for (i, node) in nodes.iter().enumerate() { - if skip[i] { - continue; - } - let kind = node_kind(node_kinds, nodes, i); - - // --- Conv → BatchNorm → Relu 3-node fusion --- - if kind == NodeKind::Conv - && let Some(next) = nodes.get(i + 1) - && node_kind(node_kinds, nodes, i + 1) == NodeKind::BatchNormalization - && !next.inputs.is_empty() - && next.inputs[0] == node.outputs[0] - && let Some((relu_idx, identity_idxs)) = - find_relu_after_identity_chain(nodes, node_kinds, i + 2, &next.outputs[0]) - { - execute_node_with_layout_kind(node, &mut env, kind)?; - execute_node_with_layout_kind(next, &mut env, node_kind(node_kinds, nodes, i + 1))?; - if let Some(tensor) = env.get_mut(&next.outputs[0]) { - relu_inplace(tensor); - } - let source = &next.outputs[0]; - for &id_idx in &identity_idxs { - env.alias(&nodes[id_idx].outputs[0], source); - } - env.alias(&nodes[relu_idx].outputs[0], source); - if i + 1 < skip.len() { - skip[i + 1] = true; - } - mark_skip_indices(&mut skip, &identity_idxs); - mark_skip_indices(&mut skip, &[relu_idx]); - continue; - } - - // --- Conv + Relu fusion --- - if kind == NodeKind::Conv - && let Some((relu_idx, identity_idxs)) = - find_relu_after_identity_chain(nodes, node_kinds, i + 1, &node.outputs[0]) - { - exec_conv(node, &mut env, yscv_kernels::Activation::Relu)?; - let source = &node.outputs[0]; - for &id_idx in &identity_idxs { - env.alias(&nodes[id_idx].outputs[0], source); - } - env.alias(&nodes[relu_idx].outputs[0], source); - mark_skip_indices(&mut skip, &identity_idxs); - mark_skip_indices(&mut skip, &[relu_idx]); - continue; - } - - // --- Conv + SiLU fusion (Conv → Sigmoid → Mul) --- - // Detect Sigmoid at i+1 and Mul at i+2 that form SiLU on Conv output. - if kind == NodeKind::Conv { - let conv_out = &node.outputs[0]; - // Look for Sigmoid(conv_out) → Mul(conv_out, sigmoid_out) pattern - let mut silu_mul_idx = None; - for sig_offset in 1..=2 { - if let Some(sig) = nodes.get(i + sig_offset) - && node_kind(node_kinds, nodes, i + sig_offset) == NodeKind::Sigmoid - && sig.inputs.len() == 1 - && sig.inputs[0] == *conv_out - { - let sig_out = &sig.outputs[0]; - for mul_offset in (sig_offset + 1)..=(sig_offset + 2) { - if let Some(mul) = nodes.get(i + mul_offset) - && node_kind(node_kinds, nodes, i + mul_offset) == NodeKind::Mul - && mul.inputs.len() == 2 - && ((mul.inputs[0] == *sig_out && mul.inputs[1] == *conv_out) - || (mul.inputs[1] == *sig_out && mul.inputs[0] == *conv_out)) - { - silu_mul_idx = Some((sig_offset, mul_offset, mul.outputs[0].clone())); - break; - } - } - if silu_mul_idx.is_some() { - break; - } - } - } - if let Some((sig_off, mul_off, mul_out)) = silu_mul_idx { - let conv_out_uses = tensor_use_count(&env, use_counts_by_id, use_counts, conv_out); - if conv_out_uses <= 2 { - // Fuse SiLU into Conv GEMM tiles (applied cache-hot after bias). - exec_conv(node, &mut env, yscv_kernels::Activation::Silu)?; - env.alias(&mul_out, conv_out); - } else { - // Other consumers need raw conv_out — can't fuse. - execute_node_with_layout_kind(node, &mut env, kind)?; - if let Some(tensor) = env.get(conv_out) { - let result = yscv_kernels::silu(tensor); - env.insert(mul_out.clone(), result); - } - } - let is_nhwc = env.is_nhwc(conv_out); - if is_nhwc { - env.mark_nhwc(&mul_out); - } - // Execute any intermediate nodes between Conv and Sigmoid, - // then mark them as done so the main loop doesn't re-execute them. - for mid in 1..sig_off { - if i + mid < skip.len() && !skip[i + mid] { - execute_node_with_layout_kind( - &nodes[i + mid], - &mut env, - node_kind(node_kinds, nodes, i + mid), - )?; - skip[i + mid] = true; - } - } - if i + sig_off < skip.len() { - skip[i + sig_off] = true; - } - if i + mul_off < skip.len() { - skip[i + mul_off] = true; - } - continue; - } - } - - // --- Conv + Add (residual connection) fusion --- - // Pattern: Conv → Add(conv_out, skip_connection), optionally → Relu - // Reuses conv_out buffer for the result, avoiding allocation. - if kind == NodeKind::Conv - && let Some(next) = nodes.get(i + 1) - && node_kind(node_kinds, nodes, i + 1) == NodeKind::Add - && next.inputs.len() == 2 - && (next.inputs[0] == node.outputs[0] || next.inputs[1] == node.outputs[0]) - { - let conv_out = &node.outputs[0]; - let skip_idx = if &next.inputs[0] == conv_out { 1 } else { 0 }; - let skip_name = &next.inputs[skip_idx]; - let conv_out_uses = tensor_use_count(&env, use_counts_by_id, use_counts, conv_out); - - // Only fuse if conv_out has 2 uses (Add is its only other consumer besides the - // initializer lookup that may happen). If it has more uses, we need to keep it - // for other consumers. - if conv_out_uses <= 2 { - execute_node_with_layout_kind(node, &mut env, kind)?; - - // Capture NHWC flag before remove (remove clears it). - let is_nhwc = env.is_nhwc(conv_out); - - // Add skip_connection in-place to conv_out - if let Some(mut conv_tensor) = env.remove(conv_out) { - if let Some(skip_tensor) = env.get(skip_name) { - yscv_kernels::add_inplace(&mut conv_tensor, skip_tensor); - let add_out = &next.outputs[0]; - - // Check if Relu follows Add - if let Some((relu_idx, identity_idxs)) = - find_relu_after_identity_chain(nodes, node_kinds, i + 2, add_out) - { - relu_inplace(&mut conv_tensor); - env.insert(add_out.clone(), conv_tensor); - if is_nhwc { - env.mark_nhwc(add_out); - } - let source = add_out; - for &id_idx in &identity_idxs { - env.alias(&nodes[id_idx].outputs[0], source); - } - env.alias(&nodes[relu_idx].outputs[0], source); - mark_skip_indices(&mut skip, &identity_idxs); - mark_skip_indices(&mut skip, &[relu_idx]); - } else { - env.insert(add_out.clone(), conv_tensor); - if is_nhwc { - env.mark_nhwc(add_out); - } - } - - if i + 1 < skip.len() { - skip[i + 1] = true; - } - continue; - } else { - env.insert(conv_out.clone(), conv_tensor); - } - } - } - } - - // --- BatchNormalization + Relu fusion --- - if kind == NodeKind::BatchNormalization - && let Some((relu_idx, identity_idxs)) = - find_relu_after_identity_chain(nodes, node_kinds, i + 1, &node.outputs[0]) - { - execute_node_with_layout_kind(node, &mut env, kind)?; - if let Some(tensor) = env.get_mut(&node.outputs[0]) { - relu_inplace(tensor); - } - let source = &node.outputs[0]; - for &id_idx in &identity_idxs { - env.alias(&nodes[id_idx].outputs[0], source); - } - env.alias(&nodes[relu_idx].outputs[0], source); - mark_skip_indices(&mut skip, &identity_idxs); - mark_skip_indices(&mut skip, &[relu_idx]); - continue; - } - - // --- Gemm + Relu fusion --- - if kind == NodeKind::Gemm - && let Some((relu_idx, identity_idxs)) = - find_relu_after_identity_chain(nodes, node_kinds, i + 1, &node.outputs[0]) - { - execute_node_with_layout_kind(node, &mut env, kind)?; - if let Some(tensor) = env.get_mut(&node.outputs[0]) { - relu_inplace(tensor); - } - let source = &node.outputs[0]; - for &id_idx in &identity_idxs { - env.alias(&nodes[id_idx].outputs[0], source); - } - env.alias(&nodes[relu_idx].outputs[0], source); - mark_skip_indices(&mut skip, &identity_idxs); - mark_skip_indices(&mut skip, &[relu_idx]); - continue; - } - - // --- Add + Relu fusion (with in-place Add when possible) --- - if kind == NodeKind::Add - && node.inputs.len() == 2 - && let Some((relu_idx, identity_idxs)) = - find_relu_after_identity_chain(nodes, node_kinds, i + 1, &node.outputs[0]) - { - let a_nhwc = env.is_nhwc(&node.inputs[0]); - let b_nhwc = env.is_nhwc(&node.inputs[1]); - let same_shape_nhwc = a_nhwc == b_nhwc - && match (env.get(&node.inputs[0]), env.get(&node.inputs[1])) { - (Some(a), Some(b)) => a.shape() == b.shape(), - _ => false, - }; - let mut did_inplace = false; - if same_shape_nhwc { - let a_uses = tensor_use_count(&env, use_counts_by_id, use_counts, &node.inputs[0]); - let b_uses = tensor_use_count(&env, use_counts_by_id, use_counts, &node.inputs[1]); - if a_uses <= 1 || b_uses <= 1 { - let (consume_idx, other_idx) = if a_uses <= 1 { (0, 1) } else { (1, 0) }; - if let Some(mut target) = env.remove(&node.inputs[consume_idx]) { - if let Some(other) = env.get(&node.inputs[other_idx]) { - yscv_kernels::add_relu_inplace(&mut target, other); - env.insert(node.outputs[0].clone(), target); - if a_nhwc { - env.mark_nhwc(&node.outputs[0]); - } - did_inplace = true; - } else { - env.insert(node.inputs[consume_idx].clone(), target); - } - } - } - } - if !did_inplace { - execute_node_with_layout_kind(node, &mut env, kind)?; - if let Some(tensor) = env.get_mut(&node.outputs[0]) { - relu_inplace(tensor); - } - } - let source = &node.outputs[0]; - for &id_idx in &identity_idxs { - env.alias(&nodes[id_idx].outputs[0], source); - } - env.alias(&nodes[relu_idx].outputs[0], source); - mark_skip_indices(&mut skip, &identity_idxs); - mark_skip_indices(&mut skip, &[relu_idx]); - continue; - } - - // --- MatMul + Add fusion (Gemm-like) --- - if kind == NodeKind::MatMul - && let Some(next) = nodes.get(i + 1) - && node_kind(node_kinds, nodes, i + 1) == NodeKind::Add - && next.inputs.len() == 2 - && (next.inputs[0] == node.outputs[0] || next.inputs[1] == node.outputs[0]) - { - execute_node_with_layout_kind(node, &mut env, kind)?; - execute_node_with_layout_kind(next, &mut env, node_kind(node_kinds, nodes, i + 1))?; - if i + 1 < skip.len() { - skip[i + 1] = true; - } - continue; - } - - // --- Sigmoid + Mul → SiLU fusion --- - // SiLU(x) = x * sigmoid(x). Pattern: Sigmoid(x)->y, Mul(x,y)->z - // Single-pass SIMD kernel avoids separate sigmoid allocation + Mul dispatch. - if kind == NodeKind::Sigmoid && node.inputs.len() == 1 { - let sig_in = &node.inputs[0]; - let sig_out = &node.outputs[0]; - // Look ahead up to 3 positions for a matching Mul (SiLU pattern). - let mut found_silu = false; - for look in 1..=3 { - if let Some(next) = nodes.get(i + look) - && node_kind(node_kinds, nodes, i + look) == NodeKind::Mul - && next.inputs.len() == 2 - { - let is_silu = (next.inputs[0] == *sig_out && next.inputs[1] == *sig_in) - || (next.inputs[1] == *sig_out && next.inputs[0] == *sig_in); - if is_silu { - let is_nhwc = env.is_nhwc(sig_in); - let mul_out = &next.outputs[0]; - // sig_in is used by Sigmoid + Mul = 2 fused consumers. - // Only remove if no other node needs it. - let sig_in_uses = - tensor_use_count(&env, use_counts_by_id, use_counts, sig_in); - if sig_in_uses <= 2 { - if let Some(mut tensor) = env.remove(sig_in) { - yscv_kernels::silu_inplace(&mut tensor); - env.insert(mul_out.clone(), tensor); - } - } else if let Some(x_tensor) = env.get(sig_in) { - let result_tensor = yscv_kernels::silu(x_tensor); - env.insert(mul_out.clone(), result_tensor); - } - if is_nhwc { - env.mark_nhwc(mul_out); - } - // Execute any intermediate nodes, then mark them done - // so the main loop doesn't re-execute them. - for mid in 1..look { - if i + mid < skip.len() && !skip[i + mid] { - execute_node_with_layout_kind( - &nodes[i + mid], - &mut env, - node_kind(node_kinds, nodes, i + mid), - )?; - skip[i + mid] = true; - } - } - if i + look < skip.len() { - skip[i + look] = true; - } - found_silu = true; - break; - } - } - } - if found_silu { - continue; - } - } - - // --- In-place Add: reuse buffer when one input is last-use --- - if kind == NodeKind::Add && node.inputs.len() == 2 { - let a_nhwc = env.is_nhwc(&node.inputs[0]); - let b_nhwc = env.is_nhwc(&node.inputs[1]); - if a_nhwc == b_nhwc { - let same_shape = match (env.get(&node.inputs[0]), env.get(&node.inputs[1])) { - (Some(a), Some(b)) => a.shape() == b.shape(), - _ => false, - }; - if same_shape { - let a_uses = - tensor_use_count(&env, use_counts_by_id, use_counts, &node.inputs[0]); - let b_uses = - tensor_use_count(&env, use_counts_by_id, use_counts, &node.inputs[1]); - if a_uses <= 1 || b_uses <= 1 { - let (consume_idx, other_idx) = if a_uses <= 1 { (0, 1) } else { (1, 0) }; - if let Some(mut target) = env.remove(&node.inputs[consume_idx]) { - if let Some(other) = env.get(&node.inputs[other_idx]) { - yscv_kernels::add_inplace(&mut target, other); - env.insert(node.outputs[0].clone(), target); - if a_nhwc { - env.mark_nhwc(&node.outputs[0]); - } - continue; - } - env.insert(node.inputs[consume_idx].clone(), target); - } - } - } - } - } - - // Zero-copy Reshape: avoid data clone when the data input has only - // one consumer (this Reshape node). - if kind == NodeKind::Reshape { - // NHWC-passthrough fast path. When the input is NHWC - // physical rank-4 `[N,H,W,C]` and the model reshapes to a rank-3 - // `[N, C, H*W]` (merge-spatial), the data layout is already in - // `[N, H*W, C]` order — which is exactly what a downstream - // `Transpose(perm=[0,2,1])+MatMul` (FusedTransposeMatMul) wants - // as the post-transpose A operand. Skip the `ensure_nchw` permute - // and propagate the NHWC tag. - // Kill switch: `YSCV_RESHAPE_NHWC_PASSTHROUGH_OFF=1`. - if !reshape_nhwc_passthrough_disabled() - && try_reshape_nhwc_passthrough(node, &mut env, use_counts)? - { - continue; - } - for name in &node.inputs { - if !name.is_empty() { - ensure_nchw(&mut env, name)?; - } - } - exec_reshape_zerocopy(node, &mut env, use_counts)?; - continue; - } - - // Fast path for Conv: use pre-computed params to skip attr FxHashMap lookups - if matches!( - kind, - NodeKind::Conv | NodeKind::ConvRelu | NodeKind::ConvSilu - ) { - let cp = model - .runtime_index - .conv_params - .get(i) - .and_then(|o| o.as_ref()); - let activation = match kind { - NodeKind::ConvRelu => yscv_kernels::Activation::Relu, - NodeKind::ConvSilu => yscv_kernels::Activation::Silu, - _ => yscv_kernels::Activation::None, - }; - - // --- DW+PW fusion: detect depthwise Conv followed by pointwise 1x1 --- - if cp.is_some_and(|p| p.is_depthwise) { - // Look ahead for pointwise 1x1 Conv consuming our output exclusively - let dw_out = &node.outputs[0]; - let dw_uses = use_counts.get(dw_out).copied().unwrap_or(0); - if dw_uses == 1 - && let Some(next_idx) = (i + 1..nodes.len()).find(|&j| !skip[j]) - { - let next_cp = model - .runtime_index - .conv_params - .get(next_idx) - .and_then(|o| o.as_ref()); - let next = &nodes[next_idx]; - let next_kind = node_kind(node_kinds, nodes, next_idx); - if next_cp.is_some_and(|p| p.is_pointwise && !p.has_padding) - && next.inputs.first().map(|s| s.as_str()) == Some(dw_out.as_str()) - && matches!( - next_kind, - NodeKind::Conv | NodeKind::ConvRelu | NodeKind::ConvSilu - ) - { - let pw_activation = match next_kind { - NodeKind::ConvRelu => yscv_kernels::Activation::Relu, - NodeKind::ConvSilu => yscv_kernels::Activation::Silu, - _ => yscv_kernels::Activation::None, - }; - let dw_input_ids_slice: &[Option] = model - .runtime_index - .node_input_ids - .get(i) - .map(|v| v.as_slice()) - .unwrap_or(&[]); - exec_fused_dw_pw( - node, - next, - &mut env, - activation, - pw_activation, - cp, - next_cp, - dw_input_ids_slice, - &mut remaining_uses, - &output_id_mask, - // Legacy single-pass dispatch path (used by - // calibration/quantize flows). M3 enclave - // lookahead is wired only in the plan-based - // `execute_plan_branch`; here we always - // convert back to NHWC at the boundary. - false, - )?; - // Decrement PW inputs so DW output also gets - // freed. Mirrors the post-action cleanup below, - // but tailored to the fused pair since the - // outer `continue` below would otherwise skip - // it. - let pw_pre_ids = model - .runtime_index - .node_input_ids - .get(next_idx) - .map(|v| v.as_slice()) - .unwrap_or(&[]); - for (inp_idx, inp) in next.inputs.iter().enumerate() { - if inp.is_empty() { - continue; - } - let id = pw_pre_ids - .get(inp_idx) - .and_then(|opt| *opt) - .or_else(|| env.resolve_id(inp)); - if let Some(id) = id - && id < remaining_uses.len() - { - remaining_uses[id] = remaining_uses[id].saturating_sub(1); - if remaining_uses[id] == 0 && !output_id_mask[id] { - env.remove_by_id(id); - } - } - } - skip[next_idx] = true; - continue; - } - } - } - - let prepacked = prepacked_for_conv_node(model, i); - exec_conv_with_params(node, &mut env, activation, cp, prepacked)?; - env.mark_nhwc(&node.outputs[0]); - } else { - execute_node_with_layout_kind(node, &mut env, kind)?; - } - - // --- Early deallocation: free tensors whose last consumer was this node --- - let input_ids = &model.runtime_index.node_input_ids; - let pre_ids = if i < input_ids.len() { - &input_ids[i] - } else { - &[][..] - }; - for (inp_idx, inp) in node.inputs.iter().enumerate() { - if inp.is_empty() { - continue; - } - // Use pre-resolved ID (O(1)) when available, fallback to FxHashMap. - let id = pre_ids - .get(inp_idx) - .and_then(|opt| *opt) - .or_else(|| env.resolve_id(inp)); - if let Some(id) = id - && id < remaining_uses.len() - { - remaining_uses[id] = remaining_uses[id].saturating_sub(1); - if remaining_uses[id] == 0 && !output_id_mask[id] { - env.remove_by_id(id); - } - } - } - } - - // Optional per-op trace for debugging inference divergence. - if std::env::var("CPU_TRACE").is_ok() { - for node in nodes { - for out_name in &node.outputs { - if let Some(t) = env.get(out_name) { - let d = t.data(); - if d.is_empty() { - continue; - } - let max = d.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)); - let min = d.iter().fold(f32::INFINITY, |a, &b| a.min(b)); - let mean = d.iter().sum::() / d.len() as f32; - let nhwc = if env.is_nhwc(out_name) { " [NHWC]" } else { "" }; - eprintln!( - "[{:>20}] {:60} shape={:?} min={:>10.4} max={:>10.4} mean={:>10.4}{}", - node.op_type, - out_name, - t.shape(), - min, - max, - mean, - nhwc, - ); - } - } - } - } - - // Ensure all outputs are in NCHW (ONNX standard layout) - for name in &model.outputs { - env.materialize_quant_i8_raw(name)?; - ensure_nchw(&mut env, name)?; - } - - let mut result = FxHashMap::with_capacity_and_hasher(model.outputs.len(), FxBuildHasher); - for name in &model.outputs { - if let Some(t) = env.remove(name) { - result.insert(name.clone(), t); - } else if let Some(t) = env.get(name) { - result.insert(name.clone(), t.clone()); - } else { - eprintln!("warning: ONNX output '{}' not found in environment", name); - } - } - Ok(result) -} diff --git a/crates/yscv-onnx/src/runner/layout.rs b/crates/yscv-onnx/src/runner/layout.rs index 9a1705c..941f594 100644 --- a/crates/yscv-onnx/src/runner/layout.rs +++ b/crates/yscv-onnx/src/runner/layout.rs @@ -49,113 +49,16 @@ pub(crate) fn reshape_nhwc_passthrough_disabled() -> bool { *CACHED.get_or_init(|| std::env::var_os("YSCV_RESHAPE_NHWC_PASSTHROUGH_OFF").is_some()) } -/// Fast path: avoid the `ensure_nchw` permute before a -/// `Reshape` when the input is NHWC physical rank-4 `[N,H,W,C]` and the -/// reshape merges spatial dims into `[N, C, H*W]` (the model's NCHW -/// logical view). The NHWC memory order is already `[N, H*W, C]` which -/// is what a downstream `Transpose(perm=[0,2,1])+MatMul` consumes as -/// its post-transpose A — `exec_fused_transpose_matmul` honours the -/// NHWC tag and switches to a non-transposed matmul kernel. +/// Fast path: avoid the `ensure_nchw` permute before a `Reshape` when the +/// input is NHWC physical rank-4 `[N,H,W,C]` and the reshape merges spatial +/// dims into `[N, C, H*W]` (the model's NCHW logical view). The NHWC memory +/// order is already `[N, H*W, C]`, which is what a downstream +/// `Transpose(perm=[0,2,1])+MatMul` consumes as its post-transpose A — +/// `exec_fused_transpose_matmul` honours the NHWC tag and switches to a +/// non-transposed matmul kernel. /// -/// Returns `Ok(true)` when the fast path handled the reshape (caller -/// should skip the default ensure_nchw+reshape path); `Ok(false)` -/// otherwise. -pub(crate) fn try_reshape_nhwc_passthrough( - node: &OnnxNode, - env: &mut TensorEnv, - use_counts: &FxHashMap, -) -> Result { - if node.inputs.len() < 2 || node.inputs[0].is_empty() { - return Ok(false); - } - if !env.is_nhwc(&node.inputs[0]) { - return Ok(false); - } - let output_name = match node.outputs.first() { - Some(n) if !n.is_empty() => n, - _ => return Ok(false), - }; - if !env.reshape_nhwc_passthrough_safe.contains(output_name) { - return Ok(false); - } - let in_shape = match env.get(&node.inputs[0]) { - Some(t) if t.rank() == 4 => t.shape().to_vec(), - _ => return Ok(false), - }; - let n = in_shape[0]; - let h = in_shape[1]; - let w = in_shape[2]; - let c = in_shape[3]; - let total: usize = in_shape.iter().product(); - let target_raw: Vec = match env.get(&node.inputs[1]) { - Some(t) => t.data().iter().map(|&v| v as i64).collect(), - None => return Ok(false), - }; - let mut target: Vec = Vec::with_capacity(target_raw.len()); - let mut neg_idx: Option = None; - for (i, &d) in target_raw.iter().enumerate() { - if d == -1 { - neg_idx = Some(i); - target.push(1); - } else if d == 0 { - target.push(if i < in_shape.len() { in_shape[i] } else { 1 }); - } else { - target.push(d as usize); - } - } - if let Some(idx) = neg_idx { - let known: usize = target - .iter() - .enumerate() - .filter(|&(i, _)| i != idx) - .map(|(_, &d)| d) - .product(); - target[idx] = total.checked_div(known.max(1)).unwrap_or(total); - } - if target.len() != 3 { - return Ok(false); - } - if target[0] != n || target[1] != c || target[2] != h * w { - return Ok(false); - } - // Metadata-only reshape: keep the NHWC physical data, set the - // model's NCHW logical shape, and keep the NHWC tag so the - // FusedTransposeMatMul consumer can adjust. - let sole_consumer = use_counts - .get(node.inputs[0].as_str()) - .copied() - .unwrap_or(0) - <= 1; - let new_shape = vec![n, c, h * w]; - let out = if sole_consumer { - let input = env - .remove(&node.inputs[0]) - .ok_or_else(|| OnnxError::MissingInput { - node: node.name.clone(), - input: node.inputs[0].clone(), - })?; - input - .into_reshape(new_shape) - .map_err(|e| OnnxError::DecodeFailed { - message: e.to_string(), - })? - } else { - let input = get_tensor(env, &node.name, &node.inputs[0])?; - input - .reshape(new_shape) - .map_err(|e| OnnxError::DecodeFailed { - message: e.to_string(), - })? - }; - env.insert(node.outputs[0].clone(), out); - env.mark_nhwc(&node.outputs[0]); - Ok(true) -} - -/// Same as [`try_reshape_nhwc_passthrough`] but without a borrowed -/// `use_counts` table — falls back to `reshape` (CoW clone) rather -/// than the `remove`+`into_reshape` zero-copy path. Used by the -/// plan-based `execute_node_with_layout_kind_inner` dispatch. +/// Returns `Ok(true)` when the fast path handled the reshape (caller should +/// skip the default ensure_nchw+reshape path); `Ok(false)` otherwise. fn try_reshape_nhwc_passthrough_inner( node: &OnnxNode, env: &mut TensorEnv, diff --git a/crates/yscv-onnx/src/runner/mod.rs b/crates/yscv-onnx/src/runner/mod.rs index f69f7ad..b344ff5 100644 --- a/crates/yscv-onnx/src/runner/mod.rs +++ b/crates/yscv-onnx/src/runner/mod.rs @@ -446,45 +446,6 @@ fn tensor_use_count( .unwrap_or_else(|| fallback_use_counts.get(name).copied().unwrap_or(0)) } -#[inline] -fn find_relu_after_identity_chain( - nodes: &[OnnxNode], - node_kinds: &[NodeKind], - start_idx: usize, - expected_input: &str, -) -> Option<(usize, Vec)> { - let mut idx = start_idx; - let mut current_input = expected_input.to_string(); - let mut identity_idxs = Vec::new(); - while let Some(node) = nodes.get(idx) { - let kind = node_kind(node_kinds, nodes, idx); - if node.op_type == "Identity" - && node.inputs.len() == 1 - && !node.outputs.is_empty() - && node.inputs[0] == current_input - { - identity_idxs.push(idx); - current_input = node.outputs[0].clone(); - idx += 1; - continue; - } - if kind == NodeKind::Relu && node.inputs.len() == 1 && node.inputs[0] == current_input { - return Some((idx, identity_idxs)); - } - break; - } - None -} - -#[inline] -fn mark_skip_indices(skip: &mut [bool], indices: &[usize]) { - for &idx in indices { - if idx < skip.len() { - skip[idx] = true; - } - } -} - fn run_onnx_model_inner( model: &OnnxModel, env: TensorEnv<'_, '_>, @@ -497,12 +458,10 @@ fn run_onnx_model_inner_specialized( env: TensorEnv<'_, '_>, specialization: Option<&ShapeSpecialization>, ) -> Result, OnnxError> { - // Use JIT execution plan if available (pre-compiled dispatch, no per-node matching) - if !model.runtime_index.execution_plan.is_empty() { - return run_onnx_model_jit(model, env, specialization); - } - - run_onnx_model_sequential(model, env) + // The execution plan is built at load time and has one entry per node, so + // it covers every graph; an empty plan just means an empty graph, which the + // JIT path handles by walking nothing and collecting the outputs. + run_onnx_model_jit(model, env, specialization) } #[inline] diff --git a/crates/yscv-onnx/src/runner/reshape.rs b/crates/yscv-onnx/src/runner/reshape.rs index 0750ed1..b2a6197 100644 --- a/crates/yscv-onnx/src/runner/reshape.rs +++ b/crates/yscv-onnx/src/runner/reshape.rs @@ -17,10 +17,6 @@ pub(super) fn exec_flatten(node: &OnnxNode, env: &mut TensorEnv) -> Result<(), O Ok(()) } -pub(super) fn exec_reshape(node: &OnnxNode, env: &mut TensorEnv) -> Result<(), OnnxError> { - exec_reshape_inner(node, env, None) -} - /// Reshape using the runner's cached shape-inference result. This is used only /// for ordinary NCHW inputs; layout-aware reshapes retain the generic path. pub(super) fn exec_reshape_known( @@ -42,22 +38,7 @@ pub(super) fn exec_reshape_known( Ok(()) } -/// Reshape with optional use-count awareness: when the data input has exactly -/// one remaining consumer we can `remove` it from the environment and call -/// `into_reshape` (zero-copy). Otherwise we fall back to a cloning `reshape`. -pub(super) fn exec_reshape_zerocopy( - node: &OnnxNode, - env: &mut TensorEnv, - use_counts: &FxHashMap, -) -> Result<(), OnnxError> { - exec_reshape_inner(node, env, Some(use_counts)) -} - -fn exec_reshape_inner( - node: &OnnxNode, - env: &mut TensorEnv, - use_counts: Option<&FxHashMap>, -) -> Result<(), OnnxError> { +pub(super) fn exec_reshape(node: &OnnxNode, env: &mut TensorEnv) -> Result<(), OnnxError> { // Compute new_shape without cloning the shape tensor data. We iterate the // shape-tensor slice directly inside the borrow scope, then drop the // borrow before touching env mutably. @@ -94,31 +75,16 @@ fn exec_reshape_inner( }; let _ = total; - // With Tensor::reshape now O(1) (Arc-shared storage with copy-on-write on - // subsequent writes), both paths are cheap. We still prefer the explicit - // `remove` + `into_reshape` when the node is the sole remaining consumer, - // because that drops the env slot early so downstream writes never even - // consider a CoW clone. - let sole_consumer = use_counts - .map(|uc| uc.get(node.inputs[0].as_str()).copied().unwrap_or(0) <= 1) - .unwrap_or(false); - - if sole_consumer && let Some(input) = env.remove(&node.inputs[0]) { - let out = input - .into_reshape(new_shape) - .map_err(|e| OnnxError::DecodeFailed { - message: e.to_string(), - })?; - env.insert(node.outputs[0].clone(), out); - } else { - let input = get_tensor(env, &node.name, &node.inputs[0])?; - let out = input - .reshape(new_shape) - .map_err(|e| OnnxError::DecodeFailed { - message: e.to_string(), - })?; - env.insert(node.outputs[0].clone(), out); - } + // `Tensor::reshape` is O(1) — Arc-shared storage with copy-on-write on + // subsequent writes — so there is nothing to gain from taking the tensor + // out of the env first. + let input = get_tensor(env, &node.name, &node.inputs[0])?; + let out = input + .reshape(new_shape) + .map_err(|e| OnnxError::DecodeFailed { + message: e.to_string(), + })?; + env.insert(node.outputs[0].clone(), out); Ok(()) } diff --git a/crates/yscv-onnx/src/runner/tensor_env.rs b/crates/yscv-onnx/src/runner/tensor_env.rs index 32105c1..9fe695d 100644 --- a/crates/yscv-onnx/src/runner/tensor_env.rs +++ b/crates/yscv-onnx/src/runner/tensor_env.rs @@ -399,6 +399,10 @@ impl<'m, 'i> TensorEnv<'m, 'i> { /// Get a mutable reference to a tensor by name. /// Clone-on-write: if the tensor is only in initializers, clone it into /// the slot first. + /// + /// Only the GPU backend mutates tensors in place through the env; the CPU + /// plan dispatch builds outputs and `insert`s them. + #[cfg(feature = "gpu")] #[inline] pub(crate) fn get_mut(&mut self, name: &str) -> Option<&mut Tensor> { let id = self.resolve_id(name)?; diff --git a/crates/yscv-onnx/src/tests/optimizer.rs b/crates/yscv-onnx/src/tests/optimizer.rs index 4383a93..4429e25 100644 --- a/crates/yscv-onnx/src/tests/optimizer.rs +++ b/crates/yscv-onnx/src/tests/optimizer.rs @@ -39,6 +39,114 @@ fn optimize_removes_dropout_nodes() { assert_eq!(model.nodes[1].inputs[0], "relu_out"); } +/// `NodeProto.name` is optional. The pass used to delete Dropouts by name via +/// `retain`, so a single unnamed Dropout put the empty string in the delete set +/// and took every other unnamed node in the graph with it. +#[test] +fn remove_dropout_keeps_unnamed_siblings() { + let nodes = vec![ + onnx::NodeProto { + op_type: Some("Relu".into()), + input: vec!["x".into()], + output: vec!["relu_out".into()], + ..Default::default() + }, + onnx::NodeProto { + op_type: Some("Dropout".into()), + input: vec!["relu_out".into()], + output: vec!["drop_out".into()], + ..Default::default() + }, + onnx::NodeProto { + op_type: Some("Relu".into()), + input: vec!["drop_out".into()], + output: vec!["y".into()], + ..Default::default() + }, + ]; + let bytes = build_minimal_onnx_model(nodes, vec![], vec!["x"], vec!["y"]); + let mut model = load_onnx_model(&bytes).unwrap(); + assert!( + model.nodes.iter().all(|n| n.name.is_empty()), + "fixture must have unnamed nodes for this to be a regression test" + ); + + optimize_onnx_graph(&mut model); + + assert_eq!( + model.node_count(), + 2, + "only the Dropout should go, not every unnamed node" + ); + assert!(model.nodes.iter().all(|n| n.op_type == "Relu")); + assert_eq!(model.nodes[1].inputs[0], "relu_out"); +} + +/// A Squeeze -> Unsqueeze -> Squeeze chain with matching axes matches the +/// inverse-pair predicate at both `i` and `i + 1`. Accepting both made the +/// reverse-removal loop delete already-shifted indices. +#[test] +fn eliminate_squeeze_unsqueeze_handles_overlapping_chain() { + let axes_attr = || make_ints_attr("axes", vec![0]); + let nodes = vec![ + onnx::NodeProto { + op_type: Some("Squeeze".into()), + name: Some("sq0".into()), + input: vec!["x".into()], + output: vec!["a".into()], + attribute: vec![axes_attr()], + ..Default::default() + }, + onnx::NodeProto { + op_type: Some("Unsqueeze".into()), + name: Some("unsq0".into()), + input: vec!["a".into()], + output: vec!["b".into()], + attribute: vec![axes_attr()], + ..Default::default() + }, + onnx::NodeProto { + op_type: Some("Squeeze".into()), + name: Some("sq1".into()), + input: vec!["b".into()], + output: vec!["c".into()], + attribute: vec![axes_attr()], + ..Default::default() + }, + onnx::NodeProto { + op_type: Some("Relu".into()), + name: Some("sink".into()), + input: vec!["c".into()], + output: vec!["y".into()], + ..Default::default() + }, + ]; + let bytes = build_minimal_onnx_model(nodes, vec![], vec!["x"], vec!["y"]); + let mut model = load_onnx_model(&bytes).unwrap(); + + optimize_onnx_graph(&mut model); + + // The Relu is the only node that must survive; whichever inverse pair the + // pass claims, it must not corrupt the graph around it. + let sink = model + .nodes + .iter() + .find(|n| n.op_type == "Relu") + .expect("Relu sink must survive"); + assert_eq!(model.outputs, vec!["y".to_string()]); + assert!( + !sink.inputs[0].is_empty(), + "sink input must stay wired to a real value" + ); + for node in &model.nodes { + assert!( + matches!(node.op_type.as_str(), "Squeeze" | "Unsqueeze" | "Relu"), + "unrelated node type {} appeared", + node.op_type + ); + } +} + #[test] fn optimize_eliminates_dead_nodes() { let nodes = vec![