onnx_runtime_ep_api/kernel.rs
1//! The [`Kernel`] trait and kernel-match / cost types (§4.2).
2
3use onnx_runtime_ir::TensorLayout;
4
5use crate::error::Result;
6use crate::tensor::{TensorMut, TensorView};
7
8/// A cost estimate for running a kernel, consumed by the placement cost model
9/// (`docs/ORT2.md` §6). All time fields are in **microseconds**; a fuller model
10/// (roofline, calibration) lands in `onnx-runtime-cost-model` (Phase 2).
11///
12/// The struct is `#[non_exhaustive]`: the Phase-2 cost model may add fields
13/// (e.g. energy, occupancy) without breaking EP crates. Construct it via
14/// [`Cost::ZERO`], [`Cost::new`], or the `with_*` builders rather than a struct
15/// literal so those additions stay source-compatible.
16///
17/// The three time components (`compute_us`, `memory_us`, `transfer_us`) map onto
18/// the design's *compute*, *memory-traffic*, and *layout/transfer* estimates;
19/// `launch_us` captures fixed dispatch latency (§6.2 `launch_overhead`) and
20/// `bytes_moved` carries the raw memory-traffic figure a roofline model needs
21/// (§6.3, mirroring the design `Cost::memory_bytes`).
22#[derive(Clone, Copy, Debug, Default, PartialEq)]
23#[non_exhaustive]
24pub struct Cost {
25 /// Estimated compute time (µs).
26 pub compute_us: f64,
27 /// Estimated memory-traffic time (µs).
28 pub memory_us: f64,
29 /// Estimated layout-conversion / cross-device copy time at boundaries (µs).
30 pub transfer_us: f64,
31 /// Fixed kernel-launch / dispatch latency (µs), independent of size.
32 pub launch_us: f64,
33 /// Estimated bytes of memory traffic (for roofline / bandwidth models).
34 pub bytes_moved: u64,
35}
36
37impl Cost {
38 /// The zero cost (free op).
39 pub const ZERO: Cost = Cost {
40 compute_us: 0.0,
41 memory_us: 0.0,
42 transfer_us: 0.0,
43 launch_us: 0.0,
44 bytes_moved: 0,
45 };
46
47 /// A cost from its three time components; `launch_us` and `bytes_moved`
48 /// default to zero (set them via the builders).
49 pub fn new(compute_us: f64, memory_us: f64, transfer_us: f64) -> Self {
50 Self {
51 compute_us,
52 memory_us,
53 transfer_us,
54 ..Self::ZERO
55 }
56 }
57
58 /// Set the fixed launch/dispatch latency.
59 pub fn with_launch_us(mut self, launch_us: f64) -> Self {
60 self.launch_us = launch_us;
61 self
62 }
63
64 /// Set the estimated memory-traffic volume.
65 pub fn with_bytes_moved(mut self, bytes_moved: u64) -> Self {
66 self.bytes_moved = bytes_moved;
67 self
68 }
69
70 /// Total estimated wall time (µs): the sum of all time components.
71 pub fn total_us(&self) -> f64 {
72 self.compute_us + self.memory_us + self.transfer_us + self.launch_us
73 }
74}
75
76/// Result of [`crate::ExecutionProvider::supports_op`].
77pub enum KernelMatch {
78 Supported {
79 cost: Cost,
80 /// Layouts the kernel requires for each input, if constrained.
81 required_input_layouts: Option<Vec<TensorLayout>>,
82 /// Layouts the kernel produces for each output.
83 output_layouts: Vec<TensorLayout>,
84 },
85 Unsupported,
86}
87
88impl KernelMatch {
89 /// Whether the op is supported.
90 pub fn is_supported(&self) -> bool {
91 matches!(self, KernelMatch::Supported { .. })
92 }
93}
94
95/// A zero-copy **view output**: a kernel's declaration that one of its outputs
96/// is a strided view aliasing one of its inputs' buffers, rather than freshly
97/// computed bytes (`docs/ORT2.md` §5.4, lazy PyTorch-style views).
98///
99/// The `shape` / `strides` / `byte_offset` describe the output tensor relative
100/// to the **same base pointer** as the referenced input view (i.e. relative to
101/// the input's backing allocation, honoring any offset that input itself
102/// already carried). Strides are in **elements** and may be negative (DLPack).
103/// The executor records this metadata against the output value and does **not**
104/// allocate a buffer or invoke the compute path for that slot; the source
105/// buffer is kept alive until the view's consumers have run.
106#[derive(Clone, Debug, PartialEq, Eq)]
107pub struct ViewOutput {
108 /// Positional index (into the kernel's `inputs`) of the aliased input.
109 pub input_index: usize,
110 /// Output shape.
111 pub shape: Vec<usize>,
112 /// Output element strides relative to the aliased input's base pointer.
113 pub strides: Vec<i64>,
114 /// Byte offset of the output element origin from the aliased input's base.
115 pub byte_offset: usize,
116}
117
118/// A kernel ready to execute a specific op with specific shapes (§4.2).
119pub trait Kernel: Send {
120 /// Execute over device-resident inputs/outputs.
121 fn execute(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()>;
122
123 /// Estimated FLOPs, if known (for the cost model).
124 fn estimated_flops(&self) -> Option<u64> {
125 None
126 }
127
128 /// Attempt to express this node's outputs as zero-copy [`ViewOutput`]s over
129 /// its inputs instead of computing bytes (the layout/movement-op fast path).
130 ///
131 /// `inputs` carries the real (possibly already-strided) input views and
132 /// `num_outputs` is the node's output arity. Returning:
133 /// * `None` — the default — means "compute normally": the executor allocates
134 /// output buffers and calls [`Kernel::execute`].
135 /// * `Some(specs)` means every output is a view; `specs.len()` MUST equal
136 /// `num_outputs`. A kernel that can view some but not all outputs must
137 /// return `None` (all-or-nothing) so correctness never regresses.
138 ///
139 /// When `Some` is returned, [`Kernel::execute`] is **not** invoked.
140 fn view_outputs(&self, inputs: &[TensorView], num_outputs: usize) -> Option<Vec<ViewOutput>> {
141 let _ = (inputs, num_outputs);
142 None
143 }
144
145 /// Whether the kernel accepts a non-contiguous (strided) input at `idx`.
146 fn supports_strided_input(&self, input_idx: usize) -> bool {
147 let _ = input_idx;
148 false
149 }
150
151 /// The layout the kernel writes most efficiently, if it has a preference.
152 fn preferred_output_layout(&self) -> Option<TensorLayout> {
153 None
154 }
155
156 /// Whether this kernel can be captured inside a CUDA graph.
157 fn cuda_graph_compatible(&self) -> bool {
158 false
159 }
160}
161
162#[cfg(test)]
163mod tests {
164 use super::*;
165
166 #[test]
167 fn cost_zero_and_total() {
168 assert_eq!(Cost::ZERO.total_us(), 0.0);
169 let c = Cost::new(10.0, 5.0, 2.0);
170 assert_eq!(c.total_us(), 17.0);
171 assert_eq!(c.bytes_moved, 0);
172 assert_eq!(c.launch_us, 0.0);
173 }
174
175 #[test]
176 fn cost_builders_are_additive() {
177 let c = Cost::new(10.0, 5.0, 2.0)
178 .with_launch_us(3.0)
179 .with_bytes_moved(4096);
180 // launch time folds into the total; bytes_moved is metadata for roofline.
181 assert_eq!(c.total_us(), 20.0);
182 assert_eq!(c.bytes_moved, 4096);
183 }
184
185 #[test]
186 fn kernel_match_reports_support() {
187 let supported = KernelMatch::Supported {
188 cost: Cost::new(1.0, 0.0, 0.0),
189 required_input_layouts: None,
190 output_layouts: vec![],
191 };
192 assert!(supported.is_supported());
193 assert!(!KernelMatch::Unsupported.is_supported());
194 }
195}