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 0f321fb22a..b868ad9a0f 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 @@ -5794,6 +5794,173 @@ opset_import: {opset_imports} } } +/// A chain of `depth` identical single-input nodes, `X -> t1 -> … -> Y`. +/// +/// Depth is the axis that separates the two costs #1077 conflates. Whatever +/// ORT charges per `Run` — session bookkeeping, fetch/feed marshalling — is +/// paid once regardless of depth, while our per-node dispatch is paid `depth` +/// times. Comparing depth 1, 10 and 100 against ORT therefore reads the +/// per-node slope directly, instead of a single number in which the fixed and +/// variable parts are indistinguishable. +/// +/// `dynamic` leaves the row count symbolic, so ORT cannot fold shapes at +/// session build and our shape inference runs for real on every `Run`. A +/// dispatch path that is only fast on static shapes is not fast. +fn chain_case( + name: &'static str, + op: &'static str, + depth: usize, + len: usize, + dynamic: bool, +) -> MatmulFamilyCase { + assert!(depth >= 1, "a chain needs at least one node"); + let dims = vec![1, len as i64]; + let decl = |n: &str| -> String { + if dynamic { + format!( + "{{ name: \"{n}\" type {{ tensor_type {{ elem_type: {ELEM_F32} \ + shape {{ dim: [{{ dim_param: \"batch\" }}, {{ dim_value: {len} }}] }} }} }} }}" + ) + } else { + textproto_value_info(n, ELEM_F32, &dims) + } + }; + let nodes: Vec = (0..depth) + .map(|i| { + let input = if i == 0 { + "X".to_owned() + } else { + format!("t{i}") + }; + let output = if i + 1 == depth { + "Y".to_owned() + } else { + format!("t{}", i + 1) + }; + format!("{{ input: [\"{input}\"] output: [\"{output}\"] op_type: \"{op}\" }}") + }) + .collect(); + let model = format!( + r#"ir_version: 10 +graph {{ + node: [{}] + name: "{name}" + input: [{}] + output: [{}] +}} +opset_import: [{{ version: 17 }}] +"#, + nodes.join(", "), + decl("X"), + decl("Y"), + ); + let x = activation_f32(1, len); + MatmulFamilyCase { + name, + op, + model, + inputs: vec![GeneratedInput { + name: "X", + elem: ELEM_F32, + dims, + data: f32_slice_to_bytes(&x), + }], + output: "Y", + output_elem: ELEM_F32, + ort_can_build: true, + // A Relu chain is exact in both implementations; Identity is a copy. + // Anything above zero here would be hiding a real disagreement. + tolerance: 0.0, + } +} + +/// A single small `MatMul`, `[1, k] x [k, n]`, with both operands supplied at +/// run time. +/// +/// The elementwise cases all have trivial kernels, which is what makes them +/// good dispatch probes — but it also means a win there could be an artefact of +/// how little work the kernel does. `MatMul` at these sizes still has a real +/// kernel while staying small enough for dispatch to matter, so it is the case +/// that says whether the fast path generalises beyond one-liners. +/// +/// Both operands are runtime inputs rather than initializers, which keeps the +/// weight out of the TextFormat source (an octal-escaped 1024x1024 f32 blob is +/// megabytes of test fixture) and denies both sides any prepacking, so the +/// comparison stays about dispatch. +fn small_matmul_case(name: &'static str, k: usize, n: usize) -> MatmulFamilyCase { + let model = format!( + r#"ir_version: 10 +graph {{ + node: [{{ input: ["X", "W"] output: ["Y"] op_type: "MatMul" }}] + name: "{name}" + input: [{}, {}] + output: [{}] +}} +opset_import: [{{ version: 17 }}] +"#, + textproto_value_info("X", ELEM_F32, &[1, k as i64]), + textproto_value_info("W", ELEM_F32, &[k as i64, n as i64]), + textproto_value_info("Y", ELEM_F32, &[1, n as i64]), + ); + MatmulFamilyCase { + name, + op: "MatMul", + model, + inputs: vec![ + GeneratedInput { + name: "X", + elem: ELEM_F32, + dims: vec![1, k as i64], + data: f32_slice_to_bytes(&activation_f32(1, k)), + }, + GeneratedInput { + name: "W", + elem: ELEM_F32, + dims: vec![k as i64, n as i64], + data: f32_slice_to_bytes(&activation_f32(k, n)), + }, + ], + output: "Y", + output_elem: ELEM_F32, + ort_can_build: true, + tolerance: 1e-3, + } +} + +/// The grid issue #1077 is closed against. +/// +/// Every row isolates one variable. Read together they say where the remaining +/// dispatch cost is; read individually none of them would. +/// +/// * **`identity_1`** is the floor. An `Identity` node does a memcpy, so +/// essentially everything the timer sees is the cost of getting there and +/// back. If we are at parity anywhere, it has to be here first. +/// * **`relu_1` / `relu_10` / `relu_100`** vary only depth, giving the +/// per-node slope with the per-`Run` constant divided out. +/// * **`*_dyn`** repeat the static rows with a symbolic batch dimension, so a +/// fast path that only works when shapes are known at build time cannot pass +/// unnoticed. +/// * **`matmul_*`** check the result is not an artefact of trivial kernels. +/// * The `4k` width keeps every kernel small enough that dispatch is a visible +/// fraction; at 1 Mi it would be rounding error, which is the whole reason +/// the earlier activation sweep could not answer this question. +fn dispatch_grid_cases() -> Vec { + const W: usize = 4096; + vec![ + chain_case("grid_identity_1_static", "Identity", 1, W, false), + chain_case("grid_identity_10_static", "Identity", 10, W, false), + chain_case("grid_relu_1_static", "Relu", 1, W, false), + chain_case("grid_relu_10_static", "Relu", 10, W, false), + chain_case("grid_relu_100_static", "Relu", 100, W, false), + chain_case("grid_identity_1_dyn", "Identity", 1, W, true), + chain_case("grid_relu_1_dyn", "Relu", 1, W, true), + chain_case("grid_relu_10_dyn", "Relu", 10, W, true), + chain_case("grid_relu_100_dyn", "Relu", 100, W, true), + small_matmul_case("grid_matmul_128x128", 128, 128), + small_matmul_case("grid_matmul_512x512", 512, 512), + ] +} + /// The elementwise grid `CPU_ACTIVATION_GAPS.md` is written against, as /// session-level cases. /// @@ -6180,7 +6347,25 @@ fn plugin_path_ab_vs_plain_ort() { 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().into_iter().chain(unary_bench_cases()) { + // The #1077 grid is off by default: it is 11 more sessions, and the + // matmul/unary rows answer a different question (kernel throughput) than + // these do (dispatch overhead). `NXRT_MM_BENCH_GRID=1` adds them; + // `NXRT_MM_BENCH_GRID=only` runs them alone, which is what a dispatch + // measurement wants — nothing else resident, nothing else warming caches. + let grid = std::env::var("NXRT_MM_BENCH_GRID").unwrap_or_default(); + let cases: Vec = match grid.as_str() { + "only" => dispatch_grid_cases(), + "1" => matmul_bench_cases() + .into_iter() + .chain(unary_bench_cases()) + .chain(dispatch_grid_cases()) + .collect(), + _ => matmul_bench_cases() + .into_iter() + .chain(unary_bench_cases()) + .collect(), + }; + for case in cases { if !filter.is_empty() && !case.name.contains(&filter) { continue; } diff --git a/crates/onnx-runtime-ep-plugin/src/compute.rs b/crates/onnx-runtime-ep-plugin/src/compute.rs index 3906614eec..2621f66aef 100644 --- a/crates/onnx-runtime-ep-plugin/src/compute.rs +++ b/crates/onnx-runtime-ep-plugin/src/compute.rs @@ -1544,6 +1544,24 @@ fn staging_trace_enabled() -> bool { *ENABLED.get_or_init(|| std::env::var("ONNX_GENAI_PLUGIN_TRANSFER_TRACE").as_deref() == Ok("1")) } +/// Emit one staging-trace line, formatting the message only if the trace is on. +/// +/// `staging_log` already returns immediately when the trace is disabled, but +/// the argument is evaluated first: `staging_log!(…)` builds a +/// `String` and formats every interpolated value on the way to a function that +/// throws it away. One of those sits on the per-`Run` dispatch path, so it was +/// paid by every dispatch in production to serve a debugging aid that is off. +/// +/// This checks the gate first, so a disabled trace costs one relaxed load of a +/// `OnceLock` and nothing else. +macro_rules! staging_log { + ($($arg:tt)*) => { + if staging_trace_enabled() { + staging_log(&format!($($arg)*)); + } + }; +} + /// Emit one staging-trace line. Goes to stderr (buffered by pipes, so it can be /// lost when the process hangs) *and*, when `ONNX_GENAI_PLUGIN_TRANSFER_TRACE_FILE` /// is set, is appended to that file with an immediate flush — the file path is @@ -1694,54 +1712,50 @@ unsafe fn stage_host_boundary_inputs( ) -> Result<(), String> { // First pass: which operands are host-resident data operands needing upload? let mut to_stage: Vec = Vec::new(); - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: enter operands={} host_operands={:?}", ort_indices.len(), host_operands - )); + ); for (p, slot) in ort_indices.iter().enumerate() { let Some(ort_idx) = *slot else { - staging_log(&format!( - "[plugin/staging #982] {label}: op[{p}] skip (not ORT-bound)" - )); + staging_log!("[plugin/staging #982] {label}: op[{p}] skip (not ORT-bound)"); continue; }; if host_operands.contains(&p) { - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: op[{p}] ort_idx={ort_idx} skip (host control operand)" - )); + ); continue; } let view = &kernel_inputs[p]; if view.is_absent() || view.data.0.is_null() { - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: op[{p}] ort_idx={ort_idx} skip (absent/null)" - )); + ); continue; } let numel: usize = view.shape.iter().product(); if numel == 0 || view.dtype.byte_size() == 0 { - staging_log(&format!( - "[plugin/staging #982] {label}: op[{p}] ort_idx={ort_idx} skip (empty)" - )); + staging_log!("[plugin/staging #982] {label}: op[{p}] ort_idx={ort_idx} skip (empty)"); continue; } let Some(mem_info) = (unsafe { ort_input_mem_info(api, ctx, ort_idx) }) else { - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: op[{p}] ort_idx={ort_idx} skip (no mem info)" - )); + ); continue; }; let dev_type = unsafe { mem_info_device_type(api, mem_info) }; if unsafe { mem_info_is_device(api, mem_info) } { - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: op[{p}] ort_idx={ort_idx} skip (already device, dev_type={dev_type:?})" - )); + ); continue; // already on device — nothing to do } - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: op[{p}] ort_idx={ort_idx} HOST (dev_type={dev_type:?}) → will stage" - )); + ); to_stage.push(p); } @@ -1772,10 +1786,10 @@ unsafe fn stage_host_boundary_inputs( && let Some(recon) = staging.recon_mem_info.as_ref() { device_mi = recon.ptr; - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: no device OrtValue found; \ falling back to reconstructed EP memory info recon={device_mi:?}" - )); + ); } if device_mi.is_null() { return Err(format!( @@ -1794,19 +1808,17 @@ unsafe fn stage_host_boundary_inputs( let byte_len = numel * view.dtype.byte_size(); let dst = unsafe { alloc_scratch(api, ctx, device_mi, byte_len) } .map_err(|e| format!("{label}: staging scratch alloc failed: {e}"))?; - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: staging op[{p}] byte_len={byte_len} \ mem_dev={mem_dev:?} dst={dst:?} — issuing host→device upload" - )); + ); // SAFETY: the source is a host tensor of `byte_len` contiguous bytes // (ORT-provided, contiguous strides); `dst` is device scratch of the // same size. let src = unsafe { std::slice::from_raw_parts(view.data.0.cast::(), byte_len) }; unsafe { staging.copier.copy_host_to_device(src, dst) } .map_err(|e| format!("{label}: host→device upload failed: {e}"))?; - staging_log(&format!( - "[plugin/staging #982] {label}: op[{p}] host→device upload complete" - )); + staging_log!("[plugin/staging #982] {label}: op[{p}] host→device upload complete"); kernel_inputs[p].data = DevicePtr(dst.cast_const()); } @@ -2184,12 +2196,12 @@ unsafe extern "C" fn compute_execute( return fail_status("Compute: no kernels compiled for this subgraph"); } - staging_log(&format!( + staging_log!( "[plugin/staging #982] compute_execute enter: entries={} routing={} device_staging={}", exported.entries.len(), exported.routing.is_some(), exported.device_staging.is_some() - )); + ); let api = crate::status::host_api(); if api.is_null() { diff --git a/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs b/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs index 5385d85a13..de35f76036 100644 --- a/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs +++ b/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs @@ -72,7 +72,7 @@ pub const CPU_EP_SUPPORTED_DTYPES: &[DataType] = &[ pub(crate) fn validate_dims( dims: &[i64], dtype: DataType, - context: &str, + context: impl std::fmt::Display, ) -> Result<(Vec, usize, usize), String> { let mut shape: Vec = Vec::with_capacity(dims.len()); for (dim_idx, &d) in dims.iter().enumerate() { @@ -252,9 +252,10 @@ pub(crate) unsafe fn read_inputs( unsafe { release_type_shape(type_shape) }; // Validate ORT-supplied dims: reject negatives and detect overflow. - // Three allocations: the eager `format!` label, `shape`, and `strides`. - crate::dispatch_probe::count_n(crate::dispatch_probe::Event::DispatchAlloc, 3); - let (shape, _, _) = validate_dims(&dims, dtype, &format!("input {i}"))?; + // Two allocations: `shape` and `strides`. The label is `format_args!`, + // which borrows its arguments and formats only if an error is raised. + crate::dispatch_probe::count_n(crate::dispatch_probe::Event::DispatchAlloc, 2); + let (shape, _, _) = validate_dims(&dims, dtype, format_args!("input {i}"))?; let strides = onnx_runtime_ir::compute_contiguous_strides(&shape); // Data pointer. @@ -711,20 +712,21 @@ mod dispatch_cost { ); } - /// One input also costs five heap allocations: the `Vec` once - /// per `Run`, then per input the `dims` scratch, the eager - /// `format!("input {i}")` label built even when validation succeeds, - /// `shape`, and `strides`. The label is pure waste on the success path and - /// is the first thing to remove. + /// One input costs four heap allocations: the `Vec` once per + /// `Run`, then per input the `dims` scratch, `shape`, and `strides`. + /// + /// It was five until the error label stopped being built eagerly. That is + /// the counter earning its keep: the improvement is visible as a number + /// changing in a test, not as a claim in a commit message. #[test] - fn reading_one_input_costs_exactly_five_allocations() { + fn reading_one_input_costs_exactly_four_allocations() { let api = fake_api(); let before = dispatch_probe::snapshot(); let _ = unsafe { read_inputs(&api, std::ptr::null_mut()) }.expect("fake api is total"); let d = dispatch_probe::snapshot().since(&before); assert_eq!( d.event(Event::DispatchAlloc), - 1 + 4, + 1 + 3, "allocations on the per-Run input path changed" ); } @@ -751,7 +753,7 @@ mod dispatch_cost { assert_eq!(inputs.len(), 3); assert_eq!(d.event(Event::OrtFfiCall), 1 + 3 * 7); - assert_eq!(d.event(Event::DispatchAlloc), 1 + 3 * 4); + assert_eq!(d.event(Event::DispatchAlloc), 1 + 3 * 3); } /// An absent optional input short-circuits: ORT hands back a null value and diff --git a/crates/onnx-runtime-ep-plugin/src/transfer.rs b/crates/onnx-runtime-ep-plugin/src/transfer.rs index e0387e3303..949958cfdd 100644 --- a/crates/onnx-runtime-ep-plugin/src/transfer.rs +++ b/crates/onnx-runtime-ep-plugin/src/transfer.rs @@ -770,7 +770,7 @@ unsafe extern "C" fn transfer_full_copy_tensors( let (_, _, byte_len) = match crate::kernel_ctx::validate_dims( &dims, dtype, - &format!("CopyTensors[{i}]"), + format_args!("CopyTensors[{i}]"), ) { Ok(v) => v, Err(e) => return fail_status(&e),