Carry the weight-layout tag through an Identity fold

crates/yscv-onnx/src/ir/lower.rs

@@ -12,7 +12,7 @@
 //! means fixing the loader and the exporter together, which is not this
 //! change's job.
 
-use rustc_hash::FxHashMap;
+use rustc_hash::{FxHashMap, FxHashSet};
 
 use super::{Graph, Node, Op, ValueKind, WeightLayout};
 use crate::loader::{OnnxModel, OnnxNode};
@@ -129,27 +129,46 @@ impl Graph {
 impl OnnxModel {
     /// Writes an optimized graph back over this model.
     ///
-    /// Rebuilds `nodes` and `initializers` from the IR. Deliberately leaves the
-    /// loader's weight-layout side tables (`khwc_weights` and friends) and the
-    /// runtime index alone: layout tags are keyed by weight name and passes
-    /// preserve those names, and the index is rebuilt once by the driver after
-    /// the whole pipeline.
+    /// Rebuilds `nodes`, `initializers` and the loader's weight-layout side
+    /// tables from the IR. The runtime index is left alone — the driver rebuilds
+    /// it once after the whole pipeline.
+    ///
+    /// The side tables cannot be carried over by name. A pass that rewrites a
+    /// weight is free to give the result a new one, and the runtime reads the
+    /// tables by the name the Conv actually points at, so a renamed weight
+    /// silently reverts to being read as ONNX-native OIHW — a pre-permuted
+    /// depthwise `[KH, KW, C, 1]` then parses as `[O, I, KH, KW]`, and the
+    /// output-channel count comes back as the kernel height. `weight_layouts`
+    /// is keyed by value, survives the rename, and is the authority here.
     pub(crate) fn apply_ir(&mut self, graph: &Graph) {
         let nodes: Vec<OnnxNode> = graph
             .node_ids()
             .filter_map(|id| graph.to_onnx_node(id))
             .collect();
 
         let mut initializers = FxHashMap::default();
+        let mut khwc_weights = FxHashSet::default();
+        let mut dw_khwc_weights = FxHashSet::default();
+        let mut group_khwc_weights = FxHashSet::default();
         for idx in 0..graph.value_count() {
-            let value = graph.value(super::ValueId(idx as u32));
+            let id = super::ValueId(idx as u32);
+            let value = graph.value(id);
             if let ValueKind::Constant(tensor) = &value.kind {
                 initializers.insert(value.name.clone(), tensor.clone());
+                match graph.weight_layout(id) {
+                    WeightLayout::Khwc => khwc_weights.insert(value.name.clone()),
+                    WeightLayout::DepthwiseKhwc => dw_khwc_weights.insert(value.name.clone()),
+                    WeightLayout::GroupKhwc => group_khwc_weights.insert(value.name.clone()),
+                    WeightLayout::Oihw => false,
+                };
             }
         }
 
         self.nodes = nodes;
         self.initializers = initializers;
+        self.khwc_weights = khwc_weights;
+        self.dw_khwc_weights = dw_khwc_weights;
+        self.group_khwc_weights = group_khwc_weights;
         self.inputs = graph
             .graph_inputs()
             .iter()

crates/yscv-onnx/src/optimizer/fold_constants.rs

@@ -84,7 +84,18 @@ impl Pass for FoldConstants {
                 continue;
             };
             let output = node.outputs[0];
+            // An `Identity` folded away leaves its operand's bytes under a new
+            // name. The layout tag has to travel with them, or the consumer
+            // reads a permuted weight as ONNX-native OIHW.
+            let carried = (node.op == Op::Identity)
+                .then(|| node.inputs.first().copied().flatten())
+                .flatten()
+                .map(|src| graph.weight_layout(src))
+                .filter(|l| *l != WeightLayout::Oihw);
             graph.set_constant(output, folded);
+            if let Some(layout) = carried {
+                graph.set_weight_layout(output, layout);
+            }
             graph.remove_node(node_id);
             changed = true;
         }
@@ -112,8 +123,12 @@ fn foldable_inputs(graph: &Graph, node_id: NodeId) -> Option<Vec<(String, ValueI
         graph.constant(*value)?;
         // A pre-permuted weight's bytes no longer match the graph's logical
         // view of it, so evaluating an operator against one would compute
-        // against the wrong layout.
-        if graph.weight_layout(*value) != WeightLayout::Oihw {
+        // against the wrong layout. `Identity` is the exception: it copies
+        // bytes without interpreting them, and refusing it is not the safe
+        // choice — the node then survives to run time, where it hands its
+        // consumer a permuted tensor carrying no layout tag, and the Conv
+        // reading it parses `[KH, KW, C, 1]` as `[O, I, KH, KW]`.
+        if graph.weight_layout(*value) != WeightLayout::Oihw && node.op != Op::Identity {
             return None;
         }
         operands.push((graph.value(*value).name.clone(), *value));

crates/yscv-onnx/src/tests/equivalence.rs

@@ -136,6 +136,9 @@ pub(in crate::tests) fn assert_transform_preserves_numerics(
 /// Takes an [`EnvGuard`] rather than locking internally: the caller has to hold
 /// the lock across its *own* plan assertions too, and a second acquisition here
 /// would deadlock.
+// Only the pre-permuted (CPU-only) weight layout produces the fused plans
+// this checks; GPU builds keep ONNX-native OIHW and take a different path.
+#[cfg(not(any(feature = "metal-backend", feature = "gpu")))]
 pub(in crate::tests) fn assert_plan_fusion_preserves_numerics(
     _env: &EnvGuard,
     label: &str,

crates/yscv-onnx/src/tests/optimizer.rs

@@ -580,6 +580,89 @@ fn const_tensor(name: &str, dims: Vec<i64>, data: Vec<f32>) -> onnx::TensorProto
     }
 }
 
+/// An `Identity` aliasing a pre-permuted depthwise weight folds away, and the
+/// layout tag follows the bytes to the new name.
+///
+/// Leaving it unfolded is not the conservative choice: the node survives to run
+/// time and hands the Conv a `[KH, KW, C, 1]` tensor the runtime has no tag for,
+/// so the Conv parses it as `[O, I, KH, KW]` and reports the kernel height as
+/// its output-channel count.
+#[cfg(not(any(feature = "metal-backend", feature = "gpu")))]
+#[test]
+fn fold_constants_carries_weight_layout_through_identity() {
+    // The weight is read twice: once straight, which is what makes the loader
+    // permute and register it, and once through an `Identity`. Only the second
+    // reader loses the tag, which is why a model has to contain both to fail.
+    let c = 16i64;
+    let dw_attrs = || {
+        vec![
+            make_ints_attr("kernel_shape", vec![3, 3]),
+            make_ints_attr("pads", vec![1, 1, 1, 1]),
+            make_int_attr("group", c),
+        ]
+    };
+    let nodes = vec![
+        onnx::NodeProto {
+            op_type: Some("Conv".into()),
+            name: Some("dw_direct".into()),
+            input: vec!["x".into(), "dw_w".into()],
+            output: vec!["h".into()],
+            attribute: dw_attrs(),
+            ..Default::default()
+        },
+        onnx::NodeProto {
+            op_type: Some("Identity".into()),
+            name: Some("alias".into()),
+            input: vec!["dw_w".into()],
+            output: vec!["dw_w_alias".into()],
+            ..Default::default()
+        },
+        onnx::NodeProto {
+            op_type: Some("Conv".into()),
+            name: Some("dw_aliased".into()),
+            input: vec!["h".into(), "dw_w_alias".into()],
+            output: vec!["y".into()],
+            attribute: dw_attrs(),
+            ..Default::default()
+        },
+    ];
+    let bytes = build_minimal_onnx_model(
+        nodes,
+        vec![onnx::TensorProto {
+            name: Some("dw_w".into()),
+            dims: vec![c, 1, 3, 3],
+            data_type: Some(1),
+            float_data: vec![0.05; 16 * 9],
+            ..Default::default()
+        }],
+        vec!["x"],
+        vec!["y"],
+    );
+    let mut model = load_onnx_model(&bytes).unwrap();
+    assert!(
+        model.dw_khwc_weights.contains("dw_w"),
+        "fixture must start from a pre-permuted depthwise weight"
+    );
+
+    optimize_onnx_graph(&mut model).unwrap();
+
+    for node in model.nodes.iter().filter(|n| n.op_type.starts_with("Conv")) {
+        let w = &node.inputs[1];
+        assert!(
+            model.dw_khwc_weights.contains(w),
+            "`{}` reads `{w}`, which lost its depthwise-KHWC tag across the \
+             fold; registered: {:?}",
+            node.name,
+            model.dw_khwc_weights
+        );
+        assert_eq!(
+            model.initializers[w].shape(),
+            &[3, 3, 16, 1],
+            "and the bytes `{w}` points at are the permuted ones"
+        );
+    }
+}
+
 /// A chain of constant nodes folds in a single sweep: each fold makes its
 /// output constant, which can only enable nodes later in topological order.
 #[test]
@@ -1503,6 +1586,11 @@ fn plan_fuses_conv_add_with_the_later_producer_of_a_residual() {
 ///
 /// `channels` picks whether the blocked kernel is eligible at all — it needs
 /// both channel counts to be a multiple of 16.
+// Depends on the loader pre-permuting conv weights: the fusion and
+// handoff gates below match on the permuted shapes. GPU builds keep the
+// ONNX-native OIHW layout, so the plan legitimately comes out different
+// there and the shape this pins is not the one to expect.
+#[cfg(not(any(feature = "metal-backend", feature = "gpu")))]
 fn nchwc_handoff_for_chained_blocks(channels: usize) -> Vec<bool> {
     let c = channels as i64;
     let dw_w = |name: &str| onnx::TensorProto {
@@ -1575,6 +1663,7 @@ fn nchwc_handoff_for_chained_blocks(channels: usize) -> Vec<bool> {
 /// Checked here rather than through a run because the kernel that consumes the
 /// flag is gated on AVX-512, which the development host does not have — the
 /// predicate itself is what changed, so the predicate is what gets pinned.
+#[cfg(not(any(feature = "metal-backend", feature = "gpu")))]
 #[test]
 fn nchwc_handoff_is_resolved_at_plan_time() {
     // c = 16: both blocks clear the blocked kernel's channel gate, so the
@@ -1600,6 +1689,7 @@ fn nchwc_handoff_is_resolved_at_plan_time() {
 ///
 /// The DW output already has to have exactly one reader for the merge to be
 /// legal, so that reader is unique and is now looked up directly.
+#[cfg(not(any(feature = "metal-backend", feature = "gpu")))]
 #[test]
 fn plan_merges_pw_dw_pw_reduce_across_an_unrelated_node() {
     // Held for the whole test: the plan assertion below reads the same switch
@@ -1797,6 +1887,7 @@ fn pw_dw_pw_reduce_residual_fixture(
 /// with the graph output missing — no error anywhere.
 ///
 /// The merge now takes over whatever `ConvAdd` resolved, or declines.
+#[cfg(not(any(feature = "metal-backend", feature = "gpu")))]
 #[test]
 fn plan_merge_keeps_a_non_adjacent_residual_add() {
     let env = crate::tests::equivalence::lock_env();