From b89c2493806ee04412c7c1899fa31458e73d026d Mon Sep 17 00:00:00 2001 From: human9000 Date: Tue, 7 Jul 2026 22:28:25 +0300 Subject: [PATCH 1/3] Shape inference --- crates/yscv-onnx/src/lib.rs | 4 + crates/yscv-onnx/src/optimizer/mod.rs | 10 +- crates/yscv-onnx/src/shape_infer.rs | 794 ++++++++++++++++++++++ crates/yscv-onnx/src/tests/mod.rs | 1 + crates/yscv-onnx/src/tests/shape_infer.rs | 149 ++++ 5 files changed, 957 insertions(+), 1 deletion(-) create mode 100644 crates/yscv-onnx/src/shape_infer.rs create mode 100644 crates/yscv-onnx/src/tests/shape_infer.rs diff --git a/crates/yscv-onnx/src/lib.rs b/crates/yscv-onnx/src/lib.rs index c8b8d157..0b561c53 100644 --- a/crates/yscv-onnx/src/lib.rs +++ b/crates/yscv-onnx/src/lib.rs @@ -63,6 +63,10 @@ pub use runner::run_onnx_model; pub use runner::run_onnx_model_borrowed; pub use runner::run_onnx_model_borrowed_slice; pub use runner::{QuantRuntimeStats, quant_runtime_stats, reset_quant_runtime_stats}; +pub use shape_infer::{ + Dim, ShapeDiagnostic, ShapeInference, ShapeMap, TensorShape, infer_shapes, + infer_shapes_from_tensors, +}; #[cfg(all(target_os = "macos", feature = "metal-backend"))] pub use runner::metal_runner::{ diff --git a/crates/yscv-onnx/src/optimizer/mod.rs b/crates/yscv-onnx/src/optimizer/mod.rs index e08aae5b..22a51d45 100644 --- a/crates/yscv-onnx/src/optimizer/mod.rs +++ b/crates/yscv-onnx/src/optimizer/mod.rs @@ -12,7 +12,15 @@ mod remove_dropout_nodes; mod reorder_nodes_for_fusion; mod strip_qdq_within_fusion_chains; -use crate::{loader::OnnxModel, optimizer::{eliminate_dead_code::eliminate_dead_code, eliminate_squeeze_unsqueeze_pairs::eliminate_squeeze_unsqueeze_pairs, remove_dropout_nodes::remove_dropout_nodes, reorder_nodes_for_fusion::reorder_nodes_for_fusion}}; +use crate::{ + loader::OnnxModel, + optimizer::{ + eliminate_dead_code::eliminate_dead_code, + eliminate_squeeze_unsqueeze_pairs::eliminate_squeeze_unsqueeze_pairs, + remove_dropout_nodes::remove_dropout_nodes, + reorder_nodes_for_fusion::reorder_nodes_for_fusion, + }, +}; pub use analyze_nchwc::analyze_nchwc; pub use fold_constants::fold_constants; diff --git a/crates/yscv-onnx/src/shape_infer.rs b/crates/yscv-onnx/src/shape_infer.rs new file mode 100644 index 00000000..74024525 --- /dev/null +++ b/crates/yscv-onnx/src/shape_infer.rs @@ -0,0 +1,794 @@ +use std::collections::HashMap; + +use thiserror::Error; +use yscv_tensor::Tensor; + +use crate::loader::{OnnxAttribute, OnnxModel, OnnxNode}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Dim { + Known(usize), + Unknown, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct TensorShape { + pub dims: Vec, +} + +impl TensorShape { + pub fn known(dims: Vec) -> Self { + Self { + dims: dims.into_iter().map(Dim::Known).collect(), + } + } + + pub fn unknown_rank(rank: usize) -> Self { + Self { + dims: vec![Dim::Unknown; rank], + } + } + + pub fn rank(&self) -> usize { + self.dims.len() + } + + pub fn dim(&self, idx: usize) -> Option { + match self.dims.get(idx) { + Some(Dim::Known(v)) => Some(*v), + _ => None, + } + } + + pub fn num_elements(&self) -> Option { + let mut n = 1u64; + for dim in &self.dims { + match dim { + Dim::Known(v) => n = n.checked_mul(*v as u64)?, + Dim::Unknown => return None, + } + } + Some(n) + } + + pub fn as_known_dims(&self) -> Option> { + self.dims + .iter() + .map(|d| match d { + Dim::Known(v) => Some(*v), + Dim::Unknown => None, + }) + .collect() + } +} + +pub type ShapeMap = HashMap; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ShapeDiagnostic { + pub node_index: usize, + pub node_name: String, + pub op_type: String, + pub error: ShapeError, +} + +#[derive(Debug, Clone, PartialEq, Eq, Error)] +pub enum ShapeError { + #[error("unsupported shape rule for op {op_type}")] + UnsupportedOp { op_type: String }, + #[error("{op_type} missing required input {index}")] + MissingInput { op_type: String, index: usize }, + #[error("missing input shape for value {name}")] + MissingInputShape { name: String }, + #[error("missing initializer {name}")] + MissingInitializer { name: String }, + #[error("{op_type} requires rank {expected}, got rank {actual}")] + RankMismatch { + op_type: String, + expected: usize, + actual: usize, + }, + #[error("{op_type} requires rank at least {min}, got rank {actual}")] + RankTooSmall { + op_type: String, + min: usize, + actual: usize, + }, + #[error("{op_type} supports {expected}")] + UnsupportedRank { + op_type: String, + expected: &'static str, + }, + #[error("incompatible broadcast dimensions {left} and {right}")] + BroadcastIncompatible { left: usize, right: usize }, + #[error("axis {axis} out of range for rank {rank}")] + AxisOutOfRange { axis: i64, rank: usize }, + #[error("{op_type} missing required attribute {name}")] + MissingAttribute { op_type: String, name: &'static str }, + #[error("Constant without tensor value attribute")] + ConstantWithoutTensor, + #[error("{op_type} input {name} must be constant")] + NonConstantInput { op_type: String, name: String }, + #[error("Reshape target contains unsupported dimension {dim}")] + InvalidReshapeDim { dim: i64 }, + #[error("{op_type} input ranks differ")] + InputRanksDiffer { op_type: String }, + #[error("Transpose perm rank {perm_rank} does not match input rank {input_rank}")] + TransposePermRankMismatch { perm_rank: usize, input_rank: usize }, +} + +#[derive(Debug, Clone)] +pub struct ShapeInference { + pub shapes: ShapeMap, + pub diagnostics: Vec, +} + +pub fn infer_shapes(model: &OnnxModel, input_shapes: &ShapeMap) -> ShapeInference { + let mut shapes = input_shapes.clone(); + let mut diagnostics = Vec::new(); + + for (name, tensor) in &model.initializers { + shapes.insert(name.clone(), TensorShape::known(tensor.shape().to_vec())); + } + + for (idx, node) in model.nodes.iter().enumerate() { + match infer_node(model, &shapes, node) { + Ok(outputs) => { + for (name, shape) in node.outputs.iter().zip(outputs) { + if !name.is_empty() { + shapes.insert(name.clone(), shape); + } + } + } + Err(error) => diagnostics.push(ShapeDiagnostic { + node_index: idx, + node_name: node.name.clone(), + op_type: node.op_type.clone(), + error, + }), + } + } + + ShapeInference { + shapes, + diagnostics, + } +} + +pub fn infer_shapes_from_tensors( + model: &OnnxModel, + inputs: &HashMap, +) -> ShapeInference { + let input_shapes: ShapeMap = inputs + .iter() + .map(|(name, tensor)| (name.clone(), TensorShape::known(tensor.shape().to_vec()))) + .collect(); + infer_shapes(model, &input_shapes) +} + +fn infer_node( + model: &OnnxModel, + shapes: &ShapeMap, + node: &OnnxNode, +) -> Result, ShapeError> { + match node.op_type.as_str() { + "Constant" => infer_constant(node), + "Identity" | "Dropout" => unary_same_shape(shapes, node), + + "Relu" | "Sigmoid" | "Tanh" | "Clip" | "BatchNormalization" | "BatchNormalization_Relu" => { + unary_same_shape(shapes, node) + } + "Conv" | "Conv_Relu" | "Conv_SiLU" => infer_conv(model, shapes, node), + "Add" | "Sub" | "Mul" | "Div" | "Pow" => infer_broadcast(shapes, node), + "Concat" => infer_concat(shapes, node), + "Transpose" => infer_transpose(shapes, node), + "Reshape" => infer_reshape(model, shapes, node), + "Flatten" => infer_flatten(shapes, node), + "Squeeze" => infer_squeeze(model, shapes, node), + "Unsqueeze" => infer_unsqueeze(model, shapes, node), + "MatMul" => infer_matmul(shapes, node), + "Gemm" => infer_gemm(model, shapes, node), + "MaxPool" | "AveragePool" => infer_pool(shapes, node), + "GlobalAveragePool" => infer_global_pool(shapes, node), + _ => Err(ShapeError::UnsupportedOp { + op_type: node.op_type.clone(), + }), + } +} + +fn unary_same_shape(shapes: &ShapeMap, node: &OnnxNode) -> Result, ShapeError> { + let input = node + .inputs + .first() + .ok_or_else(|| ShapeError::MissingInput { + op_type: node.op_type.clone(), + index: 0, + })?; + let shape = shapes + .get(input) + .cloned() + .ok_or_else(|| ShapeError::MissingInputShape { + name: input.clone(), + })?; + Ok(vec![shape; node.outputs.len().max(1)]) +} + +fn infer_constant(node: &OnnxNode) -> Result, ShapeError> { + match node.attributes.get("value") { + Some(OnnxAttribute::Tensor(t)) => Ok(vec![TensorShape::known(t.shape().to_vec())]), + _ => Err(ShapeError::ConstantWithoutTensor), + } +} + +fn infer_conv( + model: &OnnxModel, + shapes: &ShapeMap, + node: &OnnxNode, +) -> Result, ShapeError> { + let input_name = node + .inputs + .first() + .ok_or_else(|| ShapeError::MissingInput { + op_type: node.op_type.clone(), + index: 0, + })?; + let weight_name = node.inputs.get(1).ok_or_else(|| ShapeError::MissingInput { + op_type: node.op_type.clone(), + index: 1, + })?; + let input = shapes + .get(input_name) + .ok_or_else(|| ShapeError::MissingInputShape { + name: input_name.clone(), + })?; + let weight = + model + .initializers + .get(weight_name) + .ok_or_else(|| ShapeError::MissingInitializer { + name: weight_name.clone(), + })?; + let w_shape = weight.shape(); + if input.rank() != 4 || w_shape.len() != 4 { + return Err(ShapeError::UnsupportedRank { + op_type: node.op_type.clone(), + expected: "rank-4 NCHW inputs and rank-4 weights", + }); + } + + let n = input.dims[0]; + let ih = input.dims[2]; + let iw = input.dims[3]; + let oc = if model.khwc_weights.contains(weight_name) { + w_shape[3] + } else if model.dw_khwc_weights.contains(weight_name) { + w_shape[2] + } else { + w_shape[0] + }; + let (kh, kw) = if model.khwc_weights.contains(weight_name) + || model.dw_khwc_weights.contains(weight_name) + { + (w_shape[0], w_shape[1]) + } else { + (w_shape[2], w_shape[3]) + }; + + let strides = ints_attr(node, "strides").unwrap_or_else(|| vec![1, 1]); + let dilations = ints_attr(node, "dilations").unwrap_or_else(|| vec![1, 1]); + let pads = ints_attr(node, "pads").unwrap_or_else(|| vec![0, 0, 0, 0]); + let sh = strides.first().copied().unwrap_or(1).max(1) as usize; + let sw = strides.get(1).copied().unwrap_or(1).max(1) as usize; + let dh = dilations.first().copied().unwrap_or(1).max(1) as usize; + let dw = dilations.get(1).copied().unwrap_or(1).max(1) as usize; + let pt = pads.first().copied().unwrap_or(0).max(0) as usize; + let pl = pads.get(1).copied().unwrap_or(0).max(0) as usize; + let pb = pads.get(2).copied().unwrap_or(0).max(0) as usize; + let pr = pads.get(3).copied().unwrap_or(0).max(0) as usize; + + let oh = conv_output_dim(ih, kh, sh, dh, pt, pb); + let ow = conv_output_dim(iw, kw, sw, dw, pl, pr); + Ok(vec![TensorShape { + dims: vec![n, Dim::Known(oc), oh, ow], + }]) +} + +fn conv_output_dim( + input: Dim, + kernel: usize, + stride: usize, + dilation: usize, + pad_begin: usize, + pad_end: usize, +) -> Dim { + match input { + Dim::Known(v) => { + let effective = dilation + .saturating_mul(kernel.saturating_sub(1)) + .saturating_add(1); + let padded = v.saturating_add(pad_begin).saturating_add(pad_end); + if padded < effective { + Dim::Known(0) + } else { + Dim::Known((padded - effective) / stride + 1) + } + } + Dim::Unknown => Dim::Unknown, + } +} + +fn infer_broadcast(shapes: &ShapeMap, node: &OnnxNode) -> Result, ShapeError> { + if node.inputs.len() < 2 { + return Err(ShapeError::MissingInput { + op_type: node.op_type.clone(), + index: 1, + }); + } + let a = shapes + .get(&node.inputs[0]) + .ok_or_else(|| ShapeError::MissingInputShape { + name: node.inputs[0].clone(), + })?; + let b = shapes + .get(&node.inputs[1]) + .ok_or_else(|| ShapeError::MissingInputShape { + name: node.inputs[1].clone(), + })?; + Ok(vec![broadcast_shapes(a, b)?]) +} + +fn broadcast_shapes(a: &TensorShape, b: &TensorShape) -> Result { + let rank = a.rank().max(b.rank()); + let mut dims = Vec::with_capacity(rank); + for i in 0..rank { + let ad = dim_from_right(a, i); + let bd = dim_from_right(b, i); + let out = match (ad, bd) { + (Dim::Known(1), d) | (d, Dim::Known(1)) => d, + (Dim::Known(x), Dim::Known(y)) if x == y => Dim::Known(x), + (Dim::Unknown, _) | (_, Dim::Unknown) => Dim::Unknown, + (Dim::Known(x), Dim::Known(y)) => { + return Err(ShapeError::BroadcastIncompatible { left: x, right: y }); + } + }; + dims.push(out); + } + dims.reverse(); + Ok(TensorShape { dims }) +} + +fn dim_from_right(shape: &TensorShape, idx_from_right: usize) -> Dim { + if idx_from_right >= shape.rank() { + Dim::Known(1) + } else { + shape.dims[shape.rank() - 1 - idx_from_right] + } +} + +fn infer_concat(shapes: &ShapeMap, node: &OnnxNode) -> Result, ShapeError> { + let first_name = node + .inputs + .first() + .ok_or_else(|| ShapeError::MissingInput { + op_type: node.op_type.clone(), + index: 0, + })?; + let first = shapes + .get(first_name) + .cloned() + .ok_or_else(|| ShapeError::MissingInputShape { + name: first_name.clone(), + })?; + let rank = first.rank(); + let axis = normalize_axis(int_attr(node, "axis").unwrap_or(0), rank)?; + let mut out = first; + let mut axis_sum = 0usize; + let mut axis_known = true; + for name in &node.inputs { + let shape = shapes + .get(name) + .ok_or_else(|| ShapeError::MissingInputShape { name: name.clone() })?; + if shape.rank() != rank { + return Err(ShapeError::InputRanksDiffer { + op_type: node.op_type.clone(), + }); + } + for d in 0..rank { + if d == axis { + if let Some(v) = shape.dim(d) { + axis_sum = axis_sum.saturating_add(v); + } else { + axis_known = false; + } + } else if out.dims[d] != shape.dims[d] { + out.dims[d] = Dim::Unknown; + } + } + } + out.dims[axis] = if axis_known { + Dim::Known(axis_sum) + } else { + Dim::Unknown + }; + Ok(vec![out]) +} + +fn infer_transpose(shapes: &ShapeMap, node: &OnnxNode) -> Result, ShapeError> { + let input_name = node + .inputs + .first() + .ok_or_else(|| ShapeError::MissingInput { + op_type: node.op_type.clone(), + index: 0, + })?; + let input = shapes + .get(input_name) + .ok_or_else(|| ShapeError::MissingInputShape { + name: input_name.clone(), + })?; + let perm = ints_attr(node, "perm").unwrap_or_else(|| (0..input.rank() as i64).rev().collect()); + if perm.len() != input.rank() { + return Err(ShapeError::TransposePermRankMismatch { + perm_rank: perm.len(), + input_rank: input.rank(), + }); + } + let mut dims = Vec::with_capacity(input.rank()); + for p in perm { + let idx = normalize_axis(p, input.rank())?; + dims.push(input.dims[idx]); + } + Ok(vec![TensorShape { dims }]) +} + +fn infer_reshape( + model: &OnnxModel, + shapes: &ShapeMap, + node: &OnnxNode, +) -> Result, ShapeError> { + let input_name = node + .inputs + .first() + .ok_or_else(|| ShapeError::MissingInput { + op_type: node.op_type.clone(), + index: 0, + })?; + let input = shapes + .get(input_name) + .ok_or_else(|| ShapeError::MissingInputShape { + name: input_name.clone(), + })?; + let shape_name = node.inputs.get(1).ok_or_else(|| ShapeError::MissingInput { + op_type: node.op_type.clone(), + index: 1, + })?; + let target = + model + .initializers + .get(shape_name) + .ok_or_else(|| ShapeError::NonConstantInput { + op_type: node.op_type.clone(), + name: shape_name.clone(), + })?; + let raw: Vec = target.data().iter().map(|v| *v as i64).collect(); + let input_numel = input.num_elements(); + let mut dims = Vec::with_capacity(raw.len()); + let mut known_product = 1u64; + let mut minus_one = None; + for (idx, &d) in raw.iter().enumerate() { + if d == 0 { + let copied = input.dims.get(idx).copied().unwrap_or(Dim::Unknown); + if let Dim::Known(v) = copied { + known_product = known_product.saturating_mul(v as u64); + } + dims.push(copied); + } else if d == -1 { + minus_one = Some(idx); + dims.push(Dim::Unknown); + } else if d > 0 { + known_product = known_product.saturating_mul(d as u64); + dims.push(Dim::Known(d as usize)); + } else { + return Err(ShapeError::InvalidReshapeDim { dim: d }); + } + } + if let (Some(total), Some(idx)) = (input_numel, minus_one) + && known_product != 0 + && total % known_product == 0 + { + dims[idx] = Dim::Known((total / known_product) as usize); + } + Ok(vec![TensorShape { dims }]) +} + +fn infer_flatten(shapes: &ShapeMap, node: &OnnxNode) -> Result, ShapeError> { + let input_name = node + .inputs + .first() + .ok_or_else(|| ShapeError::MissingInput { + op_type: node.op_type.clone(), + index: 0, + })?; + let input = shapes + .get(input_name) + .ok_or_else(|| ShapeError::MissingInputShape { + name: input_name.clone(), + })?; + let axis = normalize_axis(int_attr(node, "axis").unwrap_or(1), input.rank())?; + let left = product_dims(&input.dims[..axis]); + let right = product_dims(&input.dims[axis..]); + Ok(vec![TensorShape { + dims: vec![left, right], + }]) +} + +fn infer_squeeze( + model: &OnnxModel, + shapes: &ShapeMap, + node: &OnnxNode, +) -> Result, ShapeError> { + let input_name = node + .inputs + .first() + .ok_or_else(|| ShapeError::MissingInput { + op_type: node.op_type.clone(), + index: 0, + })?; + let input = shapes + .get(input_name) + .ok_or_else(|| ShapeError::MissingInputShape { + name: input_name.clone(), + })?; + let axes = node_axes(model, node)?; + let dims = if axes.is_empty() { + input + .dims + .iter() + .copied() + .filter(|d| *d != Dim::Known(1)) + .collect() + } else { + let axes: Vec = axes + .into_iter() + .map(|a| normalize_axis(a, input.rank())) + .collect::>()?; + input + .dims + .iter() + .enumerate() + .filter_map(|(idx, d)| if axes.contains(&idx) { None } else { Some(*d) }) + .collect() + }; + Ok(vec![TensorShape { dims }]) +} + +fn infer_unsqueeze( + model: &OnnxModel, + shapes: &ShapeMap, + node: &OnnxNode, +) -> Result, ShapeError> { + let input_name = node + .inputs + .first() + .ok_or_else(|| ShapeError::MissingInput { + op_type: node.op_type.clone(), + index: 0, + })?; + let input = shapes + .get(input_name) + .ok_or_else(|| ShapeError::MissingInputShape { + name: input_name.clone(), + })?; + let out_rank = input.rank().saturating_add(node_axes(model, node)?.len()); + let mut axes: Vec = node_axes(model, node)? + .into_iter() + .map(|a| normalize_axis(a, out_rank)) + .collect::>()?; + axes.sort_unstable(); + let mut input_idx = 0usize; + let mut dims = Vec::with_capacity(out_rank); + for out_idx in 0..out_rank { + if axes.contains(&out_idx) { + dims.push(Dim::Known(1)); + } else { + dims.push(input.dims[input_idx]); + input_idx += 1; + } + } + Ok(vec![TensorShape { dims }]) +} + +fn infer_matmul(shapes: &ShapeMap, node: &OnnxNode) -> Result, ShapeError> { + if node.inputs.len() < 2 { + return Err(ShapeError::MissingInput { + op_type: node.op_type.clone(), + index: 1, + }); + } + let a = shapes + .get(&node.inputs[0]) + .ok_or_else(|| ShapeError::MissingInputShape { + name: node.inputs[0].clone(), + })?; + let b = shapes + .get(&node.inputs[1]) + .ok_or_else(|| ShapeError::MissingInputShape { + name: node.inputs[1].clone(), + })?; + if a.rank() < 2 || b.rank() < 2 { + return Err(ShapeError::RankTooSmall { + op_type: node.op_type.clone(), + min: 2, + actual: a.rank().min(b.rank()), + }); + } + let a_batch = TensorShape { + dims: a.dims[..a.rank() - 2].to_vec(), + }; + let b_batch = TensorShape { + dims: b.dims[..b.rank() - 2].to_vec(), + }; + let batch = broadcast_shapes(&a_batch, &b_batch)?.dims; + let mut dims = batch; + dims.push(a.dims[a.rank() - 2]); + dims.push(b.dims[b.rank() - 1]); + Ok(vec![TensorShape { dims }]) +} + +fn infer_gemm( + model: &OnnxModel, + shapes: &ShapeMap, + node: &OnnxNode, +) -> Result, ShapeError> { + if node.inputs.len() < 2 { + return Err(ShapeError::MissingInput { + op_type: node.op_type.clone(), + index: 1, + }); + } + let a = shapes + .get(&node.inputs[0]) + .ok_or_else(|| ShapeError::MissingInputShape { + name: node.inputs[0].clone(), + })?; + let b = shapes + .get(&node.inputs[1]) + .ok_or_else(|| ShapeError::MissingInputShape { + name: node.inputs[1].clone(), + })?; + if a.rank() != 2 || b.rank() != 2 { + return Err(ShapeError::UnsupportedRank { + op_type: node.op_type.clone(), + expected: "rank-2 inputs", + }); + } + let trans_a = int_attr(node, "transA").unwrap_or(0) != 0; + let trans_b = int_attr(node, "transB").unwrap_or(0) != 0; + let m = if trans_a { a.dims[1] } else { a.dims[0] }; + let n = if trans_b { b.dims[0] } else { b.dims[1] }; + let out = TensorShape { dims: vec![m, n] }; + if let Some(c_name) = node.inputs.get(2) + && !c_name.is_empty() + && let Some(c) = model.initializers.get(c_name) + { + let c_shape = TensorShape::known(c.shape().to_vec()); + let _ = broadcast_shapes(&out, &c_shape)?; + } + Ok(vec![out]) +} + +fn infer_pool(shapes: &ShapeMap, node: &OnnxNode) -> Result, ShapeError> { + let input_name = node + .inputs + .first() + .ok_or_else(|| ShapeError::MissingInput { + op_type: node.op_type.clone(), + index: 0, + })?; + let input = shapes + .get(input_name) + .ok_or_else(|| ShapeError::MissingInputShape { + name: input_name.clone(), + })?; + if input.rank() != 4 { + return Err(ShapeError::UnsupportedRank { + op_type: node.op_type.clone(), + expected: "rank-4 NCHW inputs", + }); + } + let kernels = ints_attr(node, "kernel_shape").ok_or(ShapeError::MissingAttribute { + op_type: node.op_type.clone(), + name: "kernel_shape", + })?; + let strides = ints_attr(node, "strides").unwrap_or_else(|| vec![1, 1]); + let pads = ints_attr(node, "pads").unwrap_or_else(|| vec![0, 0, 0, 0]); + let kh = kernels.first().copied().unwrap_or(1).max(1) as usize; + let kw = kernels.get(1).copied().unwrap_or(1).max(1) as usize; + let sh = strides.first().copied().unwrap_or(1).max(1) as usize; + let sw = strides.get(1).copied().unwrap_or(1).max(1) as usize; + let pt = pads.first().copied().unwrap_or(0).max(0) as usize; + let pl = pads.get(1).copied().unwrap_or(0).max(0) as usize; + let pb = pads.get(2).copied().unwrap_or(0).max(0) as usize; + let pr = pads.get(3).copied().unwrap_or(0).max(0) as usize; + Ok(vec![TensorShape { + dims: vec![ + input.dims[0], + input.dims[1], + conv_output_dim(input.dims[2], kh, sh, 1, pt, pb), + conv_output_dim(input.dims[3], kw, sw, 1, pl, pr), + ], + }]) +} + +fn infer_global_pool(shapes: &ShapeMap, node: &OnnxNode) -> Result, ShapeError> { + let input_name = node + .inputs + .first() + .ok_or_else(|| ShapeError::MissingInput { + op_type: node.op_type.clone(), + index: 0, + })?; + let input = shapes + .get(input_name) + .ok_or_else(|| ShapeError::MissingInputShape { + name: input_name.clone(), + })?; + if input.rank() != 4 { + return Err(ShapeError::UnsupportedRank { + op_type: node.op_type.clone(), + expected: "rank-4 NCHW inputs", + }); + } + Ok(vec![TensorShape { + dims: vec![input.dims[0], input.dims[1], Dim::Known(1), Dim::Known(1)], + }]) +} + +fn product_dims(dims: &[Dim]) -> Dim { + let mut product = 1usize; + for dim in dims { + match dim { + Dim::Known(v) => product = product.saturating_mul(*v), + Dim::Unknown => return Dim::Unknown, + } + } + Dim::Known(product) +} + +fn node_axes(model: &OnnxModel, node: &OnnxNode) -> Result, ShapeError> { + if node.inputs.len() > 1 && !node.inputs[1].is_empty() { + let t = model.initializers.get(&node.inputs[1]).ok_or_else(|| { + ShapeError::NonConstantInput { + op_type: node.op_type.clone(), + name: node.inputs[1].clone(), + } + })?; + return Ok(t.data().iter().map(|v| *v as i64).collect()); + } + if let Some(OnnxAttribute::Ints(axes)) = node.attributes.get("axes") { + return Ok(axes.clone()); + } + Ok(Vec::new()) +} + +fn normalize_axis(axis: i64, rank: usize) -> Result { + let normalized = if axis < 0 { rank as i64 + axis } else { axis }; + if normalized < 0 || normalized >= rank as i64 { + return Err(ShapeError::AxisOutOfRange { axis, rank }); + } + Ok(normalized as usize) +} + +fn int_attr(node: &OnnxNode, name: &str) -> Option { + match node.attributes.get(name) { + Some(OnnxAttribute::Int(v)) => Some(*v), + _ => None, + } +} + +fn ints_attr(node: &OnnxNode, name: &str) -> Option> { + match node.attributes.get(name) { + Some(OnnxAttribute::Ints(v)) => Some(v.clone()), + _ => None, + } +} diff --git a/crates/yscv-onnx/src/tests/mod.rs b/crates/yscv-onnx/src/tests/mod.rs index 6565d00c..f7e3b8ca 100644 --- a/crates/yscv-onnx/src/tests/mod.rs +++ b/crates/yscv-onnx/src/tests/mod.rs @@ -11,6 +11,7 @@ mod normalization; mod optimizer; mod qlinear; mod reshape; +mod shape_infer; use prost::Message; diff --git a/crates/yscv-onnx/src/tests/shape_infer.rs b/crates/yscv-onnx/src/tests/shape_infer.rs new file mode 100644 index 00000000..1544697d --- /dev/null +++ b/crates/yscv-onnx/src/tests/shape_infer.rs @@ -0,0 +1,149 @@ +use std::collections::HashMap; + +use yscv_tensor::Tensor; + +use super::{build_minimal_onnx_model, make_ints_attr}; +use crate::proto::onnx; +use crate::{Dim, TensorShape, graph_cost, infer_shapes, load_onnx_model}; + +fn tensor_proto(name: &str, shape: Vec, data: Vec) -> onnx::TensorProto { + onnx::TensorProto { + name: Some(name.to_string()), + dims: shape, + data_type: Some(1), + float_data: data, + ..Default::default() + } +} + +#[test] +fn infers_conv_shape_and_cost() { + let node = onnx::NodeProto { + op_type: Some("Conv".into()), + name: Some("conv".into()), + input: vec!["x".into(), "w".into()], + output: vec!["y".into()], + attribute: vec![ + make_ints_attr("strides", vec![2, 2]), + make_ints_attr("pads", vec![1, 1, 1, 1]), + ], + ..Default::default() + }; + let model = load_onnx_model(&build_minimal_onnx_model( + vec![node], + vec![tensor_proto( + "w", + vec![8, 3, 3, 3], + vec![0.0; 8 * 3 * 3 * 3], + )], + vec!["x"], + vec!["y"], + )) + .unwrap(); + let input_shapes = HashMap::from([("x".to_string(), TensorShape::known(vec![1, 3, 32, 32]))]); + + let inferred = infer_shapes(&model, &input_shapes); + assert!( + inferred.diagnostics.is_empty(), + "{:?}", + inferred.diagnostics + ); + assert_eq!( + inferred.shapes["y"].dims, + vec![Dim::Known(1), Dim::Known(8), Dim::Known(16), Dim::Known(16)] + ); + + let cost = graph_cost(&model, &inferred); + assert_eq!(cost.unknown_nodes, 0); + assert_eq!(cost.estimated_macs, 8 * 16 * 16 * 3 * 3 * 3); +} + +#[test] +fn infers_reshape_and_matmul_shapes() { + let reshape = onnx::NodeProto { + op_type: Some("Reshape".into()), + name: Some("reshape".into()), + input: vec!["x".into(), "shape".into()], + output: vec!["flat".into()], + ..Default::default() + }; + let matmul = onnx::NodeProto { + op_type: Some("MatMul".into()), + name: Some("matmul".into()), + input: vec!["flat".into(), "w".into()], + output: vec!["y".into()], + ..Default::default() + }; + let model = load_onnx_model(&build_minimal_onnx_model( + vec![reshape, matmul], + vec![ + tensor_proto("shape", vec![2], vec![1.0, 48.0]), + tensor_proto("w", vec![48, 10], vec![0.0; 48 * 10]), + ], + vec!["x"], + vec!["y"], + )) + .unwrap(); + let input_shapes = HashMap::from([("x".to_string(), TensorShape::known(vec![1, 3, 4, 4]))]); + + let inferred = infer_shapes(&model, &input_shapes); + assert!( + inferred.diagnostics.is_empty(), + "{:?}", + inferred.diagnostics + ); + assert_eq!(inferred.shapes["flat"].as_known_dims(), Some(vec![1, 48])); + assert_eq!(inferred.shapes["y"].as_known_dims(), Some(vec![1, 10])); + + let cost = graph_cost(&model, &inferred); + assert_eq!(cost.unknown_nodes, 0); + assert_eq!(cost.estimated_macs, 480); +} + +#[test] +fn reports_unknowns_without_guessing() { + let node = onnx::NodeProto { + op_type: Some("Resize".into()), + name: Some("resize".into()), + input: vec!["x".into()], + output: vec!["y".into()], + ..Default::default() + }; + let model = load_onnx_model(&build_minimal_onnx_model( + vec![node], + vec![], + vec!["x"], + vec!["y"], + )) + .unwrap(); + let input_shapes = HashMap::from([("x".to_string(), TensorShape::known(vec![1, 3, 8, 8]))]); + + let inferred = infer_shapes(&model, &input_shapes); + assert_eq!(inferred.diagnostics.len(), 1); + + let cost = graph_cost(&model, &inferred); + assert_eq!(cost.unknown_nodes, 1); + assert_eq!(cost.known_nodes, 0); +} + +#[test] +fn infers_shapes_from_initializer_inputs() { + let node = onnx::NodeProto { + op_type: Some("Relu".into()), + name: Some("relu".into()), + input: vec!["x".into()], + output: vec!["y".into()], + ..Default::default() + }; + let init = Tensor::from_vec(vec![2], vec![1.0, -1.0]).unwrap(); + let model = load_onnx_model(&build_minimal_onnx_model( + vec![node], + vec![tensor_proto("x", vec![2], init.data().to_vec())], + vec!["x"], + vec!["y"], + )) + .unwrap(); + + let inferred = infer_shapes(&model, &HashMap::new()); + assert_eq!(inferred.shapes["y"].as_known_dims(), Some(vec![2])); +} From 96c2d4dbdff6683a3533b4ed40588269cfb44e25 Mon Sep 17 00:00:00 2001 From: human9000 Date: Tue, 7 Jul 2026 22:28:45 +0300 Subject: [PATCH 2/3] graph cost eval --- crates/yscv-onnx/src/lib.rs | 6 +- crates/yscv-onnx/src/optimizer/graph_cost.rs | 378 +++++++++++++++++++ crates/yscv-onnx/src/optimizer/mod.rs | 2 + 3 files changed, 384 insertions(+), 2 deletions(-) create mode 100644 crates/yscv-onnx/src/optimizer/graph_cost.rs diff --git a/crates/yscv-onnx/src/lib.rs b/crates/yscv-onnx/src/lib.rs index 0b561c53..2d03c83b 100644 --- a/crates/yscv-onnx/src/lib.rs +++ b/crates/yscv-onnx/src/lib.rs @@ -15,6 +15,7 @@ mod optimizer; mod proto; pub mod quantize; mod runner; +pub mod shape_infer; pub use dtype::{OnnxDtype, OnnxTensorData}; pub use error::OnnxError; @@ -27,8 +28,9 @@ pub use loader::{ OnnxAttribute, OnnxModel, OnnxNode, OnnxTensor, load_onnx_model, load_onnx_model_from_file, }; pub use optimizer::{ - GraphStats, fold_constants, fold_conv_bn, fuse_bn_relu, fuse_conv_relu, graph_stats, - optimize_onnx_graph, strip_qdq_within_fusion_chains, + GraphCost, GraphCostDiff, GraphStats, NodeCost, fold_constants, fold_conv_bn, fuse_bn_relu, + fuse_conv_relu, graph_cost, graph_cost_diff, graph_stats, optimize_onnx_graph, + strip_qdq_within_fusion_chains, }; pub use quantize::quantize_weights_int4; pub use quantize::{CalibrationCollector, CalibrationScope, MinMax}; diff --git a/crates/yscv-onnx/src/optimizer/graph_cost.rs b/crates/yscv-onnx/src/optimizer/graph_cost.rs new file mode 100644 index 00000000..e337a18d --- /dev/null +++ b/crates/yscv-onnx/src/optimizer/graph_cost.rs @@ -0,0 +1,378 @@ +use crate::loader::{OnnxAttribute, OnnxModel, OnnxNode}; +use crate::shape_infer::{ShapeInference, TensorShape}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct NodeCost { + pub index: usize, + pub name: String, + pub op_type: String, + pub output_shape: Option, + pub macs: u64, + pub element_ops: u64, + pub bytes_read: u64, + pub bytes_written: u64, + pub score: u64, + pub reason: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GraphCost { + pub node_count: usize, + pub known_nodes: usize, + pub unknown_nodes: usize, + pub estimated_macs: u64, + pub estimated_element_ops: u64, + pub estimated_bytes_read: u64, + pub estimated_bytes_written: u64, + pub score: u64, + pub nodes: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GraphCostDiff { + pub before_score: u64, + pub after_score: u64, + pub delta_score: i128, + pub before_nodes: usize, + pub after_nodes: usize, + pub delta_nodes: isize, + pub before_macs: u64, + pub after_macs: u64, + pub delta_macs: i128, +} + +pub fn graph_cost(model: &OnnxModel, shapes: &ShapeInference) -> GraphCost { + let mut nodes = Vec::with_capacity(model.nodes.len()); + let mut known_nodes = 0usize; + let mut unknown_nodes = 0usize; + let mut estimated_macs = 0u64; + let mut estimated_element_ops = 0u64; + let mut estimated_bytes_read = 0u64; + let mut estimated_bytes_written = 0u64; + let diagnostic_by_index: std::collections::HashMap = shapes + .diagnostics + .iter() + .map(|d| (d.node_index, d.error.to_string())) + .collect(); + + for (index, node) in model.nodes.iter().enumerate() { + let output_shape = node + .outputs + .first() + .and_then(|name| shapes.shapes.get(name)) + .cloned(); + let mut cost = cost_node(model, node, output_shape.as_ref()); + cost.index = index; + cost.name = node.name.clone(); + cost.op_type = node.op_type.clone(); + if cost.reason.is_none() { + cost.reason = diagnostic_by_index.get(&index).cloned(); + } + if cost.reason.is_some() { + unknown_nodes += 1; + } else { + known_nodes += 1; + } + estimated_macs = estimated_macs.saturating_add(cost.macs); + estimated_element_ops = estimated_element_ops.saturating_add(cost.element_ops); + estimated_bytes_read = estimated_bytes_read.saturating_add(cost.bytes_read); + estimated_bytes_written = estimated_bytes_written.saturating_add(cost.bytes_written); + nodes.push(cost); + } + + let score = estimated_macs + .saturating_add(estimated_element_ops) + .saturating_add(estimated_bytes_read.saturating_add(estimated_bytes_written) / 4); + + GraphCost { + node_count: model.nodes.len(), + known_nodes, + unknown_nodes, + estimated_macs, + estimated_element_ops, + estimated_bytes_read, + estimated_bytes_written, + score, + nodes, + } +} + +pub fn graph_cost_diff(before: &GraphCost, after: &GraphCost) -> GraphCostDiff { + GraphCostDiff { + before_score: before.score, + after_score: after.score, + delta_score: after.score as i128 - before.score as i128, + before_nodes: before.node_count, + after_nodes: after.node_count, + delta_nodes: after.node_count as isize - before.node_count as isize, + before_macs: before.estimated_macs, + after_macs: after.estimated_macs, + delta_macs: after.estimated_macs as i128 - before.estimated_macs as i128, + } +} + +fn cost_node(model: &OnnxModel, node: &OnnxNode, output_shape: Option<&TensorShape>) -> NodeCost { + let mut cost = NodeCost { + index: 0, + name: String::new(), + op_type: String::new(), + output_shape: output_shape.cloned(), + macs: 0, + element_ops: 0, + bytes_read: 0, + bytes_written: 0, + score: 0, + reason: None, + }; + + match node.op_type.as_str() { + "Conv" | "Conv_Relu" | "Conv_SiLU" => cost_conv(model, node, output_shape, &mut cost), + "MatMul" => cost_matmul(model, node, output_shape, &mut cost), + "Gemm" => cost_gemm(model, node, output_shape, &mut cost), + "Relu" + | "Clip" + | "Sigmoid" + | "Tanh" + | "Add" + | "Sub" + | "Mul" + | "Div" + | "Pow" + | "BatchNormalization" + | "BatchNormalization_Relu" => { + let weight = match node.op_type.as_str() { + "Sigmoid" | "Tanh" => 8, + "BatchNormalization" | "BatchNormalization_Relu" => 4, + _ => 1, + }; + cost_elementwise(output_shape, weight, &mut cost); + } + "Transpose" => cost_materialized_copy(output_shape, &mut cost), + "Reshape" | "Flatten" | "Squeeze" | "Unsqueeze" | "Identity" | "Dropout" | "Constant" => { + cost_metadata_only(output_shape, &mut cost); + } + "Concat" => cost_concat(output_shape, &mut cost), + "MaxPool" | "AveragePool" => cost_pool(node, output_shape, &mut cost), + "GlobalAveragePool" => cost_global_pool(model, node, output_shape, &mut cost), + _ => cost.reason = Some("unsupported cost rule".to_string()), + } + + cost.score = cost + .macs + .saturating_add(cost.element_ops) + .saturating_add(cost.bytes_read.saturating_add(cost.bytes_written) / 4); + cost +} + +fn cost_conv( + model: &OnnxModel, + node: &OnnxNode, + output_shape: Option<&TensorShape>, + cost: &mut NodeCost, +) { + let Some(out) = output_shape else { + cost.reason = Some("missing Conv output shape".to_string()); + return; + }; + let Some(weight_name) = node.inputs.get(1) else { + cost.reason = Some("Conv missing weight".to_string()); + return; + }; + let Some(weight) = model.initializers.get(weight_name) else { + cost.reason = Some("Conv weight is not constant".to_string()); + return; + }; + let Some(out_numel) = out.num_elements() else { + cost.reason = Some("Conv output has unknown dimensions".to_string()); + return; + }; + let w_shape = weight.shape(); + if w_shape.len() != 4 { + cost.reason = Some("Conv weight is not rank 4".to_string()); + return; + } + let group = int_attr(node, "group").unwrap_or(1).max(1) as usize; + let (ic_per_group, kh, kw) = if model.khwc_weights.contains(weight_name) { + (w_shape[2], w_shape[0], w_shape[1]) + } else if model.dw_khwc_weights.contains(weight_name) { + (1, w_shape[0], w_shape[1]) + } else { + (w_shape[1], w_shape[2], w_shape[3]) + }; + let kernel_work = if group == w_shape[0] && ic_per_group == 1 { + kh.saturating_mul(kw) + } else { + ic_per_group.saturating_mul(kh).saturating_mul(kw) + }; + cost.macs = out_numel.saturating_mul(kernel_work as u64); + if matches!(node.op_type.as_str(), "Conv_Relu" | "Conv_SiLU") { + cost.element_ops = out_numel; + } + cost.bytes_written = out_numel.saturating_mul(4); + cost.bytes_read = cost + .bytes_written + .saturating_add((weight.data().len() as u64).saturating_mul(4)); +} + +fn cost_matmul( + model: &OnnxModel, + node: &OnnxNode, + output_shape: Option<&TensorShape>, + cost: &mut NodeCost, +) { + let (Some(a), Some(b)) = (node.inputs.first(), node.inputs.get(1)) else { + cost.reason = Some("MatMul missing inputs".to_string()); + return; + }; + let Some(out) = output_shape else { + cost.reason = Some("missing MatMul output shape".to_string()); + return; + }; + let Some(out_numel) = out.num_elements() else { + cost.reason = Some("MatMul output has unknown dimensions".to_string()); + return; + }; + let k = model + .initializers + .get(a) + .and_then(|t| t.shape().last().copied()) + .or_else(|| { + model + .initializers + .get(b) + .and_then(|t| t.shape().first().copied()) + }); + let Some(k) = k else { + cost.reason = Some("MatMul reduction dimension unknown".to_string()); + return; + }; + cost.macs = out_numel.saturating_mul(k as u64); + cost.bytes_written = out_numel.saturating_mul(4); + cost.bytes_read = cost.bytes_written; + for input in [&a, &b] { + if let Some(t) = model.initializers.get(*input) { + cost.bytes_read = cost + .bytes_read + .saturating_add((t.data().len() as u64).saturating_mul(4)); + } + } +} + +fn cost_gemm( + model: &OnnxModel, + node: &OnnxNode, + output_shape: Option<&TensorShape>, + cost: &mut NodeCost, +) { + let Some(out) = output_shape else { + cost.reason = Some("missing Gemm output shape".to_string()); + return; + }; + let Some(out_numel) = out.num_elements() else { + cost.reason = Some("Gemm output has unknown dimensions".to_string()); + return; + }; + let Some(b_name) = node.inputs.get(1) else { + cost.reason = Some("Gemm missing B input".to_string()); + return; + }; + let Some(b) = model.initializers.get(b_name) else { + cost.reason = Some("Gemm B shape unknown".to_string()); + return; + }; + if b.shape().len() != 2 { + cost.reason = Some("Gemm B is not rank 2".to_string()); + return; + } + let trans_b = int_attr(node, "transB").unwrap_or(0) != 0; + let k = if trans_b { b.shape()[1] } else { b.shape()[0] }; + cost.macs = out_numel.saturating_mul(k as u64); + cost.bytes_written = out_numel.saturating_mul(4); + cost.bytes_read = cost + .bytes_written + .saturating_add((b.data().len() as u64).saturating_mul(4)); +} + +fn cost_elementwise(output_shape: Option<&TensorShape>, weight: u64, cost: &mut NodeCost) { + let Some(out_numel) = output_shape.and_then(TensorShape::num_elements) else { + cost.reason = Some("elementwise output shape unknown".to_string()); + return; + }; + cost.element_ops = out_numel.saturating_mul(weight); + cost.bytes_read = out_numel.saturating_mul(4); + cost.bytes_written = out_numel.saturating_mul(4); +} + +fn cost_materialized_copy(output_shape: Option<&TensorShape>, cost: &mut NodeCost) { + let Some(out_numel) = output_shape.and_then(TensorShape::num_elements) else { + cost.reason = Some("copy output shape unknown".to_string()); + return; + }; + cost.bytes_read = out_numel.saturating_mul(4); + cost.bytes_written = out_numel.saturating_mul(4); +} + +fn cost_metadata_only(output_shape: Option<&TensorShape>, cost: &mut NodeCost) { + if output_shape.is_none() { + cost.reason = Some("metadata op output shape unknown".to_string()); + } +} + +fn cost_concat(output_shape: Option<&TensorShape>, cost: &mut NodeCost) { + cost_materialized_copy(output_shape, cost); +} + +fn cost_pool(node: &OnnxNode, output_shape: Option<&TensorShape>, cost: &mut NodeCost) { + let Some(out_numel) = output_shape.and_then(TensorShape::num_elements) else { + cost.reason = Some("pool output shape unknown".to_string()); + return; + }; + let kernel = ints_attr(node, "kernel_shape") + .and_then(|v| v.first().zip(v.get(1)).map(|(h, w)| h.max(&1) * w.max(&1))) + .unwrap_or(1) as u64; + cost.element_ops = out_numel.saturating_mul(kernel); + cost.bytes_read = out_numel.saturating_mul(kernel).saturating_mul(4); + cost.bytes_written = out_numel.saturating_mul(4); +} + +fn cost_global_pool( + model: &OnnxModel, + node: &OnnxNode, + output_shape: Option<&TensorShape>, + cost: &mut NodeCost, +) { + let Some(out_numel) = output_shape.and_then(TensorShape::num_elements) else { + cost.reason = Some("global pool output shape unknown".to_string()); + return; + }; + let input_spatial = node + .inputs + .first() + .and_then(|name| model.initializers.get(name)) + .and_then(|t| { + t.shape() + .get(2) + .zip(t.shape().get(3)) + .map(|(h, w)| h.saturating_mul(*w)) + }) + .unwrap_or(1) as u64; + cost.element_ops = out_numel.saturating_mul(input_spatial); + cost.bytes_read = cost.element_ops.saturating_mul(4); + cost.bytes_written = out_numel.saturating_mul(4); +} + +fn int_attr(node: &OnnxNode, name: &str) -> Option { + if let Some(OnnxAttribute::Int(v)) = node.attributes.get(name) { + Some(*v) + } else { + None + } +} + +fn ints_attr(node: &OnnxNode, name: &str) -> Option> { + if let Some(OnnxAttribute::Ints(v)) = node.attributes.get(name) { + Some(v.clone()) + } else { + None + } +} diff --git a/crates/yscv-onnx/src/optimizer/mod.rs b/crates/yscv-onnx/src/optimizer/mod.rs index 22a51d45..e0ba07c0 100644 --- a/crates/yscv-onnx/src/optimizer/mod.rs +++ b/crates/yscv-onnx/src/optimizer/mod.rs @@ -7,6 +7,7 @@ mod fold_conv_bn; mod fold_conv_mul; mod fuse_bn_relu; mod fuse_conv_relu; +mod graph_cost; mod graph_stats; mod remove_dropout_nodes; mod reorder_nodes_for_fusion; @@ -29,6 +30,7 @@ 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::{GraphCost, GraphCostDiff, NodeCost, graph_cost, graph_cost_diff}; pub use graph_stats::{GraphStats, graph_stats}; pub use strip_qdq_within_fusion_chains::strip_qdq_within_fusion_chains; From 067b459d346d8178376bf58f25cab11b5066e363 Mon Sep 17 00:00:00 2001 From: human9000 Date: Fri, 10 Jul 2026 18:24:29 +0300 Subject: [PATCH 3/3] graph cost diff for opt efficiency eval --- .../perf-runners/onnx-pr-bench/src/main.rs | 24 +++- crates/yscv-onnx/src/lib.rs | 4 +- crates/yscv-onnx/src/optimizer/graph_cost.rs | 121 +++++++++++++++++- crates/yscv-onnx/src/optimizer/mod.rs | 4 +- crates/yscv-onnx/src/tests/optimizer.rs | 113 +++++++++++++++- 5 files changed, 258 insertions(+), 8 deletions(-) diff --git a/benchmarks/perf-runners/onnx-pr-bench/src/main.rs b/benchmarks/perf-runners/onnx-pr-bench/src/main.rs index f7bab5be..ed811a82 100644 --- a/benchmarks/perf-runners/onnx-pr-bench/src/main.rs +++ b/benchmarks/perf-runners/onnx-pr-bench/src/main.rs @@ -8,8 +8,9 @@ use std::time::Instant; use serde_json::json; use yscv_onnx::{ - OnnxExportAttr, OnnxExportGraph, OnnxExportNode, OnnxExportValueInfo, OnnxRunner, - export_onnx_model_to_file, load_onnx_model_from_file, optimize_onnx_graph, + export_onnx_model_to_file, graph_cost, graph_cost_report, infer_shapes_from_tensors, + load_onnx_model_from_file, optimize_onnx_graph, OnnxExportAttr, OnnxExportGraph, + OnnxExportNode, OnnxExportValueInfo, OnnxRunner, }; use yscv_tensor::Tensor; @@ -685,6 +686,10 @@ fn run_case( }) .collect::>()?; let feed: Vec<(&str, &Tensor)> = inputs.iter().map(|(n, t)| (n.as_str(), t)).collect(); + let input_map: std::collections::HashMap = inputs.iter().cloned().collect(); + let shape_inference = infer_shapes_from_tensors(&model, &input_map); + let graph_cost = graph_cost(&model, &shape_inference); + let graph_cost_text = graph_cost_report(&graph_cost); let mut run_summaries = Vec::with_capacity(runs); for run_idx in 0..runs { @@ -723,6 +728,16 @@ fn run_case( "median_p50_us": median_p50_us, "nodes_before_opt": nodes_before, "nodes_after_opt": nodes_after, + "graph_cost": { + "node_count": graph_cost.node_count, + "known_nodes": graph_cost.known_nodes, + "unknown_nodes": graph_cost.unknown_nodes, + "estimated_macs": graph_cost.estimated_macs, + "estimated_element_ops": graph_cost.estimated_element_ops, + "estimated_bytes_read": graph_cost.estimated_bytes_read, + "estimated_bytes_written": graph_cost.estimated_bytes_written, + "score": graph_cost.score, + }, "dispatch": yscv_kernels::runtime_dispatch_report().to_string(), "profile_summary": profile_summary, }); @@ -733,7 +748,10 @@ fn run_case( output, serde_json::to_vec_pretty(&report).map_err(|e| format!("encode json: {e}"))?, ) - .map_err(|e| format!("write {}: {e}", output.display())) + .map_err(|e| format!("write {}: {e}", output.display()))?; + let graph_cost_path = output.with_extension("graph-cost.txt"); + std::fs::write(&graph_cost_path, graph_cost_text) + .map_err(|e| format!("write {}: {e}", graph_cost_path.display())) } fn run() -> Result<(), String> { diff --git a/crates/yscv-onnx/src/lib.rs b/crates/yscv-onnx/src/lib.rs index 2d03c83b..adb0cd06 100644 --- a/crates/yscv-onnx/src/lib.rs +++ b/crates/yscv-onnx/src/lib.rs @@ -29,8 +29,8 @@ pub use loader::{ }; pub use optimizer::{ GraphCost, GraphCostDiff, GraphStats, NodeCost, fold_constants, fold_conv_bn, fuse_bn_relu, - fuse_conv_relu, graph_cost, graph_cost_diff, graph_stats, optimize_onnx_graph, - strip_qdq_within_fusion_chains, + fuse_conv_relu, graph_cost, graph_cost_diff, graph_cost_report, graph_stats, + optimize_onnx_graph, strip_qdq_within_fusion_chains, }; pub use quantize::quantize_weights_int4; pub use quantize::{CalibrationCollector, CalibrationScope, MinMax}; diff --git a/crates/yscv-onnx/src/optimizer/graph_cost.rs b/crates/yscv-onnx/src/optimizer/graph_cost.rs index e337a18d..5184d80a 100644 --- a/crates/yscv-onnx/src/optimizer/graph_cost.rs +++ b/crates/yscv-onnx/src/optimizer/graph_cost.rs @@ -1,11 +1,14 @@ use crate::loader::{OnnxAttribute, OnnxModel, OnnxNode}; -use crate::shape_infer::{ShapeInference, TensorShape}; +use crate::shape_infer::{Dim, ShapeInference, TensorShape}; +use std::fmt::Write as _; #[derive(Debug, Clone, PartialEq, Eq)] pub struct NodeCost { pub index: usize, pub name: String, pub op_type: String, + pub inputs: Vec, + pub outputs: Vec, pub output_shape: Option, pub macs: u64, pub element_ops: u64, @@ -65,6 +68,8 @@ pub fn graph_cost(model: &OnnxModel, shapes: &ShapeInference) -> GraphCost { cost.index = index; cost.name = node.name.clone(); cost.op_type = node.op_type.clone(); + cost.inputs = node.inputs.clone(); + cost.outputs = node.outputs.clone(); if cost.reason.is_none() { cost.reason = diagnostic_by_index.get(&index).cloned(); } @@ -111,11 +116,75 @@ pub fn graph_cost_diff(before: &GraphCost, after: &GraphCost) -> GraphCostDiff { } } +/// Formats a deterministic text snapshot of graph-cost estimates. +/// +/// Intended for benchmarking and review workflows where the output is saved +/// and diffed across commits to check whether optimizer changes make the graph +/// lighter. The per-node section is sorted by a stable textual key rather than +/// execution order so incidental reordering produces less diff noise. +pub fn graph_cost_report(cost: &GraphCost) -> String { + let mut out = String::new(); + let _ = writeln!(out, "graph_cost v1"); + let _ = writeln!(out, "summary.node_count={}", cost.node_count); + let _ = writeln!(out, "summary.known_nodes={}", cost.known_nodes); + let _ = writeln!(out, "summary.unknown_nodes={}", cost.unknown_nodes); + let _ = writeln!(out, "summary.estimated_macs={}", cost.estimated_macs); + let _ = writeln!( + out, + "summary.estimated_element_ops={}", + cost.estimated_element_ops + ); + let _ = writeln!( + out, + "summary.estimated_bytes_read={}", + cost.estimated_bytes_read + ); + let _ = writeln!( + out, + "summary.estimated_bytes_written={}", + cost.estimated_bytes_written + ); + let _ = writeln!(out, "summary.score={}", cost.score); + out.push('\n'); + out.push_str("nodes\n"); + + let mut nodes: Vec<&NodeCost> = cost.nodes.iter().collect(); + nodes.sort_by(|a, b| { + stable_node_key(a) + .cmp(&stable_node_key(b)) + .then(a.index.cmp(&b.index)) + }); + + for node in nodes { + let _ = writeln!(out, "- key={}", stable_node_key(node)); + let _ = writeln!(out, " index={}", node.index); + let _ = writeln!(out, " op_type={}", node.op_type); + let _ = writeln!(out, " name={}", printable_field(&node.name)); + let _ = writeln!(out, " inputs={}", join_fields(&node.inputs)); + let _ = writeln!(out, " outputs={}", join_fields(&node.outputs)); + let _ = writeln!( + out, + " output_shape={}", + format_output_shape(node.output_shape.as_ref()) + ); + let _ = writeln!(out, " macs={}", node.macs); + let _ = writeln!(out, " element_ops={}", node.element_ops); + let _ = writeln!(out, " bytes_read={}", node.bytes_read); + let _ = writeln!(out, " bytes_written={}", node.bytes_written); + let _ = writeln!(out, " score={}", node.score); + let _ = writeln!(out, " reason={}", format_reason(node.reason.as_deref())); + } + + out +} + fn cost_node(model: &OnnxModel, node: &OnnxNode, output_shape: Option<&TensorShape>) -> NodeCost { let mut cost = NodeCost { index: 0, name: String::new(), op_type: String::new(), + inputs: Vec::new(), + outputs: Vec::new(), output_shape: output_shape.cloned(), macs: 0, element_ops: 0, @@ -369,6 +438,56 @@ fn int_attr(node: &OnnxNode, name: &str) -> Option { } } +fn stable_node_key(node: &NodeCost) -> String { + format!( + "{}|{}|{}|{}", + node.op_type, + printable_field(&node.name), + join_fields(&node.outputs), + join_fields(&node.inputs) + ) +} + +fn printable_field(s: &str) -> &str { + if s.is_empty() { "-" } else { s } +} + +fn join_fields(fields: &[String]) -> String { + if fields.is_empty() { + "-".to_string() + } else { + fields + .iter() + .map(|s| printable_field(s)) + .collect::>() + .join(",") + } +} + +fn format_output_shape(shape: Option<&TensorShape>) -> String { + let Some(shape) = shape else { + return "?".to_string(); + }; + let mut out = String::from("["); + for (i, dim) in shape.dims.iter().enumerate() { + if i > 0 { + out.push(','); + } + match dim { + Dim::Known(v) => { + let _ = write!(out, "{v}"); + } + Dim::Unknown => out.push('?'), + } + } + out.push(']'); + out +} + +fn format_reason(reason: Option<&str>) -> &str { + reason.unwrap_or("-") +} + fn ints_attr(node: &OnnxNode, name: &str) -> Option> { if let Some(OnnxAttribute::Ints(v)) = node.attributes.get(name) { Some(v.clone()) diff --git a/crates/yscv-onnx/src/optimizer/mod.rs b/crates/yscv-onnx/src/optimizer/mod.rs index e0ba07c0..04dd0589 100644 --- a/crates/yscv-onnx/src/optimizer/mod.rs +++ b/crates/yscv-onnx/src/optimizer/mod.rs @@ -30,7 +30,9 @@ 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::{GraphCost, GraphCostDiff, NodeCost, graph_cost, graph_cost_diff}; +pub use graph_cost::{ + GraphCost, GraphCostDiff, NodeCost, graph_cost, graph_cost_diff, graph_cost_report, +}; pub use graph_stats::{GraphStats, graph_stats}; pub use strip_qdq_within_fusion_chains::strip_qdq_within_fusion_chains; diff --git a/crates/yscv-onnx/src/tests/optimizer.rs b/crates/yscv-onnx/src/tests/optimizer.rs index eaad3a2e..682f8323 100644 --- a/crates/yscv-onnx/src/tests/optimizer.rs +++ b/crates/yscv-onnx/src/tests/optimizer.rs @@ -1,5 +1,8 @@ use super::*; -use crate::optimizer::{fuse_conv_relu, graph_stats, optimize_onnx_graph}; +use crate::optimizer::{ + fuse_conv_relu, graph_cost, graph_cost_report, graph_stats, optimize_onnx_graph, +}; +use crate::{TensorShape, infer_shapes}; #[test] fn optimize_removes_dropout_nodes() { @@ -143,6 +146,114 @@ fn reorder_enables_fusion_on_interleaved_branches() { .collect::>() ); } +#[test] +fn graph_cost_report_is_deterministic_and_sorted_by_stable_key() { + let nodes = vec![ + onnx::NodeProto { + op_type: Some("Relu".into()), + name: Some("relu_b".into()), + input: vec!["mid_b".into()], + output: vec!["out_b".into()], + ..Default::default() + }, + onnx::NodeProto { + op_type: Some("Relu".into()), + name: Some("relu_a".into()), + input: vec!["mid_a".into()], + output: vec!["out_a".into()], + ..Default::default() + }, + ]; + let bytes = build_minimal_onnx_model( + nodes, + vec![], + vec!["mid_a", "mid_b"], + vec!["out_a", "out_b"], + ); + let model = load_onnx_model(&bytes).unwrap(); + let shapes = infer_shapes( + &model, + &std::collections::HashMap::from([ + ("mid_a".to_string(), TensorShape::known(vec![1, 4])), + ("mid_b".to_string(), TensorShape::known(vec![1, 4])), + ]), + ); + let cost = graph_cost(&model, &shapes); + let report_a = graph_cost_report(&cost); + let report_b = graph_cost_report(&cost); + + assert_eq!( + report_a, report_b, + "report formatting must be deterministic" + ); + let pos_a = report_a.find("key=Relu|relu_a|out_a|mid_a").unwrap(); + let pos_b = report_a.find("key=Relu|relu_b|out_b|mid_b").unwrap(); + assert!( + pos_a < pos_b, + "nodes should be sorted by stable textual key" + ); +} + +#[test] +fn graph_cost_report_shows_lighter_graph_after_optimization() { + let nodes = vec![ + onnx::NodeProto { + op_type: Some("Conv".into()), + name: Some("conv0".into()), + input: vec!["x".into(), "w".into()], + output: vec!["conv_out".into()], + ..Default::default() + }, + onnx::NodeProto { + op_type: Some("Relu".into()), + name: Some("relu0".into()), + input: vec!["conv_out".into()], + output: vec!["relu_out".into()], + ..Default::default() + }, + onnx::NodeProto { + op_type: Some("Dropout".into()), + name: Some("drop0".into()), + input: vec!["relu_out".into()], + output: vec!["y".into()], + ..Default::default() + }, + ]; + let w = onnx::TensorProto { + name: Some("w".into()), + dims: vec![8, 3, 3, 3], + data_type: Some(1), + float_data: vec![0.0; 8 * 3 * 3 * 3], + ..Default::default() + }; + let bytes = build_minimal_onnx_model(nodes, vec![w], vec!["x"], vec!["y"]); + let mut model = load_onnx_model(&bytes).unwrap(); + let input_shapes = + std::collections::HashMap::from([("x".to_string(), TensorShape::known(vec![1, 3, 8, 8]))]); + + let before_shapes = infer_shapes(&model, &input_shapes); + let before = graph_cost(&model, &before_shapes); + let before_report = graph_cost_report(&before); + + optimize_onnx_graph(&mut model); + + let after_shapes = infer_shapes(&model, &input_shapes); + let after = graph_cost(&model, &after_shapes); + let after_report = graph_cost_report(&after); + + assert!( + after.node_count < before.node_count, + "optimizer should remove/fuse nodes" + ); + assert!( + after.score <= before.score, + "optimized graph should not get heavier in this case" + ); + assert!(before_report.contains("summary.node_count=3")); + assert!(after_report.contains("summary.node_count=1")); + assert!(after_report.contains("op_type=Conv_Relu")); + assert!(!after_report.contains("op_type=Dropout")); +} #[test] fn graph_stats_reports_op_counts() {