tensor_wasm_exec/
auto_offload.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2026 Craton Software Company
3//! Pre-instantiation auto-offload analysis.
4//!
5//! Walks Wasm bytes via [`wasmparser`], extracts per-function operator
6//! sequences shaped like basic blocks, and runs each through
7//! [`tensor_wasm_jit::detector::classify`]. Verdicts are emitted as `tracing`
8//! events so operators can see which kernels *would* be GPU-offloaded
9//! once Wasmtime exposes the per-block compile hook documented in
10//! `docs/WASMTIME-FORK.md`.
11//!
12//! This is consultation-only: the analysis runs, but Wasmtime's
13//! Cranelift output is NOT replaced. Activating the swap requires
14//! the rewrite pipeline in `tensor_wasm_jit::rewrite`.
15
16use tensor_wasm_jit::detector::{classify, BlockIR, DetectorConfig, DetectorVerdict, Op};
17use tensor_wasm_jit::ir::ElemType;
18use tracing::{debug, info, instrument};
19use wasmparser::{Operator, Parser, Payload};
20
21/// A single classification result for one function body.
22#[derive(Debug, Clone, PartialEq)]
23pub struct OffloadVerdict {
24    /// Function index inside the Wasm module.
25    pub function_index: u32,
26    /// Detector classification.
27    pub verdict: DetectorVerdict,
28    /// Number of operators inspected.
29    pub op_count: usize,
30    /// Fraction of v128 ops (0.0-1.0).
31    pub v128_ratio: f32,
32}
33
34/// Errors raised by the analyser.
35#[derive(Debug, thiserror::Error)]
36pub enum AnalyseError {
37    /// The Wasm bytes failed to parse.
38    #[error("wasmparser: {0}")]
39    Parse(String),
40}
41
42/// Map a wasmparser [`Operator`] to the simplified [`tensor_wasm_jit::detector::Op`].
43///
44/// Most ops fall outside the detector's taxonomy — they're collapsed into
45/// [`Op::Load`] / [`Op::Store`] (memory) or [`Op::Branch`] (control flow) so
46/// the ratio computation is meaningful but never panics. Anything else
47/// becomes [`Op::Other`] so the v128 ratio is computed honestly without
48/// inflating arithmetic density from non-arithmetic ops like `local.get`.
49fn op_to_detector_op(op: &Operator<'_>) -> Op {
50    use wasmparser::Operator::*;
51    match op {
52        V128Load { .. } => Op::Load,
53        V128Store { .. } => Op::Store,
54        F32Add | I32Add | I64Add | F64Add => Op::ScalarAdd,
55        F32Mul | I32Mul | I64Mul | F64Mul => Op::ScalarMul,
56        I32Load { .. } | I64Load { .. } | F32Load { .. } | F64Load { .. } => Op::Load,
57        I32Store { .. } | I64Store { .. } | F32Store { .. } | F64Store { .. } => Op::Store,
58        Br { .. }
59        | BrIf { .. }
60        | BrTable { .. }
61        | If { .. }
62        | Else
63        | Loop { .. }
64        | Block { .. } => Op::Branch,
65        Call { .. } | CallIndirect { .. } | ReturnCall { .. } => Op::Call,
66        F32x4Add => Op::V128Add {
67            lane_ty: ElemType::F32,
68            lanes: 4,
69        },
70        I32x4Add => Op::V128Add {
71            lane_ty: ElemType::I32,
72            lanes: 4,
73        },
74        F32x4Mul => Op::V128Mul {
75            lane_ty: ElemType::F32,
76            lanes: 4,
77        },
78        I32x4Mul => Op::V128Mul {
79            lane_ty: ElemType::I32,
80            lanes: 4,
81        },
82        // Fail-closed: widths the emitter cannot yet lower correctly stay on the CPU path.
83        F64x2Add | I64x2Add | I16x8Add | I8x16Add | F64x2Mul | I64x2Mul | I16x8Mul => Op::Other,
84        _ => Op::Other,
85    }
86}
87
88/// Analyse Wasm bytes, emit per-function `tracing` events, and return the
89/// list of verdicts, using the default detector configuration
90/// ([`DetectorConfig::default`]).
91///
92/// Thin wrapper over [`analyse_with_config`] preserving the historical
93/// signature for the consultation-only call sites.
94#[instrument(skip(wasm), fields(wasm_bytes = wasm.len()))]
95pub fn analyse(wasm: &[u8]) -> Result<Vec<OffloadVerdict>, AnalyseError> {
96    analyse_with_config(wasm, &DetectorConfig::default())
97}
98
99/// Analyse Wasm bytes against an explicit detector configuration, emit
100/// per-function `tracing` events, and return the list of verdicts.
101///
102/// The default-config wrapper [`analyse`] keeps the historical
103/// consultation-only behaviour; this entry point lets the executor's
104/// auto-offload activation path consult with the *same* detector thresholds
105/// the rewrite will use, so the consultation verdict and the rewrite agree
106/// on which functions are offload candidates.
107#[instrument(skip(wasm, cfg), fields(wasm_bytes = wasm.len()))]
108pub fn analyse_with_config(
109    wasm: &[u8],
110    cfg: &DetectorConfig,
111) -> Result<Vec<OffloadVerdict>, AnalyseError> {
112    let mut verdicts = Vec::new();
113    let mut func_idx: u32 = 0;
114    for payload in Parser::new(0).parse_all(wasm) {
115        let payload = payload.map_err(|e| AnalyseError::Parse(format!("{e}")))?;
116        if let Payload::CodeSectionEntry(body) = payload {
117            let mut ops = body
118                .get_operators_reader()
119                .map_err(|e| AnalyseError::Parse(format!("{e}")))?;
120            let mut detector_ops = Vec::new();
121            let mut saw_loop = false;
122            while !ops.eof() {
123                match ops.read() {
124                    Ok(op) => {
125                        if matches!(op, wasmparser::Operator::Loop { .. }) {
126                            saw_loop = true;
127                        }
128                        detector_ops.push(op_to_detector_op(&op));
129                    }
130                    Err(_) => break,
131                }
132            }
133            // Only attach a static-loop guess when the parser actually saw a
134            // `Loop` — straight-line functions inheriting `Some(128)` made
135            // the detector approve cold code that would never amortise the
136            // offload setup cost.
137            let trip_count = if saw_loop { Some(128) } else { None };
138            let block = BlockIR::new(format!("func{func_idx}"), detector_ops.clone(), trip_count);
139            let v = classify(&block, cfg);
140            let op_count = detector_ops.len();
141            let v128 = detector_ops.iter().filter(|o| o.is_v128()).count();
142            let v128_ratio = if op_count == 0 {
143                0.0
144            } else {
145                v128 as f32 / op_count as f32
146            };
147            match v {
148                DetectorVerdict::Offload => info!(
149                    target: "tensor_wasm_exec::auto_offload",
150                    function = func_idx,
151                    op_count,
152                    v128_ratio,
153                    "verdict: Offload"
154                ),
155                DetectorVerdict::KeepOnCpu => debug!(
156                    target: "tensor_wasm_exec::auto_offload",
157                    function = func_idx,
158                    op_count,
159                    v128_ratio,
160                    "verdict: KeepOnCpu"
161                ),
162            }
163            verdicts.push(OffloadVerdict {
164                function_index: func_idx,
165                verdict: v,
166                op_count,
167                v128_ratio,
168            });
169            func_idx += 1;
170        }
171    }
172    Ok(verdicts)
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178
179    const TRIVIAL_WAT: &str = r#"(module (func (export "noop")))"#;
180
181    #[test]
182    fn analyse_noop_module() {
183        let wasm = wat::parse_str(TRIVIAL_WAT).unwrap();
184        let v = analyse(&wasm).expect("analyse");
185        assert_eq!(v.len(), 1);
186        assert_eq!(v[0].function_index, 0);
187        // A noop function has no v128 ops worth offloading.
188        assert_eq!(v[0].verdict, DetectorVerdict::KeepOnCpu);
189    }
190
191    #[test]
192    fn analyse_v128_heavy_module() {
193        // A function full of v128.add — should be flagged.
194        let wat = r#"
195            (module
196              (memory 1)
197              (func (export "hot")
198                (drop (v128.const i32x4 1 2 3 4))
199                (drop (v128.const i32x4 5 6 7 8))
200                (drop (i32x4.add (v128.const i32x4 1 2 3 4) (v128.const i32x4 5 6 7 8)))
201                (drop (i32x4.add (v128.const i32x4 1 2 3 4) (v128.const i32x4 5 6 7 8)))
202                (drop (i32x4.add (v128.const i32x4 1 2 3 4) (v128.const i32x4 5 6 7 8)))
203                (drop (i32x4.add (v128.const i32x4 1 2 3 4) (v128.const i32x4 5 6 7 8)))
204              )
205            )
206        "#;
207        let wasm = wat::parse_str(wat).unwrap();
208        let v = analyse(&wasm).expect("analyse");
209        assert_eq!(v.len(), 1);
210        // We don't assert on the threshold outcome — the test exercises the
211        // analyser, not the threshold tuning. Just confirm we walked ops.
212        assert!(matches!(
213            v[0].verdict,
214            DetectorVerdict::Offload | DetectorVerdict::KeepOnCpu
215        ));
216        assert!(v[0].op_count > 0);
217    }
218
219    #[test]
220    fn analyse_invalid_wasm_returns_error() {
221        let bytes = [0x00, 0x61, 0x73, 0x6d, 0xff, 0xff, 0xff, 0xff];
222        let err = analyse(&bytes).unwrap_err();
223        assert!(matches!(err, AnalyseError::Parse(_)));
224    }
225
226    #[test]
227    fn straight_line_function_keeps_on_cpu_even_with_v128() {
228        // Without a `loop`, the function should NOT be offloaded —
229        // offload overhead doesn't amortise on a single straight-line call.
230        let wat = r#"
231            (module
232              (memory 1)
233              (func (export "straight") (result v128)
234                (i32x4.add (v128.const i32x4 1 2 3 4) (v128.const i32x4 5 6 7 8))
235              )
236            )
237        "#;
238        let wasm = wat::parse_str(wat).unwrap();
239        let v = analyse(&wasm).expect("analyse");
240        assert_eq!(v.len(), 1);
241        assert_eq!(
242            v[0].verdict,
243            DetectorVerdict::KeepOnCpu,
244            "straight-line v128 must not trigger offload"
245        );
246    }
247}
tensor_wasm_exec/auto_offload.rs

tensor_wasm_exec/
auto_offload.rs
