From ff85b06cdb46650629cf86a0fa0422b1c5683b79 Mon Sep 17 00:00:00 2001 From: human9000 Date: Tue, 14 Jul 2026 11:04:38 +0300 Subject: [PATCH] utilize load-time shape inference during runtime * Reshape uses inferred shape if present without re-parsing the node * metal plan can skip its own CPU inference if all nodes' shapes are known --- .../perf-runners/onnx-pr-bench/Cargo.lock | 11 +++ .../perf-runners/onnx-pr-bench/Cargo.toml | 1 + .../perf-runners/onnx-pr-bench/src/main.rs | 4 +- crates/yscv-onnx/src/runner/execute.rs | 6 ++ crates/yscv-onnx/src/runner/metal/compile.rs | 91 ++++++++++++------ crates/yscv-onnx/src/runner/mod.rs | 94 +++++++++++++++---- crates/yscv-onnx/src/runner/plan_branch.rs | 10 +- crates/yscv-onnx/src/runner/reshape.rs | 21 +++++ 8 files changed, 192 insertions(+), 46 deletions(-) diff --git a/benchmarks/perf-runners/onnx-pr-bench/Cargo.lock b/benchmarks/perf-runners/onnx-pr-bench/Cargo.lock index 0088a0f..929cd04 100644 --- a/benchmarks/perf-runners/onnx-pr-bench/Cargo.lock +++ b/benchmarks/perf-runners/onnx-pr-bench/Cargo.lock @@ -481,6 +481,12 @@ version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + [[package]] name = "rustix" version = "1.1.4" @@ -627,6 +633,7 @@ version = "0.1.10" name = "yscv-detect" version = "0.1.10" dependencies = [ + "rustc-hash", "thiserror", "yscv-tensor", "yscv-video", @@ -638,6 +645,7 @@ version = "0.1.10" dependencies = [ "cc", "rayon", + "rustc-hash", "thiserror", "yscv-cpu", "yscv-tensor", @@ -651,6 +659,7 @@ dependencies = [ "prost", "prost-build", "rayon", + "rustc-hash", "thiserror", "yscv-cpu", "yscv-kernels", @@ -664,6 +673,7 @@ version = "0.1.0" dependencies = [ "image", "mimalloc", + "rustc-hash", "serde_json", "yscv-detect", "yscv-kernels", @@ -696,6 +706,7 @@ dependencies = [ "arc-swap", "bytes", "rayon", + "rustc-hash", "thiserror", "yscv-cpu", "yscv-tensor", diff --git a/benchmarks/perf-runners/onnx-pr-bench/Cargo.toml b/benchmarks/perf-runners/onnx-pr-bench/Cargo.toml index 1625383..2a920da 100644 --- a/benchmarks/perf-runners/onnx-pr-bench/Cargo.toml +++ b/benchmarks/perf-runners/onnx-pr-bench/Cargo.toml @@ -9,6 +9,7 @@ publish = false [dependencies] image = { version = "0.25", default-features = false, features = ["jpeg", "png"] } mimalloc = { version = "0.1", default-features = false } +rustc-hash = "2.1.3" serde_json = "1" yscv-detect = { path = "../../../crates/yscv-detect" } yscv-kernels = { path = "../../../crates/yscv-kernels", default-features = false } diff --git a/benchmarks/perf-runners/onnx-pr-bench/src/main.rs b/benchmarks/perf-runners/onnx-pr-bench/src/main.rs index ed811a8..af40395 100644 --- a/benchmarks/perf-runners/onnx-pr-bench/src/main.rs +++ b/benchmarks/perf-runners/onnx-pr-bench/src/main.rs @@ -14,6 +14,8 @@ use yscv_onnx::{ }; use yscv_tensor::Tensor; +use rustc_hash::FxHashMap; + #[derive(Clone, Copy)] enum FillMode { Zero, @@ -686,7 +688,7 @@ 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 input_map: FxHashMap = 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); diff --git a/crates/yscv-onnx/src/runner/execute.rs b/crates/yscv-onnx/src/runner/execute.rs index ee2f828..bfa86c0 100644 --- a/crates/yscv-onnx/src/runner/execute.rs +++ b/crates/yscv-onnx/src/runner/execute.rs @@ -10,7 +10,9 @@ use super::*; pub(crate) fn run_onnx_model_jit( model: &OnnxModel, mut env: TensorEnv<'_, '_>, + specialization: Option<&ShapeSpecialization>, ) -> Result, OnnxError> { + let reshape_shapes = specialization.map(|plan| &plan.reshape_shapes); let branches = &model.runtime_index.node_branches; let use_counts_by_id = &model.runtime_index.use_counts_by_id; let output_id_mask = build_output_id_mask(model, &env, use_counts_by_id.len()); @@ -69,6 +71,7 @@ pub(crate) fn run_onnx_model_jit( &mut env0, &mut remaining0, &output_id_mask, + reshape_shapes, |nidx| branches_ref.get(nidx).copied() == Some(0), &mut c_ns, &mut o_ns, @@ -87,6 +90,7 @@ pub(crate) fn run_onnx_model_jit( &mut env1, &mut remaining1, &output_id_mask, + reshape_shapes, |nidx| branches_ref.get(nidx).copied() == Some(1), &mut c_ns, &mut o_ns, @@ -115,6 +119,7 @@ pub(crate) fn run_onnx_model_jit( &mut env, &mut remaining, &output_id_mask, + reshape_shapes, |nidx| { branches_ref.get(nidx).copied() != Some(0) && branches_ref.get(nidx).copied() != Some(1) @@ -132,6 +137,7 @@ pub(crate) fn run_onnx_model_jit( &mut env, &mut remaining_uses, &output_id_mask, + reshape_shapes, |_| true, &mut conv_ns, &mut other_ns, diff --git a/crates/yscv-onnx/src/runner/metal/compile.rs b/crates/yscv-onnx/src/runner/metal/compile.rs index 6fb57de..ad98a9c 100644 --- a/crates/yscv-onnx/src/runner/metal/compile.rs +++ b/crates/yscv-onnx/src/runner/metal/compile.rs @@ -10,6 +10,35 @@ use yscv_kernels::metal_backend::metal_conv::MetalInference; use crate::error::OnnxError; use crate::loader::{OnnxModel, OnnxNode}; +use crate::shape_infer::{ShapeMap, TensorShape, infer_shapes}; + +fn can_skip_cpu_shape_discovery(model: &OnnxModel, shapes: &ShapeMap) -> bool { + model.nodes.iter().all(|node| { + matches!( + node.op_type.as_str(), + "Conv" + | "Add" + | "Sub" + | "Mul" + | "Div" + | "Sigmoid" + | "Relu" + | "Concat" + | "Transpose" + | "Reshape" + | "MatMul" + ) && node + .outputs + .iter() + .filter(|name| !name.is_empty()) + .all(|name| { + shapes + .get(name) + .and_then(TensorShape::as_known_dims) + .is_some() + }) + }) +} /// Compile a Metal execution plan for the given ONNX model. /// Runs a shape-inference pass on CPU, then pre-allocates Metal buffers @@ -31,35 +60,43 @@ pub fn compile_metal_plan( let debug_metal = false; let mut env = TensorEnv::from_model(model); env.insert(input_name.to_string(), input_tensor.clone()); - // We need tensor shapes AND data for fallback ops. Some ops (Split) consume - // their inputs, so we snapshot shapes + data for fallback-eligible outputs - // immediately after each node executes. - let mut cpu_shapes: FxHashMap> = FxHashMap::default(); + let input_shapes: ShapeMap = FxHashMap::from_iter([( + input_name.to_string(), + TensorShape::known(input_tensor.shape().to_vec()), + )]); + let inferred = infer_shapes(model, &input_shapes); + // A supported, fully-known graph needs shapes but not a CPU execution. Keep + // the existing CPU walk for every other graph because fallback nodes may + // need their concrete values, not merely their output dimensions. + let skip_cpu_prepass = std::env::var("METAL_COMPARE").is_err() + && inferred.diagnostics.is_empty() + && can_skip_cpu_shape_discovery(model, &inferred.shapes); + let mut cpu_shapes: FxHashMap> = inferred + .shapes + .iter() + .filter_map(|(name, shape)| shape.as_known_dims().map(|dims| (name.clone(), dims))) + .collect(); let mut cpu_data: FxHashMap> = FxHashMap::default(); - for (ni, node) in model.nodes.iter().enumerate() { - if let Err(e) = execute_node_cpu_for_metal_compile(node, &mut env) - && debug_metal - { - eprintln!( - " [metal] CPU pass node {} {} '{}' FAILED: {}", - ni, node.op_type, node.name, e - ); - } - // Snapshot outputs that Metal will need for cpu_fallback - for out_name in &node.outputs { - if out_name.is_empty() { - continue; + if !skip_cpu_prepass { + for (ni, node) in model.nodes.iter().enumerate() { + if let Err(e) = execute_node_cpu_for_metal_compile(node, &mut env) + && debug_metal + { + eprintln!( + " [metal] CPU pass node {} {} '{}' FAILED: {}", + ni, node.op_type, node.name, e + ); } - if let Some(t) = env.get(out_name) { - cpu_shapes.insert(out_name.clone(), t.shape().to_vec()); - // Only save data for cpu_fallback-eligible ops (shape ops, etc.) - // to avoid excessive memory usage. - // Save data for any op that might need cpu_fallback - // (shape ops, unknown ops, etc.) — limit to small tensors to save memory - let n_elem = t.len(); - if n_elem <= 1_000_000 { - // ~4MB limit per tensor - cpu_data.insert(out_name.clone(), t.data().to_vec()); + // Snapshot outputs that Metal will need for cpu_fallback. + for out_name in &node.outputs { + if out_name.is_empty() { + continue; + } + if let Some(t) = env.get(out_name) { + cpu_shapes.insert(out_name.clone(), t.shape().to_vec()); + if t.len() <= 1_000_000 { + cpu_data.insert(out_name.clone(), t.data().to_vec()); + } } } } diff --git a/crates/yscv-onnx/src/runner/mod.rs b/crates/yscv-onnx/src/runner/mod.rs index a4ec3cf..f69f7ad 100644 --- a/crates/yscv-onnx/src/runner/mod.rs +++ b/crates/yscv-onnx/src/runner/mod.rs @@ -1,6 +1,8 @@ use rustc_hash::FxHashMap; pub(crate) use rustc_hash::FxHashSet; +use crate::shape_infer::{ShapeMap, TensorShape, infer_shapes}; + pub(crate) use yscv_kernels::{ BatchNorm2dParams, add as kernel_add, avg_pool2d_nhwc, batch_norm2d_nhwc, matmul_2d, max_pool2d_nhwc, mul as kernel_mul, relu, relu_inplace, sigmoid, softmax_last_dim, @@ -61,24 +63,62 @@ pub(crate) use tensor_env::*; // OnnxRunner — reusable inference session with configurable threading // --------------------------------------------------------------------------- +pub(crate) struct ShapeSpecialization { + input_shapes: Vec<(String, Vec)>, + reshape_shapes: FxHashMap>, +} + +fn shape_specialization( + model: &OnnxModel, + input_shapes: Vec<(String, Vec)>, +) -> ShapeSpecialization { + let shapes: ShapeMap = input_shapes + .iter() + .map(|(name, shape)| (name.clone(), TensorShape::known(shape.clone()))) + .collect(); + let inferred = infer_shapes(model, &shapes); + let reshape_shapes = model + .nodes + .iter() + .enumerate() + .filter(|(_, node)| node.op_type == "Reshape") + .filter_map(|(idx, node)| { + node.outputs + .first() + .and_then(|name| inferred.shapes.get(name)) + .and_then(TensorShape::as_known_dims) + .map(|shape| (idx, shape)) + }) + .collect(); + ShapeSpecialization { + input_shapes, + reshape_shapes, + } +} + +fn runtime_input_shapes( + model: &OnnxModel, + env: &TensorEnv<'_, '_>, +) -> Option)>> { + model + .inputs + .iter() + .filter(|name| !model.initializers.contains_key(*name)) + .map(|name| { + env.get(name) + .map(|tensor| (name.clone(), tensor.shape().to_vec())) + }) + .collect() +} + /// Reusable inference session with configurable threading. /// /// By default uses all CPU cores (like ONNX Runtime's `intra_op_num_threads`). /// Use `with_threads(1)` for single-threaded execution. -/// -/// ```rust,ignore -/// use yscv_onnx::*; -/// -/// let model = load_onnx_model("model.onnx")?; -/// let runner = OnnxRunner::new(&model)?; // all cores (default) -/// let runner_1t = OnnxRunner::with_threads(&model, 1)?; // single-thread -/// -/// let input = Tensor::from_vec(vec![1, 3, 640, 640], data)?; -/// let outputs = runner.run(&[("images", &input)])?; -/// ``` pub struct OnnxRunner<'m> { model: &'m OnnxModel, pool: Option, + shape_specialization: std::sync::OnceLock, /// Pool-agnostic scope for kernels migrated to `&dyn ParallelScope`. /// Built in lock-step with `pool`: when `pool == Some(p)`, this is a /// rayon-backed scope wrapping `p`; when `pool == None` it wraps @@ -147,6 +187,7 @@ impl<'m> OnnxRunner<'m> { Ok(Self { model, pool: None, + shape_specialization: std::sync::OnceLock::new(), parallel_scope: scope, single_thread: false, }) @@ -170,6 +211,7 @@ impl<'m> OnnxRunner<'m> { return Ok(Self { model, pool: None, + shape_specialization: std::sync::OnceLock::new(), parallel_scope: scope, single_thread: true, }); @@ -220,6 +262,7 @@ impl<'m> OnnxRunner<'m> { Ok(Self { model, pool, + shape_specialization: std::sync::OnceLock::new(), parallel_scope: scope, single_thread: false, }) @@ -252,6 +295,13 @@ impl<'m> OnnxRunner<'m> { } fn run_with_env(&self, env: TensorEnv<'_, '_>) -> Result, OnnxError> { + let specialization = runtime_input_shapes(self.model, &env).and_then(|input_shapes| { + let specialization = self + .shape_specialization + .get_or_init(|| shape_specialization(self.model, input_shapes.clone())); + (specialization.input_shapes == input_shapes).then_some(specialization) + }); + // Install the active ParallelScope on the inference thread so // migrated kernel sites (see `yscv_kernels::with_scope`) pick it // up without a signature-threading refactor. Dropped @@ -283,7 +333,7 @@ impl<'m> OnnxRunner<'m> { // Inline WITHOUT pinning does not help (the lone caller still // migrates), so this stays opt-in rather than the default. if std::env::var_os("YSCV_ST_INLINE").is_some() { - return run_onnx_model_inner(self.model, env); + return run_onnx_model_inner_specialized(self.model, env, specialization); } thread_local! { static ST_POOL: rayon::ThreadPool = rayon::ThreadPoolBuilder::new() @@ -291,7 +341,9 @@ impl<'m> OnnxRunner<'m> { .build() .expect("1-thread pool"); } - return ST_POOL.with(|pool| pool.install(|| run_onnx_model_inner(self.model, env))); + return ST_POOL.with(|pool| { + pool.install(|| run_onnx_model_inner_specialized(self.model, env, specialization)) + }); } // Optionally enter a session-scoped parallel region: inside @@ -311,11 +363,11 @@ impl<'m> OnnxRunner<'m> { // Run directly on the caller thread; yscv workers pick // up parallel work via `scope_ctx` when kernel sites // call through. - run_onnx_model_inner(model, env) + run_onnx_model_inner_specialized(model, env, specialization) } else if let Some(pool) = &self.pool { - pool.install(|| run_onnx_model_inner(model, env)) + pool.install(|| run_onnx_model_inner_specialized(model, env, specialization)) } else { - run_onnx_model_inner(model, env) + run_onnx_model_inner_specialized(model, env, specialization) } }; @@ -436,10 +488,18 @@ fn mark_skip_indices(skip: &mut [bool], indices: &[usize]) { fn run_onnx_model_inner( model: &OnnxModel, env: TensorEnv<'_, '_>, +) -> Result, OnnxError> { + run_onnx_model_inner_specialized(model, env, None) +} + +fn run_onnx_model_inner_specialized( + model: &OnnxModel, + 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); + return run_onnx_model_jit(model, env, specialization); } run_onnx_model_sequential(model, env) diff --git a/crates/yscv-onnx/src/runner/plan_branch.rs b/crates/yscv-onnx/src/runner/plan_branch.rs index a2d9619..9c8f913 100644 --- a/crates/yscv-onnx/src/runner/plan_branch.rs +++ b/crates/yscv-onnx/src/runner/plan_branch.rs @@ -11,6 +11,7 @@ pub(crate) fn execute_plan_branch( env: &mut TensorEnv<'_, '_>, remaining_uses: &mut [usize], output_id_mask: &[bool], + reshape_shapes: Option<&FxHashMap>>, mut accept: impl FnMut(usize) -> bool, conv_ns: &mut u64, other_ns: &mut u64, @@ -991,7 +992,14 @@ pub(crate) fn execute_plan_branch( NodeAction::Generic { node_idx, kind } => { let node = &nodes[*node_idx]; - execute_node_with_layout_kind(node, env, *kind)?; + if *kind == NodeKind::Reshape + && !env.is_nhwc(&node.inputs[0]) + && let Some(shape) = reshape_shapes.and_then(|shapes| shapes.get(node_idx)) + { + exec_reshape_known(node, env, shape)?; + } else { + execute_node_with_layout_kind(node, env, *kind)?; + } } } diff --git a/crates/yscv-onnx/src/runner/reshape.rs b/crates/yscv-onnx/src/runner/reshape.rs index cbfc111..0750ed1 100644 --- a/crates/yscv-onnx/src/runner/reshape.rs +++ b/crates/yscv-onnx/src/runner/reshape.rs @@ -21,6 +21,27 @@ pub(super) fn exec_reshape(node: &OnnxNode, env: &mut TensorEnv) -> Result<(), O 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( + node: &OnnxNode, + env: &mut TensorEnv, + shape: &[usize], +) -> Result<(), OnnxError> { + let input = get_tensor(env, &node.name, &node.inputs[0])?; + let expected_len = shape.iter().product::(); + if input.len() != expected_len { + return exec_reshape(node, env); + } + let out = input + .reshape(shape.to_vec()) + .map_err(|e| OnnxError::DecodeFailed { + message: e.to_string(), + })?; + env.insert(node.outputs[0].clone(), out); + 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`.