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 kernel ready to execute a specific op with specific shapes (§4.2).
96pub trait Kernel: Send {
97    /// Execute over device-resident inputs/outputs.
98    fn execute(&self, inputs: &[TensorView], outputs: &mut [TensorMut]) -> Result<()>;
99
100    /// Estimated FLOPs, if known (for the cost model).
101    fn estimated_flops(&self) -> Option<u64> {
102        None
103    }
104
105    /// Whether the kernel accepts a non-contiguous (strided) input at `idx`.
106    fn supports_strided_input(&self, input_idx: usize) -> bool {
107        let _ = input_idx;
108        false
109    }
110
111    /// The layout the kernel writes most efficiently, if it has a preference.
112    fn preferred_output_layout(&self) -> Option<TensorLayout> {
113        None
114    }
115
116    /// Whether this kernel can be captured inside a CUDA graph.
117    fn cuda_graph_compatible(&self) -> bool {
118        false
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    #[test]
127    fn cost_zero_and_total() {
128        assert_eq!(Cost::ZERO.total_us(), 0.0);
129        let c = Cost::new(10.0, 5.0, 2.0);
130        assert_eq!(c.total_us(), 17.0);
131        assert_eq!(c.bytes_moved, 0);
132        assert_eq!(c.launch_us, 0.0);
133    }
134
135    #[test]
136    fn cost_builders_are_additive() {
137        let c = Cost::new(10.0, 5.0, 2.0)
138            .with_launch_us(3.0)
139            .with_bytes_moved(4096);
140        // launch time folds into the total; bytes_moved is metadata for roofline.
141        assert_eq!(c.total_us(), 20.0);
142        assert_eq!(c.bytes_moved, 4096);
143    }
144
145    #[test]
146    fn kernel_match_reports_support() {
147        let supported = KernelMatch::Supported {
148            cost: Cost::new(1.0, 0.0, 0.0),
149            required_input_layouts: None,
150            output_layouts: vec![],
151        };
152        assert!(supported.is_supported());
153        assert!(!KernelMatch::Unsupported.is_supported());
154    }
155}
onnx_runtime_ep_api/kernel.rs

onnx_runtime_ep_api/
kernel.rs
