tensor_wasm_jit/
clif_lower.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2026 Craton Software Company
3//! Lower a [`crate::detector::BlockIR`] into a
4//! [`crate::ir::TensorWasmKernelBlueprint`].
5
6use thiserror::Error;
7
8use crate::detector::{BlockIR, Op};
9use crate::ir::{ElemType, GridHint, TensorWasmKernelBlueprint, TensorWasmOp};
10
11/// Errors produced by the lowering pass.
12#[derive(Debug, Error, PartialEq)]
13pub enum LowerError {
14    /// The block contained operations the lowering pass doesn't yet handle.
15    #[error("unsupported op in block `{block}`: {op:?}")]
16    UnsupportedOp {
17        /// Block name from `BlockIR::name`.
18        block: String,
19        /// The op the lowering didn't know how to handle.
20        op: Op,
21    },
22    /// The lowering refused because of a memory-safety concern (e.g. an
23    /// out-of-bounds load pattern).
24    #[error("memory-reference out of bounds in block `{block}`")]
25    OutOfBoundsMemory {
26        /// Block name.
27        block: String,
28    },
29    /// jit CRITICAL fix: the block mixed SIMD element types (e.g. an
30    /// `f32x4.add` next to an `i32x4.add`). The flat blueprint op stream
31    /// has a single load/store element width inferred from the block's SIMD
32    /// ops, so a heterogeneous block cannot be lowered without miscompiling
33    /// the loads/stores. We fail closed; the caller keeps the function on
34    /// the CPU path.
35    #[error("mixed SIMD element types in block `{block}` ({first} vs {second})")]
36    MixedElementTypes {
37        /// Block name.
38        block: String,
39        /// First element type seen.
40        first: ElemType,
41        /// Conflicting element type seen later.
42        second: ElemType,
43    },
44}
45
46/// Default lane width used when a block has no SIMD op to infer width from
47/// (e.g. a pure load/store block); kept for the f32x4 production shape.
48const DEFAULT_LANES: u32 = 4;
49
50/// A memory reference seen in the source block: an absolute byte offset
51/// into the instance's linear memory and an access width in bytes. The
52/// lowering pass uses these to refuse offload candidates whose accesses
53/// extend past the instance's allocated `UnifiedBuffer` extent.
54#[derive(Debug, Clone, Copy, PartialEq, Eq)]
55pub struct MemRef {
56    /// Byte offset into the instance linear memory.
57    pub offset: u64,
58    /// Number of bytes accessed.
59    pub len: u64,
60}
61
62/// Infer the single SIMD element type + lane count used by a block.
63///
64/// jit CRITICAL fix: the blueprint op stream is flat and the load/store
65/// element width must match the block's SIMD compute width. We scan the
66/// SIMD ops and require they all agree on element type and lane count; a
67/// mismatch is failed closed via [`LowerError::MixedElementTypes`]. Blocks
68/// with no SIMD op (pure load/store) fall back to the f32x4 default — the
69/// only shape the production auto-offload path emits for such a block.
70fn infer_block_shape(block: &BlockIR) -> Result<(ElemType, u32), LowerError> {
71    let mut shape: Option<(ElemType, u32)> = None;
72    for op in &block.ops {
73        let cur = match op {
74            Op::V128Add { lane_ty, lanes }
75            | Op::V128Mul { lane_ty, lanes }
76            | Op::V128Fma { lane_ty, lanes } => (*lane_ty, *lanes),
77            _ => continue,
78        };
79        match shape {
80            None => shape = Some(cur),
81            Some((seen_ty, _)) if seen_ty != cur.0 => {
82                return Err(LowerError::MixedElementTypes {
83                    block: block.name.clone(),
84                    first: seen_ty,
85                    second: cur.0,
86                });
87            }
88            Some(_) => {}
89        }
90    }
91    Ok(shape.unwrap_or((ElemType::F32, DEFAULT_LANES)))
92}
93
94/// Lower a [`BlockIR`] into a [`TensorWasmKernelBlueprint`].
95///
96/// The mapping is intentionally narrow. Each SIMD op carries its element
97/// type and lane count through to the blueprint; load/store ops adopt the
98/// block's inferred SIMD element type (see [`infer_block_shape`]):
99///
100/// | Source                                       | Target                                  |
101/// |----------------------------------------------|-----------------------------------------|
102/// | `Op::V128Add { lane_ty, lanes }`             | `TensorWasmOp::VecAdd { elem, lanes }`  |
103/// | `Op::V128Mul { lane_ty, lanes }`             | `TensorWasmOp::VecMul { elem, lanes }`  |
104/// | `Op::V128Fma { lane_ty, lanes }`             | `TensorWasmOp::VecFma { elem, lanes }`  |
105/// | `Op::Load`                                   | `TensorWasmOp::LoadUnified { elem, .. }`|
106/// | `Op::Store`                                  | `TensorWasmOp::StoreUnified { elem, ..}`|
107/// | Anything else                                | `UnsupportedOp` error                   |
108pub fn lower_block(block: &BlockIR) -> Result<TensorWasmKernelBlueprint, LowerError> {
109    let (mem_elem, mem_lanes) = infer_block_shape(block)?;
110    let mut bp = TensorWasmKernelBlueprint::new(&block.name);
111    for op in &block.ops {
112        let lo = match op {
113            Op::V128Add { lane_ty, lanes } => TensorWasmOp::VecAdd {
114                elem: *lane_ty,
115                lanes: *lanes,
116            },
117            Op::V128Mul { lane_ty, lanes } => TensorWasmOp::VecMul {
118                elem: *lane_ty,
119                lanes: *lanes,
120            },
121            Op::V128Fma { lane_ty, lanes } => TensorWasmOp::VecFma {
122                elem: *lane_ty,
123                lanes: *lanes,
124            },
125            Op::Load => TensorWasmOp::LoadUnified {
126                elem: mem_elem,
127                lanes: mem_lanes,
128            },
129            Op::Store => TensorWasmOp::StoreUnified {
130                elem: mem_elem,
131                lanes: mem_lanes,
132            },
133            unsupported => {
134                return Err(LowerError::UnsupportedOp {
135                    block: block.name.clone(),
136                    op: *unsupported,
137                });
138            }
139        };
140        bp.ops.push(lo);
141    }
142
143    // Trip count clamped into u32; saturate rather than truncate so a
144    // (genuinely huge) static loop doesn't masquerade as a tiny grid.
145    let total_threads = u32::try_from(block.loop_trip_count.unwrap_or(256)).unwrap_or(u32::MAX);
146    bp.grid_hint = GridHint {
147        total_threads,
148        preferred_block_size: 128,
149    };
150    Ok(bp)
151}
152
153/// Lower a [`BlockIR`] with explicit memory-reference validation.
154///
155/// `mem_refs` is the list of memory accesses extracted from the source
156/// block by the caller's side-band analysis (e.g. `wasmparser`). Each
157/// `(offset, len)` pair must satisfy `offset + len <= memory_bytes`;
158/// any violation aborts with [`LowerError::OutOfBoundsMemory`] so the
159/// caller falls back to CPU execution.
160pub fn lower_block_checked(
161    block: &BlockIR,
162    mem_refs: &[MemRef],
163    memory_bytes: u64,
164) -> Result<TensorWasmKernelBlueprint, LowerError> {
165    for r in mem_refs {
166        let end = r.offset.checked_add(r.len);
167        let out_of_bounds = match end {
168            Some(end) => end > memory_bytes,
169            None => true, // overflow → treat as OOB
170        };
171        if out_of_bounds {
172            return Err(LowerError::OutOfBoundsMemory {
173                block: block.name.clone(),
174            });
175        }
176    }
177    lower_block(block)
178}
179
180#[cfg(test)]
181mod tests {
182    use super::*;
183
184    fn block(name: &str, ops: Vec<Op>, loop_n: Option<u64>) -> BlockIR {
185        BlockIR::new(name, ops, loop_n)
186    }
187
188    fn add(lane_ty: ElemType, lanes: u32) -> Op {
189        Op::V128Add { lane_ty, lanes }
190    }
191    fn mul(lane_ty: ElemType, lanes: u32) -> Op {
192        Op::V128Mul { lane_ty, lanes }
193    }
194    fn fma(lane_ty: ElemType, lanes: u32) -> Op {
195        Op::V128Fma { lane_ty, lanes }
196    }
197
198    #[test]
199    fn lower_simple_vector_add() {
200        let b = block(
201            "vec_add",
202            vec![Op::Load, Op::Load, add(ElemType::F32, 4), Op::Store],
203            Some(128),
204        );
205        let bp = lower_block(&b).unwrap();
206        assert_eq!(bp.entry, "vec_add");
207        assert_eq!(bp.ops.len(), 4);
208        assert!(matches!(
209            bp.ops[2],
210            TensorWasmOp::VecAdd {
211                elem: ElemType::F32,
212                lanes: 4
213            }
214        ));
215        // Loads/stores adopt the block's SIMD element type.
216        assert!(matches!(
217            bp.ops[0],
218            TensorWasmOp::LoadUnified {
219                elem: ElemType::F32,
220                lanes: 4
221            }
222        ));
223        assert_eq!(bp.grid_hint.total_threads, 128);
224    }
225
226    /// jit CRITICAL: an `i32x4.add` block must lower to an i32 blueprint —
227    /// element type and lane count survive the lowering, and the loads/
228    /// stores adopt the integer width.
229    #[test]
230    fn lower_i32x4_add_carries_integer_element_type() {
231        let b = block(
232            "i32_add",
233            vec![Op::Load, Op::Load, add(ElemType::I32, 4), Op::Store],
234            Some(128),
235        );
236        let bp = lower_block(&b).unwrap();
237        assert!(matches!(
238            bp.ops[2],
239            TensorWasmOp::VecAdd {
240                elem: ElemType::I32,
241                lanes: 4
242            }
243        ));
244        assert!(matches!(
245            bp.ops[0],
246            TensorWasmOp::LoadUnified {
247                elem: ElemType::I32,
248                ..
249            }
250        ));
251        // It must NOT have silently become an f32 kernel.
252        assert!(!bp.ops.iter().any(|o| matches!(
253            o,
254            TensorWasmOp::VecAdd {
255                elem: ElemType::F32,
256                ..
257            }
258        )));
259    }
260
261    /// jit CRITICAL: a block mixing element types is failed closed rather
262    /// than lowered with a single (wrong) load/store width.
263    #[test]
264    fn lower_mixed_element_types_rejected() {
265        let b = block(
266            "mixed",
267            vec![add(ElemType::F32, 4), add(ElemType::I32, 4)],
268            Some(128),
269        );
270        let err = lower_block(&b).unwrap_err();
271        assert!(matches!(err, LowerError::MixedElementTypes { .. }));
272    }
273
274    /// Non-default lane counts (f64x2) survive lowering.
275    #[test]
276    fn lower_f64x2_lane_count_survives() {
277        let b = block("f64_add", vec![add(ElemType::F64, 2)], Some(128));
278        let bp = lower_block(&b).unwrap();
279        assert!(matches!(
280            bp.ops[0],
281            TensorWasmOp::VecAdd {
282                elem: ElemType::F64,
283                lanes: 2
284            }
285        ));
286    }
287
288    #[test]
289    fn unsupported_op_rejected() {
290        let b = block("with_call", vec![add(ElemType::F32, 4), Op::Call], Some(64));
291        let err = lower_block(&b).unwrap_err();
292        assert!(matches!(err, LowerError::UnsupportedOp { .. }));
293    }
294
295    #[test]
296    fn missing_trip_count_uses_default() {
297        let b = block("dyn", vec![add(ElemType::F32, 4)], None);
298        let bp = lower_block(&b).unwrap();
299        assert_eq!(bp.grid_hint.total_threads, 256);
300    }
301
302    #[test]
303    fn all_v128_lower() {
304        let b = block(
305            "fma_loop",
306            vec![
307                fma(ElemType::F32, 4),
308                fma(ElemType::F32, 4),
309                mul(ElemType::F32, 4),
310            ],
311            Some(64),
312        );
313        let bp = lower_block(&b).unwrap();
314        for op in &bp.ops {
315            assert!(matches!(
316                op,
317                TensorWasmOp::VecFma { .. }
318                    | TensorWasmOp::VecMul { .. }
319                    | TensorWasmOp::VecAdd { .. }
320            ));
321        }
322    }
323
324    #[test]
325    fn lower_block_checked_passes_with_in_bounds_refs() {
326        let b = block(
327            "vec_add",
328            vec![Op::Load, add(ElemType::F32, 4), Op::Store],
329            Some(128),
330        );
331        let refs = &[
332            MemRef { offset: 0, len: 16 },
333            MemRef {
334                offset: 1024,
335                len: 16,
336            },
337        ];
338        let bp = lower_block_checked(&b, refs, 64 * 1024).expect("in-bounds ok");
339        assert_eq!(bp.ops.len(), 3);
340    }
341
342    #[test]
343    fn lower_block_checked_rejects_out_of_bounds_refs() {
344        let b = block("oob", vec![Op::Load], Some(64));
345        // memory_bytes = 1024 but the ref reads past it.
346        let refs = &[MemRef {
347            offset: 1020,
348            len: 16,
349        }];
350        let err = lower_block_checked(&b, refs, 1024).unwrap_err();
351        assert!(matches!(err, LowerError::OutOfBoundsMemory { .. }));
352    }
353
354    #[test]
355    fn lower_block_checked_rejects_overflow_refs() {
356        let b = block("overflow", vec![Op::Load], Some(64));
357        let refs = &[MemRef {
358            offset: u64::MAX - 4,
359            len: 16,
360        }];
361        let err = lower_block_checked(&b, refs, 1024).unwrap_err();
362        assert!(matches!(err, LowerError::OutOfBoundsMemory { .. }));
363    }
364}
tensor_wasm_jit/clif_lower.rs

tensor_wasm_jit/
clif_lower.rs
