diff --git a/crates/onnx-runtime-ep-cpu-plugin/src/lib.rs b/crates/onnx-runtime-ep-cpu-plugin/src/lib.rs index b708999c3a..7d58475e49 100644 --- a/crates/onnx-runtime-ep-cpu-plugin/src/lib.rs +++ b/crates/onnx-runtime-ep-cpu-plugin/src/lib.rs @@ -160,6 +160,21 @@ pub extern "C" fn nxrt_ep_compiled_node_count() -> usize { onnx_runtime_ep_plugin::ep::compiled_node_count() } +/// Number of node inputs this EP reported to its kernels as session-lifetime +/// constants. Read through `dlopen` by the plugin E2E suite to prove weights +/// are recognised as weights on the ORT plugin path, where ORT presents a +/// fused node's initializers as ordinary inputs. +#[unsafe(no_mangle)] +pub extern "C" fn nxrt_ep_constant_weight_inputs() -> usize { + onnx_runtime_ep_plugin::ep::constant_weight_inputs() +} + +/// Resets the constant-weight-input counter to zero. +#[unsafe(no_mangle)] +pub extern "C" fn nxrt_ep_reset_constant_weight_inputs() { + onnx_runtime_ep_plugin::ep::reset_constant_weight_inputs() +} + /// Resets the compiled-node counter to zero. #[unsafe(no_mangle)] pub extern "C" fn nxrt_ep_reset_compiled_node_count() { diff --git a/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_export_abi.rs b/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_export_abi.rs index 8a58eb3387..10b598287b 100644 --- a/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_export_abi.rs +++ b/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_export_abi.rs @@ -491,6 +491,8 @@ fn l1_required_symbols_resolve() { &b"nxrt_ep_reset_compiled_node_count"[..], &b"nxrt_ep_workspace_placement_queries"[..], &b"nxrt_ep_reset_workspace_placement_queries"[..], + &b"nxrt_ep_constant_weight_inputs"[..], + &b"nxrt_ep_reset_constant_weight_inputs"[..], ] { let _counter: libloading::Symbol<'_, unsafe extern "C" fn() -> usize> = unsafe { lib.get(symbol) }.unwrap_or_else(|e| { @@ -576,6 +578,8 @@ fn l1_no_symbol_leakage() { && *name != "nxrt_ep_reset_compiled_node_count" && *name != "nxrt_ep_workspace_placement_queries" && *name != "nxrt_ep_reset_workspace_placement_queries" + && *name != "nxrt_ep_constant_weight_inputs" + && *name != "nxrt_ep_reset_constant_weight_inputs" && !name.starts_with("_Z") && !name.starts_with("__rust") && !name.starts_with("__rdl_") diff --git a/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs b/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs index 649e69754b..4741c96c65 100644 --- a/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs +++ b/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs @@ -4403,6 +4403,7 @@ fn exp_deferral_still_claims_when_cpu_fallback_is_disabled() { // instead, which cargo provides per test binary and cleans with `cargo clean`. /// A generated matmul-family model plus everything needed to run it twice. +#[derive(Clone)] struct MatmulFamilyCase { /// Fixture name; also the file stem under `CARGO_TARGET_TMPDIR`. name: &'static str, @@ -4428,6 +4429,7 @@ struct MatmulFamilyCase { } /// One runtime input: name, element type, dims and raw little-endian bytes. +#[derive(Clone)] struct GeneratedInput { name: &'static str, elem: ort::ONNXTensorElementDataType, @@ -5012,6 +5014,34 @@ unsafe fn try_plain_ort_session( } } +/// Byte width of one element of `elem`, for the element types this file +/// generates. +fn elem_byte_size(elem: ort::ONNXTensorElementDataType) -> usize { + match elem { + ELEM_F32 => 4, + ELEM_F16 => 2, + ELEM_U8 | ELEM_I8 => 1, + other => panic!("unhandled element type {other}"), + } +} + +/// The same case with its activation tensor rotated by one whole element. +/// +/// Rotating by an element rather than by a byte keeps every value a valid +/// number of its dtype — a byte rotation of an f32 buffer would manufacture +/// NaNs and denormals and turn a parity check into a float-corner-case test. +fn with_rotated_activation(case: &MatmulFamilyCase) -> MatmulFamilyCase { + let mut rotated = case.clone(); + let activation = &mut rotated.inputs[0]; + let width = elem_byte_size(activation.elem); + assert!( + activation.data.len() > width, + "activation is a single element; rotating it is a no-op" + ); + activation.data.rotate_left(width); + rotated +} + /// The matmul-family half of the architectural rule, asserted end to end. /// /// For every case this test, in one session with `session.disable_cpu_ep_fallback=1`: @@ -5035,6 +5065,7 @@ unsafe fn try_plain_ort_session( fn no_matmul_family_node_escapes_to_the_ort_cpu_ep() { let _lock = lock_ort_ep(); let mut checked = 0usize; + let mut sensitive = 0usize; for case in matmul_family_cases() { let path = write_generated_model(case.name, &case.model); let reg = format!("cpu_ep_mm_{}", case.name); @@ -5134,6 +5165,72 @@ fn no_matmul_family_node_escapes_to_the_ort_cpu_ep() { ort_out[worst_at], case.tolerance ); + + // Second run, same sessions, different activations. The + // EP now tells each kernel which of its inputs ORT holds + // constant, and kernels use that to keep a prepacked + // weight for the life of the session instead of rebuilding + // it per call. Caching a weight is sound; caching anything + // activation-derived is not, and the two are easy to + // confuse in a kernel that only ever sees one input set. + // Re-running the *same session* with rotated activations + // separates them: a weight cache still matches ORT, a + // stale activation cache returns the first answer. + let rotated = with_rotated_activation(&case); + let ours_second = run_generated_case( + api, + session, + &rotated, + &format!("Run({} rotated)", case.name), + ); + let ort_second = run_generated_case( + api, + ort_session, + &rotated, + &format!("Run(ORT baseline {} rotated)", case.name), + ); + // ORT is the oracle for whether this case can detect a + // stale cache at all. A saturating u8 `QLinearMatMul` + // output is unchanged by a one-element rotation for both + // implementations, which makes the check vacuous *for that + // case* rather than wrong; the suite-level counter below + // is what keeps the sweep as a whole non-vacuous. + let ort_moved = ort_second + .iter() + .zip(ort_out.iter()) + .any(|(a, b)| (a - b).abs() > case.tolerance); + if ort_moved { + sensitive += 1; + assert!( + ours_second + .iter() + .zip(ours_out.iter()) + .any(|(a, b)| (a - b).abs() > case.tolerance), + "{}: ORT's result changed when the activation rotated and ours \ + did not — this EP returned a stale answer on the second run", + case.name + ); + } + let mut worst_second = 0.0f32; + let mut worst_second_at = 0usize; + for (i, (a, b)) in ours_second.iter().zip(ort_second.iter()).enumerate() { + let delta = (a - b).abs(); + if delta > worst_second { + worst_second = delta; + worst_second_at = i; + } + } + assert!( + worst_second <= case.tolerance, + "{}: second run with new activations diverged from ORT — element \ + {worst_second_at} is {} against ORT's {} (delta {worst_second}, \ + tolerance {}). A per-session weight cache must not outlive the \ + activations it was used with", + case.name, + ours_second[worst_second_at], + ort_second[worst_second_at], + case.tolerance + ); ((*api).ReleaseSession.unwrap())(ort_session); ((*api).ReleaseSessionOptions.unwrap())(ort_opts); } @@ -5161,7 +5258,15 @@ fn no_matmul_family_node_escapes_to_the_ort_cpu_ep() { checked, 10, "every generated matmul case must have been run" ); - eprintln!("\n✅ no_matmul_family_node_escapes_to_the_ort_cpu_ep: {checked} cases PASSED"); + assert!( + sensitive >= 5, + "only {sensitive} of {checked} cases changed their result when the activation \ + rotated, so the stale-cache half of this test is mostly vacuous" + ); + eprintln!( + "\n✅ no_matmul_family_node_escapes_to_the_ort_cpu_ep: {checked} cases PASSED \ + ({sensitive} of them activation-sensitive on the second run)" + ); } } @@ -5349,3 +5454,500 @@ opset_import: [{{ version: 17 }}, {{ domain: "com.microsoft" version: 1 }}] conformance_teardown(api, env, opts, session, reg); } } + +/// Weights must reach the kernels *labelled* as weights. +/// +/// ORT gives a plugin EP a fused node whose initializers are inputs of that +/// node and graph inputs of its subgraph, so at Compile time a 1 GB +/// `MatMulNBits` `B` is shaped exactly like an activation. Every kernel that +/// builds a prepack — `MatMulNBits`, `QLinearMatMul`, the f16 `MatMul` widening +/// cache — decides whether that prepack may outlive the call from +/// `Kernel::set_constant_inputs`, so getting this wrong silently converts a +/// once-per-session cost into a per-token one, and in `MatMulNBits`'s case also +/// selects a slower kernel (its MLAS SQNBit path is gated on the same flag). +/// +/// No output comparison can see any of that, which is why the EP exports a +/// counter and this test reads it from the very cdylib ORT loaded. +/// +/// The cases are chosen to move the counter in opposite directions: +/// symmetric `nbits4_decode` has two initializers (`B`, `scales`), +/// asymmetric `nbits8_prefill` has three (`zero_points` as well), and +/// `matmul_f32_decode` declares both operands as graph inputs so it must +/// report nothing. A wiring that marks everything constant passes the first +/// two assertions and fails the third — and would be a real bug, not merely a +/// slower one, because a prepack of an activation must never outlive the call. +#[test] +fn constant_weights_are_reported_to_kernels_as_constant() { + let _lock = lock_ort_ep(); + let ep_path = match find_ep_cdylib() { + Some(p) => p, + None => { + if std::env::var("NXRT_REQUIRE_ORT_TESTS").as_deref() == Ok("1") { + panic!("NXRT_REQUIRE_ORT_TESTS=1 but the EP cdylib is missing"); + } + eprintln!("*** SKIPPED: EP cdylib not found ***"); + return; + } + }; + // Same absolute path ORT registers, so this is the same mapping and the + // same statics — not a second copy of the library. + let ep_lib = unsafe { libloading::Library::new(&ep_path) }.expect("dlopen EP cdylib"); + let read: libloading::Symbol<'_, unsafe extern "C" fn() -> usize> = + unsafe { ep_lib.get(b"nxrt_ep_constant_weight_inputs") } + .expect("nxrt_ep_constant_weight_inputs not exported"); + let reset: libloading::Symbol<'_, unsafe extern "C" fn()> = + unsafe { ep_lib.get(b"nxrt_ep_reset_constant_weight_inputs") } + .expect("nxrt_ep_reset_constant_weight_inputs not exported"); + + for (name, expected) in [ + ("nbits4_decode", 2usize), + ("nbits8_prefill", 3usize), + ("matmul_f32_decode", 0usize), + ] { + let case = matmul_family_cases() + .into_iter() + .find(|c| c.name == name) + .expect("case present"); + let path = write_generated_model(case.name, &case.model); + let reg = format!("cpu_ep_constflags_{}", case.name); + unsafe { reset() }; + let Some((_lib, api, env, opts, session)) = + (unsafe { conformance_setup(®, &path, true) }) + else { + eprintln!("*** SKIPPED: {name} — ORT not found ***"); + return; + }; + let observed = unsafe { read() }; + assert_eq!( + observed, expected, + "{name}: this EP reported {observed} constant inputs, expected {expected}" + ); + unsafe { conformance_teardown(api, env, opts, session, ®) }; + eprintln!(" [{name}] {observed} constant inputs reported"); + } +} + +/// A `MatMulNBits` whose `B` and `scales` are declared **both** as +/// initializers and as graph inputs. +/// +/// From ONNX IR version 4 that combination means the initializer is only a +/// default: the caller may hand a different tensor in on any `Run`, and ORT +/// says so in its load log ("Initializer B appears in graph inputs and will +/// not be treated as constant value/weight"). `Graph_GetInitializers` still +/// lists it — its own documentation says it "includes constant and +/// non-constant initializers" — so a name-only membership test would call it a +/// weight, and this crate would cache a prepack of the default for the life of +/// the session. +/// +/// Returns the case with every operand supplied at run time, plus a second +/// `B` payload that produces different outputs. +fn overridable_nbits_case() -> (MatmulFamilyCase, Vec) { + const M: usize = 1; + const K: usize = 64; + const N: usize = 64; + const BLOCK: usize = 32; + let blocks_per_col = K / BLOCK; + let blob = BLOCK / 2; + let b_elements = N * blocks_per_col * blob; + let scale_elements = N * blocks_per_col; + let b_dims = vec![N as i64, blocks_per_col as i64, blob as i64]; + let scale_dims = vec![scale_elements as i64]; + let a_dims = vec![M as i64, K as i64]; + + let default_b = printable_weight_blob(b_elements); + let scales = printable_scale_blob(scale_elements); + let model = format!( + r#"ir_version: 10 +graph {{ + node: [{{ + input: ["A", "B", "scales"] + output: ["Y"] + op_type: "MatMulNBits" + domain: "com.microsoft" + attribute: [ + {{ name: "K" type: INT i: {K} }}, + {{ name: "N" type: INT i: {N} }}, + {{ name: "bits" type: INT i: 4 }}, + {{ name: "block_size" type: INT i: {BLOCK} }} + ] + }}] + name: "nbits4_overridable" + initializer: [{}, {}] + input: [{}, {}, {}] + output: [{}] +}} +opset_import: [{{ version: 17 }}, {{ domain: "com.microsoft" version: 1 }}] +"#, + textproto_initializer("B", ELEM_U8, &b_dims, &default_b), + textproto_initializer("scales", ELEM_F32, &scale_dims, &scales), + textproto_value_info("A", ELEM_F32, &a_dims), + textproto_value_info("B", ELEM_U8, &b_dims), + textproto_value_info("scales", ELEM_F32, &scale_dims), + textproto_value_info("Y", ELEM_F32, &[M as i64, N as i64]), + ); + + // Bit-complement of the default, so every 4-bit nibble changes and no + // output element can coincide by construction. + let other_b: Vec = default_b.as_bytes().iter().map(|b| !*b).collect(); + let case = MatmulFamilyCase { + name: "nbits4_overridable", + op: "MatMulNBits", + model, + inputs: vec![ + GeneratedInput { + name: "A", + elem: ELEM_F32, + dims: a_dims, + data: f32_slice_to_bytes(&activation_f32(M, K)), + }, + GeneratedInput { + name: "B", + elem: ELEM_U8, + dims: b_dims, + data: default_b.as_bytes().to_vec(), + }, + GeneratedInput { + name: "scales", + elem: ELEM_F32, + dims: scale_dims, + data: scales.as_bytes().to_vec(), + }, + ], + output: "Y", + output_elem: ELEM_F32, + ort_can_build: true, + tolerance: 0.5, + }; + (case, other_b) +} + +/// An initializer the caller may override is **not** a constant weight, and +/// overriding it must change the answer. +/// +/// Two independent falsifiers, because either one alone is weak: the counter +/// pins the classification (it reads 2 if membership is keyed on +/// `Graph_GetInitializers` names alone), and the numeric check pins the +/// consequence — a session-lifetime prepack of the default would return run +/// one's answer for run two, which ORT, used here as the oracle, does not. +#[test] +fn an_overridable_initializer_is_not_treated_as_a_constant_weight() { + let _lock = lock_ort_ep(); + let ep_path = match find_ep_cdylib() { + Some(p) => p, + None => { + if std::env::var("NXRT_REQUIRE_ORT_TESTS").as_deref() == Ok("1") { + panic!("NXRT_REQUIRE_ORT_TESTS=1 but the EP cdylib is missing"); + } + eprintln!("*** SKIPPED: EP cdylib not found ***"); + return; + } + }; + let ep_lib = unsafe { libloading::Library::new(&ep_path) }.expect("dlopen EP cdylib"); + let read: libloading::Symbol<'_, unsafe extern "C" fn() -> usize> = + unsafe { ep_lib.get(b"nxrt_ep_constant_weight_inputs") } + .expect("nxrt_ep_constant_weight_inputs not exported"); + let reset: libloading::Symbol<'_, unsafe extern "C" fn()> = + unsafe { ep_lib.get(b"nxrt_ep_reset_constant_weight_inputs") } + .expect("nxrt_ep_reset_constant_weight_inputs not exported"); + + let (case, other_b) = overridable_nbits_case(); + let mut overridden = case.clone(); + overridden.inputs[1].data = other_b; + let path = write_generated_model(case.name, &case.model); + let reg = "cpu_ep_overridable_init"; + unsafe { reset() }; + let Some((_lib, api, env, opts, session)) = (unsafe { conformance_setup(reg, &path, true) }) + else { + eprintln!("*** SKIPPED: ORT not found ***"); + return; + }; + let observed = unsafe { read() }; + let ours_default = unsafe { run_generated_case(api, session, &case, "overridable/default") }; + let ours_other = unsafe { run_generated_case(api, session, &overridden, "overridable/other") }; + unsafe { conformance_teardown(api, env, opts, session, reg) }; + + assert_eq!( + observed, 0, + "an initializer that is also a graph input may be replaced on any Run, \ + but this EP reported {observed} of them as constant weights" + ); + + // ORT as the oracle: same model, same two payloads, its own CPU EP. + let (ort_default, ort_other) = { + let Some((_lib2, api2, env2, opts2, session2)) = + (unsafe { conformance_setup("cpu_ep_overridable_oracle", &path, true) }) + else { + eprintln!("*** SKIPPED: ORT not found for the oracle run ***"); + return; + }; + let plain = unsafe { try_plain_ort_session(api2, env2, &path) } + .expect("ORT can build this MatMulNBits"); + let d = unsafe { run_generated_case(api2, plain.1, &case, "oracle/default") }; + let o = unsafe { run_generated_case(api2, plain.1, &overridden, "oracle/other") }; + unsafe { + ((*api2).ReleaseSession.unwrap())(plain.1); + ((*api2).ReleaseSessionOptions.unwrap())(plain.0); + } + unsafe { conformance_teardown(api2, env2, opts2, session2, "cpu_ep_overridable_oracle") }; + (d, o) + }; + + assert!( + ort_default + .iter() + .zip(ort_other.iter()) + .any(|(a, b)| (a - b).abs() > case.tolerance), + "the two B payloads must produce different outputs or this test proves nothing" + ); + for (stage, ours, theirs) in [ + ("default", &ours_default, &ort_default), + ("overridden", &ours_other, &ort_other), + ] { + assert_eq!(ours.len(), theirs.len(), "{stage}: output length"); + for (i, (a, b)) in ours.iter().zip(theirs.iter()).enumerate() { + assert!( + (a - b).abs() <= case.tolerance, + "{stage}: element {i} is {a} but ORT says {b} — a stale prepack of the \ + default initializer would look exactly like this" + ); + } + } + eprintln!(" [overridable] 0 constant inputs; both payloads match ORT"); +} + +// ─── Plugin-path A/B benchmark ──────────────────────────────────────────────── + +/// Time `iters` `Run` calls over pre-created input tensors, returning +/// per-iteration milliseconds. +/// +/// Input `OrtValue`s are built once and reused so the measurement is the +/// kernel plus ORT's per-run overhead, not the harness copying weights. +/// +/// # Safety +/// `api`, `session` and every pointer in `values` must be valid. +unsafe fn bench_runs( + api: *const ort::OrtApi, + session: *mut ort::OrtSession, + input_name_ptrs: &[*const std::os::raw::c_char], + values: &[*const ort::OrtValue], + output_name_ptrs: &[*const std::os::raw::c_char], + iters: usize, +) -> Vec { + let mut out = Vec::with_capacity(iters); + for _ in 0..iters { + let mut output: *mut ort::OrtValue = ptr::null_mut(); + let start = std::time::Instant::now(); + unsafe { + let status = ((*api).Run.unwrap())( + session, + ptr::null(), + input_name_ptrs.as_ptr(), + values.as_ptr(), + values.len(), + output_name_ptrs.as_ptr(), + 1, + &mut output, + ); + check_status(api, status, "bench Run"); + out.push(start.elapsed().as_secs_f64() * 1e3); + ((*api).ReleaseValue.unwrap())(output); + } + } + out +} + +fn percentile(samples: &mut [f64], p: f64) -> f64 { + samples.sort_by(|a, b| a.partial_cmp(b).unwrap()); + let idx = ((samples.len() - 1) as f64 * p).round() as usize; + samples[idx] +} + +/// Shapes worth measuring on the plugin path, at sizes a real decode/prefill +/// step reaches rather than the smallest size that proves correctness. +fn matmul_bench_cases() -> Vec { + vec![ + dense_case( + "bench_matmul_f32_m1", + "MatMul", + ELEM_F32, + 1, + 2048, + 2048, + 1e-3, + ), + dense_case( + "bench_matmul_f32_m128", + "MatMul", + ELEM_F32, + 128, + 2048, + 2048, + 1e-3, + ), + dense_case( + "bench_matmul_f16_m1", + "MatMul", + ELEM_F16, + 1, + 2048, + 2048, + 5e-2, + ), + dense_case( + "bench_matmul_f16_m128", + "MatMul", + ELEM_F16, + 128, + 2048, + 2048, + 5e-2, + ), + nbits_case("bench_nbits4_m1", 1, 2048, 2048, 4, 32, ELEM_F32, true), + nbits_case("bench_nbits4_m128", 128, 2048, 2048, 4, 32, ELEM_F32, true), + nbits_case("bench_nbits4_f16_m1", 1, 2048, 2048, 4, 32, ELEM_F16, true), + nbits_case("bench_nbits8_m1", 1, 2048, 2048, 8, 32, ELEM_F32, true), + nbits_case("bench_nbits8_m256", 256, 2048, 2048, 8, 32, ELEM_F32, true), + qlinear_case("bench_qlinear_u8_m1", 1, 2048, 2048, false), + qlinear_case("bench_qlinear_u8_m128", 128, 2048, 2048, false), + qlinear_case("bench_qlinear_i8_m1", 1, 2048, 2048, true), + ] +} + +/// Interleaved A/B of this EP against plain ORT **through the ORT session +/// API**, which is the only path a real model takes. +/// +/// Every previously published matmul ratio came from `bench_generic`, which +/// drives the kernels natively. Those numbers stay valid as kernel-level +/// measurements, but until the claim fixes landed the plugin path never +/// reached the quantized kernels at all, so this is the first honest +/// session-level comparison for `MatMulNBits` and `QLinearMatMul`. +/// +/// Off by default: it takes minutes and this host is shared. Run with +/// `NXRT_MM_BENCH=1 cargo test -p onnx-runtime-ep-cpu-plugin --release +/// --test plugin_ort_e2e plugin_path_ab -- --nocapture --ignored`. +#[test] +#[ignore = "benchmark, not a correctness test"] +fn plugin_path_ab_vs_plain_ort() { + if std::env::var("NXRT_MM_BENCH").as_deref() != Ok("1") { + eprintln!("set NXRT_MM_BENCH=1 to run the plugin-path A/B"); + return; + } + let iters: usize = std::env::var("NXRT_MM_BENCH_ITERS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(31); + let filter = std::env::var("NXRT_MM_BENCH_CASE").unwrap_or_default(); + let _lock = lock_ort_ep(); + println!( + "case,ours_p50_ms,ort_p50_ms,ratio_p50,ours_p90_ms,ort_p90_ms,ratio_p90,cold_ours_ms,cold_ort_ms" + ); + for case in matmul_bench_cases() { + if !filter.is_empty() && !case.name.contains(&filter) { + continue; + } + let path = write_generated_model(case.name, &case.model); + let reg = format!("cpu_ep_bench_{}", case.name); + let cold_ours = std::time::Instant::now(); + let Some((_lib, api, env, opts, session)) = + (unsafe { conformance_setup(®, &path, true) }) + else { + eprintln!( + "*** SKIPPED: {} — ORT or EP cdylib not found ***", + case.name + ); + return; + }; + let cold_ours = cold_ours.elapsed().as_secs_f64() * 1e3; + + unsafe { + let info = query_ep_assignment(api, session); + assert!( + info.ops_on_our_ep().contains(&case.op), + "{}: not assigned to this EP, refusing to report a ratio", + case.name + ); + + let cold_ort = std::time::Instant::now(); + let (ort_opts, ort_session) = match try_plain_ort_session(api, env, &path) { + Ok(pair) => pair, + Err(msg) => { + eprintln!("{}: ORT cannot build this model ({msg})", case.name); + conformance_teardown(api, env, opts, session, ®); + continue; + } + }; + let cold_ort = cold_ort.elapsed().as_secs_f64() * 1e3; + + let mut buffers: Vec> = case.inputs.iter().map(|i| i.data.clone()).collect(); + let mut values: Vec<*const ort::OrtValue> = Vec::with_capacity(buffers.len()); + for (input, buffer) in case.inputs.iter().zip(buffers.iter_mut()) { + values.push(make_raw_tensor(api, buffer, &input.dims, input.elem) as *const _); + } + let input_names: Vec = case + .inputs + .iter() + .map(|i| CString::new(i.name).unwrap()) + .collect(); + let input_name_ptrs: Vec<*const std::os::raw::c_char> = + input_names.iter().map(|c| c.as_ptr()).collect(); + let output_name = CString::new(case.output).unwrap(); + let output_name_ptrs = [output_name.as_ptr()]; + + // Warm both sides: first run pays prepack and page faults. + bench_runs( + api, + session, + &input_name_ptrs, + &values, + &output_name_ptrs, + 3, + ); + bench_runs( + api, + ort_session, + &input_name_ptrs, + &values, + &output_name_ptrs, + 3, + ); + + // Interleave one iteration each so a drifting host load lands on + // both sides rather than on whichever ran second. + let mut ours = Vec::with_capacity(iters); + let mut theirs = Vec::with_capacity(iters); + for _ in 0..iters { + ours.extend(bench_runs( + api, + session, + &input_name_ptrs, + &values, + &output_name_ptrs, + 1, + )); + theirs.extend(bench_runs( + api, + ort_session, + &input_name_ptrs, + &values, + &output_name_ptrs, + 1, + )); + } + let (o50, o90) = (percentile(&mut ours, 0.5), percentile(&mut ours, 0.9)); + let (t50, t90) = (percentile(&mut theirs, 0.5), percentile(&mut theirs, 0.9)); + println!( + "{},{o50:.4},{t50:.4},{:.3},{o90:.4},{t90:.4},{:.3},{cold_ours:.1},{cold_ort:.1}", + case.name, + o50 / t50, + o90 / t90, + ); + + for value in values { + ((*api).ReleaseValue.unwrap())(value as *mut _); + } + ((*api).ReleaseSession.unwrap())(ort_session); + ((*api).ReleaseSessionOptions.unwrap())(ort_opts); + conformance_teardown(api, env, opts, session, ®); + } + } +} diff --git a/crates/onnx-runtime-ep-plugin/src/ep.rs b/crates/onnx-runtime-ep-plugin/src/ep.rs index 9fb4b31da1..da4cf9ed79 100644 --- a/crates/onnx-runtime-ep-plugin/src/ep.rs +++ b/crates/onnx-runtime-ep-plugin/src/ep.rs @@ -72,6 +72,27 @@ pub fn reset_compiled_node_count() { COMPILED_NODE_COUNT.store(0, Ordering::Relaxed); } +/// Global counter of node inputs this EP told a kernel to treat as +/// session-lifetime constants (see [`constant_input_flags`]). +/// +/// A weight that is not reported constant is not a correctness bug, which is +/// exactly why it needs a counter: the kernel still computes the right answer, +/// it just rebuilds its prepack on every `Run` — and `MatMulNBits` also drops +/// to a slower kernel, because its MLAS SQNBit path is gated on the same flag. +/// Nothing in an output comparison can see that, so the wiring is observed +/// here instead. +static CONSTANT_WEIGHT_INPUTS: AtomicUsize = AtomicUsize::new(0); + +/// Number of node inputs reported to kernels as constant since process start. +pub fn constant_weight_inputs() -> usize { + CONSTANT_WEIGHT_INPUTS.load(Ordering::Relaxed) +} + +/// Resets the constant-weight-input counter to zero. +pub fn reset_constant_weight_inputs() { + CONSTANT_WEIGHT_INPUTS.store(0, Ordering::Relaxed); +} + use crate::compute::ExportedComputeInfo; use crate::graph_reader::OutboundGraphReader; use crate::status::{fail_status, invalid_arg_status, ok_status}; @@ -766,7 +787,21 @@ fn ep_compile_inner( let opset = ir_graph.effective_opset(node).unwrap_or(0); match exported.ep.with(|ep| ep.get_kernel(node, &shapes, opset)) { - Ok(kernel) => { + Ok(mut kernel) => { + // Tell the kernel which of its inputs are session-lifetime + // constants. Without this every kernel sees all inputs as + // runtime tensors and re-does its one-time weight work on + // every `Run`: `MatMulNBits` repacks (or, worse, declines + // the MLAS SQNBit path, which is gated on the same flag) + // and `QLinearMatMul` re-packs B, turning a load-time cost + // into a per-token cost. + let constant_inputs = + constant_input_flags(&view, node_idx, reader.constant_initializer_names()); + CONSTANT_WEIGHT_INPUTS.fetch_add( + constant_inputs.iter().filter(|c| **c).count(), + Ordering::Relaxed, + ); + kernel.set_constant_inputs(&constant_inputs); let num_inputs = view.node_inputs(node_idx).len(); let num_outputs = view.node_outputs(node_idx).len(); @@ -917,6 +952,50 @@ fn ep_compile_inner( ok_status() } +/// Which inputs of `node` are session-lifetime constants. +/// +/// A node input that ORT lists as a *constant* initializer is a weight: ORT +/// owns the buffer, materializes it once at session creation and cannot change +/// it between `Run` calls. An IR>=4 initializer that also appears as a graph +/// input is only a default value the caller may override per `Run`, and is +/// deliberately not in that set (see `read_constant_initializer_names`). That is the +/// contract +/// [`onnx_runtime_ep_api::kernel::Kernel::set_constant_inputs`] expresses, and +/// several kernels use it to decide whether a prepack may be built once and +/// kept in the kernel instance (which lives as long as the session) instead of +/// rebuilt on every call. `MatMulNBits` goes further and gates its MLAS SQNBit +/// path on the same flag, so leaving it false does not merely repeat work — it +/// selects a different, slower kernel. +/// +/// The producer/`is_graph_input` test used elsewhere in this file cannot be +/// reused here. It is correct for the whole-model graph at capability time, +/// but ORT hands a *fused node's* subgraph over with the initializers it kept +/// inside listed as graph inputs of that subgraph, so at Compile time every +/// weight looks like an activation. `Graph_GetInitializers` still distinguishes +/// them, which is why the flags are keyed by name against that set — filtered +/// to the entries `ValueInfo_IsConstantInitializer` accepts, because +/// `Graph_GetInitializers` also lists IR>=4 defaults the caller may replace. +/// +/// Absent optional inputs and unnamed values are never constant: there is +/// nothing to cache and nothing to key a cache on. +fn constant_input_flags( + view: &onnx_runtime_ir::GraphView<'_>, + node: NodeIndex, + constant_initializer_names: &std::collections::HashSet, +) -> Vec { + view.node_inputs(node) + .iter() + .map(|input| { + input.is_some_and(|value| { + view.value(value) + .name + .as_deref() + .is_some_and(|name| constant_initializer_names.contains(name)) + }) + }) + .collect() +} + /// Whether every present input of `node` can be routed by the fused-subgraph /// compile path (`build_subgraph_routing`). /// @@ -2075,6 +2154,79 @@ mod tests { ); } + /// Constant-input flags must come from ORT's initializer list, not from + /// the producer/`is_graph_input` shape of the graph. + /// + /// This is the exact shape ORT hands a fused node's subgraph over in: + /// the weight has no producer *and* is listed as a graph input of the + /// subgraph, because it is an input of the fused node. Deriving the flag + /// the way `node_inputs_all_routable` derives routability marks every + /// weight non-constant, which costs `MatMulNBits` its prepack — and, since + /// its MLAS SQNBit path is gated on the same flag, its fast kernel too. + #[test] + fn constant_initializers_are_flagged_even_when_the_subgraph_calls_them_inputs() { + use onnx_runtime_ir::{Graph, GraphView, GraphViewCache, Node, NodeId, Shape}; + let mut g = Graph::new(); + let a = g.create_named_value("a", DataType::Float32, Shape::default()); + let w = g.create_named_value("w", DataType::Uint8, Shape::default()); + let scales = g.create_named_value("scales", DataType::Float32, Shape::default()); + let y = g.create_named_value("y", DataType::Float32, Shape::default()); + // Every input of the fused node is a graph input of the subgraph ORT + // hands over, weights included. + g.add_input(a); + g.add_input(w); + g.add_input(scales); + g.add_output(y); + g.insert_node(Node::new( + NodeId(0), + "MatMulNBits", + vec![Some(a), Some(w), Some(scales), None], + vec![y], + )); + + let cache = GraphViewCache::build(&g).unwrap(); + let view = GraphView::new(&g, &cache); + let node = view.nodes().next().expect("one node"); + let constant_initializers: std::collections::HashSet = + ["w".to_owned(), "scales".to_owned()].into_iter().collect(); + + assert_eq!( + super::constant_input_flags(&view, node, &constant_initializers), + vec![false, true, true, false], + "weights are constant, the activation is not, and an absent optional \ + input is not" + ); + } + + /// With no initializers, nothing is constant — including values that have + /// no producer. A subgraph whose weights ORT kept outside it must not have + /// its activations cached as if they were weights. + #[test] + fn nothing_is_constant_without_an_initializer_list() { + use onnx_runtime_ir::{Graph, GraphView, GraphViewCache, Node, NodeId, Shape}; + let mut g = Graph::new(); + let a = g.create_named_value("a", DataType::Float32, Shape::default()); + let b = g.create_named_value("b", DataType::Float32, Shape::default()); + let y = g.create_named_value("y", DataType::Float32, Shape::default()); + g.add_input(a); + g.add_output(y); + g.insert_node(Node::new( + NodeId(0), + "MatMul", + vec![Some(a), Some(b)], + vec![y], + )); + + let cache = GraphViewCache::build(&g).unwrap(); + let view = GraphView::new(&g, &cache); + let node = view.nodes().next().expect("one node"); + assert_eq!( + super::constant_input_flags(&view, node, &std::collections::HashSet::new()), + vec![false, false], + "a producerless value is only constant when ORT says it is an initializer" + ); + } + /// B2 (claim-time routability gate): a node consuming a weight initializer /// (a value with no producer that is not a graph input) is NOT routable by /// the fused compile path and must be declined at claim time so ORT diff --git a/crates/onnx-runtime-ep-plugin/src/graph_reader.rs b/crates/onnx-runtime-ep-plugin/src/graph_reader.rs index ce722c7982..5506f42597 100644 --- a/crates/onnx-runtime-ep-plugin/src/graph_reader.rs +++ b/crates/onnx-runtime-ep-plugin/src/graph_reader.rs @@ -22,6 +22,21 @@ pub struct OutboundGraphReader { /// Owned copies of small int64 initializer tensors, keyed by value name. /// Used to resolve opset-13 Unsqueeze/Squeeze axes from constant inputs. initializer_int64: HashMap>, + /// Names of every **constant** initializer in this graph, whatever its + /// dtype. + /// + /// ORT hands a fused node's subgraph over with its initializers listed as + /// *graph inputs* of that subgraph, so `is_graph_input` cannot distinguish + /// a weight from an activation at Compile time. This set can. + /// + /// `Graph_GetInitializers` alone is not enough: its own documentation says + /// it "includes constant and non-constant initializers", and from ONNX IR + /// version 4 an initializer that is also declared as a graph input is only + /// a *default* — the caller may replace it at every `Run`. Caching a + /// prepack of one of those would return the default value's answer forever. + /// Membership here therefore also requires + /// `ValueInfo_IsConstantInitializer`. + constant_initializer_names: HashSet, /// Out-of-band set of ValueIds that represent absent optional output slots. /// This replaces the previous in-band string sentinel (`__absent_output_*`) /// which was forgeable from model content. ValueIds are graph-internal @@ -154,6 +169,8 @@ impl OutboundGraphReader { // Read initializers and copy small int64 tensors into owned data. let initializer_int64 = unsafe { Self::read_initializers_int64(api, graph_ptr).unwrap_or_default() }; + let constant_initializer_names = + unsafe { Self::read_constant_initializer_names(api, graph_ptr).unwrap_or_default() }; let mut absent_outputs: HashSet = HashSet::new(); @@ -261,6 +278,7 @@ impl OutboundGraphReader { ort_index_to_node_id, ort_node_ptrs: ort_nodes, initializer_int64, + constant_initializer_names, absent_outputs, }) } @@ -277,6 +295,17 @@ impl OutboundGraphReader { &self.absent_outputs } + /// Names of the graph's **constant** initializers, i.e. the values ORT + /// keeps fixed for the whole session. + /// + /// Deliberately not "every initializer": an IR>=4 initializer that also + /// appears as a graph input is only a default the caller may replace on + /// any `Run`, and is excluded. Anything keyed on this set may assume the + /// buffer never changes. + pub fn constant_initializer_names(&self) -> &HashSet { + &self.constant_initializer_names + } + /// Access the owned int64 initializer data (copied from ORT during construction). pub fn initializer_int64(&self) -> &HashMap> { &self.initializer_int64 @@ -749,6 +778,72 @@ impl OutboundGraphReader { // ─── Initializer reading ──────────────────────────────────────────── + /// Read the names of every **constant** initializer, of any dtype, without + /// touching the tensor data. + /// + /// `read_initializers_int64` deliberately copies only small int64 tensors, + /// because it exists to resolve constant axes. Constant-weight detection + /// needs the opposite: every initializer, including a 1 GB `MatMulNBits` + /// `B`, and none of the bytes. + /// + /// Non-constant initializers — IR>=4 defaults for a matching graph input, + /// which the caller may override on any `Run` — are excluded. If ORT + /// cannot answer whether an initializer is constant, it is treated as + /// non-constant: the cost of a false negative is a repeated prepack, the + /// cost of a false positive is a wrong answer. + unsafe fn read_constant_initializer_names( + api: *const ort::OrtApi, + graph: *const ort::OrtGraph, + ) -> Result, String> { + let get_num = unsafe { (*api).Graph_GetNumInitializers } + .ok_or("OrtApi.Graph_GetNumInitializers is null")?; + let mut num_init = 0usize; + let status = unsafe { get_num(graph, &mut num_init) }; + Self::check(status)?; + if num_init == 0 { + return Ok(HashSet::new()); + } + + let get_inits = unsafe { (*api).Graph_GetInitializers } + .ok_or("OrtApi.Graph_GetInitializers is null")?; + let mut init_infos: Vec<*const ort::OrtValueInfo> = vec![ptr::null(); num_init]; + let status = unsafe { get_inits(graph, init_infos.as_mut_ptr(), num_init) }; + Self::check(status)?; + + let get_name = + unsafe { (*api).GetValueInfoName }.ok_or("OrtApi.GetValueInfoName is null")?; + let is_constant = unsafe { (*api).ValueInfo_IsConstantInitializer } + .ok_or("OrtApi.ValueInfo_IsConstantInitializer is null")?; + let mut names = HashSet::with_capacity(num_init); + for &info in &init_infos { + if info.is_null() { + continue; + } + let mut constant = false; + let status = unsafe { is_constant(info, &mut constant) }; + if !status.is_null() { + Self::check(status).ok(); + continue; + } + if !constant { + continue; + } + let mut name_ptr: *const std::ffi::c_char = ptr::null(); + let status = unsafe { get_name(info, &mut name_ptr) }; + if !status.is_null() || name_ptr.is_null() { + Self::check(status).ok(); + continue; + } + let name = unsafe { CStr::from_ptr(name_ptr) } + .to_string_lossy() + .into_owned(); + if !name.is_empty() { + names.insert(name); + } + } + Ok(names) + } + /// Read graph initializers and extract small int64 tensors as owned data. /// This copies the data during the callback frame; no ORT pointers are cached. unsafe fn read_initializers_int64(