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![