1use std::collections::HashMap;
14use std::mem;
15
16use cranelift_codegen::ir::{self, AbiParam, InstBuilder, types};
17use cranelift_codegen::settings::{self, Configurable};
18use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
19use cranelift_jit::{JITBuilder, JITModule};
20use cranelift_module::{Linkage, Module};
21
22use crate::ast::PolydatNode;
23use crate::ast::SlotShape;
24
25use super::kernels::{JitCore, JitKernelPushPull, JitKernelRaw};
26
27extern "C" fn jit_xxh3_hash(value: u64) -> u64 {
31 guarded(|| xxhash_rust::xxh3::xxh3_64(&value.to_le_bytes()))
32}
33
34extern "C" fn jit_interleave(a: u64, b: u64) -> u64 {
36 guarded(|| {
37 let mut result: u64 = 0;
38 for i in 0..32 {
39 result |= ((a >> i) & 1) << (2 * i);
40 result |= ((b >> i) & 1) << (2 * i + 1);
41 }
42 result
43 })
44}
45
46extern "C" fn jit_lut_sample(input_bits: u64, lut_ptr: u64, lut_len: u64) -> u64 {
51 guarded(|| {
52 let u = f64::from_bits(input_bits).clamp(0.0, 1.0);
53 let n = (lut_len - 1) as f64;
54 let pos = u * n;
55 let idx = (pos as usize).min(lut_len as usize - 2);
56 let frac = pos - idx as f64;
57 let result = unsafe {
58 let ptr = lut_ptr as *const f64;
59 let a = *ptr.add(idx);
60 let b = *ptr.add(idx + 1);
61 a * (1.0 - frac) + b * frac
62 };
63 result.to_bits()
64 })
65}
66
67extern "C" fn jit_shuffle(input: u64, feedback: u64, size: u64, min: u64) -> u64 {
69 guarded(|| crate::numeric::permute::shuffle_bounded(input, feedback, size, min))
72}
73
74extern "C" fn jit_sin(bits: u64) -> u64 {
76 guarded(|| f64::from_bits(bits).sin().to_bits())
77}
78extern "C" fn jit_cos(bits: u64) -> u64 {
79 guarded(|| f64::from_bits(bits).cos().to_bits())
80}
81extern "C" fn jit_tan(bits: u64) -> u64 {
82 guarded(|| f64::from_bits(bits).tan().to_bits())
83}
84extern "C" fn jit_asin(bits: u64) -> u64 {
85 guarded(|| f64::from_bits(bits).asin().to_bits())
86}
87extern "C" fn jit_acos(bits: u64) -> u64 {
88 guarded(|| f64::from_bits(bits).acos().to_bits())
89}
90extern "C" fn jit_atan(bits: u64) -> u64 {
91 guarded(|| f64::from_bits(bits).atan().to_bits())
92}
93extern "C" fn jit_sqrt(bits: u64) -> u64 {
94 guarded(|| f64::from_bits(bits).sqrt().to_bits())
95}
96extern "C" fn jit_abs_f64(bits: u64) -> u64 {
97 guarded(|| f64::from_bits(bits).abs().to_bits())
98}
99extern "C" fn jit_ln(bits: u64) -> u64 {
100 guarded(|| f64::from_bits(bits).ln().to_bits())
101}
102extern "C" fn jit_exp(bits: u64) -> u64 {
103 guarded(|| f64::from_bits(bits).exp().to_bits())
104}
105extern "C" fn jit_floor_base10(bits: u64) -> u64 {
106 guarded(|| {
107 use crate::numeric::round_numbers::*;
108 let x = f64::from_bits(bits);
109 let r = if !positive_finite(x) {
110 0.0
111 } else {
112 floor_pow10(x)
113 };
114 r.to_bits()
115 })
116}
117extern "C" fn jit_ceiling_base10(bits: u64) -> u64 {
118 guarded(|| {
119 use crate::numeric::round_numbers::*;
120 let x = f64::from_bits(bits);
121 let r = if !positive_finite(x) {
122 0.0
123 } else {
124 let lo = floor_pow10(x);
125 if lo == x { lo } else { lo * 10.0 }
126 };
127 r.to_bits()
128 })
129}
130extern "C" fn jit_closest_base10(bits: u64) -> u64 {
131 guarded(|| {
132 use crate::numeric::round_numbers::*;
133 let x = f64::from_bits(bits);
134 let r = if !positive_finite(x) {
135 0.0
136 } else {
137 let lo = floor_pow10(x);
138 let hi = if lo == x { lo } else { lo * 10.0 };
139 pick_closest(x, lo, hi)
140 };
141 r.to_bits()
142 })
143}
144extern "C" fn jit_floor_decade(bits: u64) -> u64 {
145 guarded(|| {
146 use crate::numeric::round_numbers::*;
147 let x = f64::from_bits(bits);
148 let r = if !positive_finite(x) {
149 0.0
150 } else {
151 let base = floor_pow10(x);
152 (x / base).floor() * base
153 };
154 r.to_bits()
155 })
156}
157extern "C" fn jit_ceiling_decade(bits: u64) -> u64 {
158 guarded(|| {
159 use crate::numeric::round_numbers::*;
160 let x = f64::from_bits(bits);
161 let r = if !positive_finite(x) {
162 0.0
163 } else {
164 let base = floor_pow10(x);
165 (x / base).ceil() * base
166 };
167 r.to_bits()
168 })
169}
170extern "C" fn jit_closest_decade(bits: u64) -> u64 {
171 guarded(|| {
172 use crate::numeric::round_numbers::*;
173 let x = f64::from_bits(bits);
174 let r = if !positive_finite(x) {
175 0.0
176 } else {
177 let base = floor_pow10(x);
178 (x / base).round() * base
179 };
180 r.to_bits()
181 })
182}
183extern "C" fn jit_floor_binomial(bits: u64) -> u64 {
184 guarded(|| {
185 use crate::numeric::round_numbers::*;
186 let x = f64::from_bits(bits);
187 let r = if !positive_finite(x) {
188 0.0
189 } else {
190 floor_pow2(x)
191 };
192 r.to_bits()
193 })
194}
195extern "C" fn jit_ceiling_binomial(bits: u64) -> u64 {
196 guarded(|| {
197 use crate::numeric::round_numbers::*;
198 let x = f64::from_bits(bits);
199 let r = if !positive_finite(x) {
200 0.0
201 } else {
202 let lo = floor_pow2(x);
203 if lo == x { lo } else { lo * 2.0 }
204 };
205 r.to_bits()
206 })
207}
208extern "C" fn jit_closest_binomial(bits: u64) -> u64 {
209 guarded(|| {
210 use crate::numeric::round_numbers::*;
211 let x = f64::from_bits(bits);
212 let r = if !positive_finite(x) {
213 0.0
214 } else {
215 let lo = floor_pow2(x);
216 let hi = if lo == x { lo } else { lo * 2.0 };
217 pick_closest(x, lo, hi)
218 };
219 r.to_bits()
220 })
221}
222extern "C" fn jit_floor_fibonacci(bits: u64) -> u64 {
223 guarded(|| {
224 use crate::numeric::round_numbers::*;
225 let x = f64::from_bits(bits);
226 let r = if !positive_finite(x) {
227 0.0
228 } else {
229 floor_fibonacci_val(x)
230 };
231 r.to_bits()
232 })
233}
234extern "C" fn jit_ceiling_fibonacci(bits: u64) -> u64 {
235 guarded(|| {
236 use crate::numeric::round_numbers::*;
237 let x = f64::from_bits(bits);
238 let r = if !positive_finite(x) {
239 0.0
240 } else {
241 ceiling_fibonacci_val(x)
242 };
243 r.to_bits()
244 })
245}
246extern "C" fn jit_closest_fibonacci(bits: u64) -> u64 {
247 guarded(|| {
248 use crate::numeric::round_numbers::*;
249 let x = f64::from_bits(bits);
250 let r = if !positive_finite(x) {
251 0.0
252 } else {
253 pick_closest(x, floor_fibonacci_val(x), ceiling_fibonacci_val(x))
254 };
255 r.to_bits()
256 })
257}
258
259extern "C" fn jit_atan2(y_bits: u64, x_bits: u64) -> u64 {
260 guarded(|| {
261 f64::from_bits(y_bits)
262 .atan2(f64::from_bits(x_bits))
263 .to_bits()
264 })
265}
266extern "C" fn jit_pow(base_bits: u64, exp_bits: u64) -> u64 {
267 guarded(|| {
268 f64::from_bits(base_bits)
269 .powf(f64::from_bits(exp_bits))
270 .to_bits()
271 })
272}
273extern "C" fn jit_round_nearest(x_bits: u64, iv_bits: u64) -> u64 {
274 guarded(|| {
275 let x = f64::from_bits(x_bits);
276 let interval = f64::from_bits(iv_bits);
277 let r = if !(interval.is_finite() && interval > 0.0) {
278 x
279 } else {
280 (x / interval).round() * interval
281 };
282 r.to_bits()
283 })
284}
285extern "C" fn jit_round_floor(x_bits: u64, iv_bits: u64) -> u64 {
286 guarded(|| {
287 let x = f64::from_bits(x_bits);
288 let interval = f64::from_bits(iv_bits);
289 let r = if !(interval.is_finite() && interval > 0.0) {
290 x
291 } else {
292 (x / interval).floor() * interval
293 };
294 r.to_bits()
295 })
296}
297extern "C" fn jit_round_ceiling(x_bits: u64, iv_bits: u64) -> u64 {
298 guarded(|| {
299 let x = f64::from_bits(x_bits);
300 let interval = f64::from_bits(iv_bits);
301 let r = if !(interval.is_finite() && interval > 0.0) {
302 x
303 } else {
304 (x / interval).ceil() * interval
305 };
306 r.to_bits()
307 })
308}
309
310extern "C" fn jit_pcg(input: u64, seed: u64, stream: u64) -> u64 {
311 guarded(|| {
312 let inc = 2u64.wrapping_mul(stream).wrapping_add(1);
313 crate::numeric::pcg::pcg_seek(seed, inc, input)
314 })
315}
316extern "C" fn jit_pcg_stream(input: u64, stream: u64, seed: u64) -> u64 {
317 guarded(|| {
318 let inc = 2u64.wrapping_mul(stream).wrapping_add(1);
319 crate::numeric::pcg::pcg_seek(seed, inc, input)
320 })
321}
322extern "C" fn jit_n_of(input: u64, n: u64, m: u64) -> u64 {
323 guarded(|| {
324 if m == 0 {
325 return 0;
326 }
327 crate::numeric::n_of_m::n_of_m_eval(input, n, m)
328 })
329}
330
331extern "C" fn jit_cycle_walk(pos: u64, range: u64, seed: u64, inc: u64) -> u64 {
332 guarded(|| {
333 let stream = inc.saturating_sub(1) / 2;
334 let state = crate::numeric::pcg::build_cycle_walk_state(range, seed, stream);
335 crate::numeric::pcg::cycle_walk_inner(
336 pos,
337 range,
338 state.half_bits,
339 state.half_mask,
340 &state.round_keys,
341 )
342 })
343}
344
345extern "C" fn jit_perlin_1d(input: u64, perm_ptr: u64, freq_bits: u64) -> u64 {
346 guarded(|| {
347 let perm = unsafe { &*(perm_ptr as *const crate::numeric::noise::PermTable) };
348 let freq = f64::from_bits(freq_bits);
349 let r = crate::numeric::noise::perlin_1d_algo(perm, input as f64 * freq);
350 r.to_bits()
351 })
352}
353
354extern "C" fn jit_perlin_2d(x: u64, y: u64, perm_ptr: u64, freq_bits: u64) -> u64 {
355 guarded(|| {
356 let perm = unsafe { &*(perm_ptr as *const crate::numeric::noise::PermTable) };
357 let freq = f64::from_bits(freq_bits);
358 let r = crate::numeric::noise::perlin_2d_algo(perm, x as f64 * freq, y as f64 * freq);
359 r.to_bits()
360 })
361}
362
363extern "C" fn jit_simplex_2d(x: u64, y: u64, perm_ptr: u64, freq_bits: u64) -> u64 {
364 guarded(|| {
365 let perm = unsafe { &*(perm_ptr as *const crate::numeric::noise::PermTable) };
366 let freq = f64::from_bits(freq_bits);
367 let r = crate::numeric::noise::simplex_2d_algo(perm, x as f64 * freq, y as f64 * freq);
368 r.to_bits()
369 })
370}
371
372extern "C" fn jit_fractal_noise_1d(input: u64, perm_ptr: u64, freq_bits: u64, octaves: u64) -> u64 {
373 guarded(|| {
374 let perm = unsafe { &*(perm_ptr as *const crate::numeric::noise::PermTable) };
375 let freq = f64::from_bits(freq_bits);
376 let r = crate::numeric::noise::fbm_1d(perm, input as f64, freq, octaves as u32);
377 r.to_bits()
378 })
379}
380
381extern "C" fn jit_fractal_noise_2d(
382 x: u64,
383 y: u64,
384 perm_ptr: u64,
385 freq_bits: u64,
386 octaves: u64,
387) -> u64 {
388 guarded(|| {
389 let perm = unsafe { &*(perm_ptr as *const crate::numeric::noise::PermTable) };
390 let freq = f64::from_bits(freq_bits);
391 let r = crate::numeric::noise::fbm_2d(perm, x as f64, y as f64, freq, octaves as u32);
392 r.to_bits()
393 })
394}
395
396extern "C" fn jit_thread_id() -> u64 {
397 guarded(|| {
398 let id = std::thread::current().id();
399 let id_str = format!("{id:?}");
400 let num = id_str.trim_start_matches("ThreadId(").trim_end_matches(')');
401 num.parse().unwrap_or(0)
402 })
403}
404
405extern "C" fn jit_current_epoch_millis() -> u64 {
406 guarded(|| {
407 std::time::SystemTime::now()
408 .duration_since(std::time::UNIX_EPOCH)
409 .unwrap()
410 .as_millis() as u64
411 })
412}
413
414#[repr(C, align(16))]
451struct JitJmpBuf([u8; 512]);
452
453#[cfg(not(windows))]
454unsafe extern "C" {
455 fn _setjmp(env: *mut JitJmpBuf) -> i32;
456 fn _longjmp(env: *mut JitJmpBuf, val: i32) -> !;
457}
458
459#[cfg(windows)]
472unsafe extern "C" {
473 fn _setjmp(env: *mut JitJmpBuf, frame: *mut std::ffi::c_void) -> i32;
474 #[link_name = "longjmp"]
475 fn _longjmp(env: *mut JitJmpBuf, val: i32) -> !;
476}
477
478use std::cell::{Cell, RefCell};
479thread_local! {
480 static JIT_JMP_BUF: Cell<Option<*mut JitJmpBuf>> = const { Cell::new(None) };
487 static JIT_VIOLATION_MSG: RefCell<Option<String>> = const { RefCell::new(None) };
491}
492
493fn jit_violation_longjmp(msg: String) -> ! {
500 JIT_VIOLATION_MSG.with(|m| *m.borrow_mut() = Some(msg.clone()));
501 let buf_ptr: Option<*mut JitJmpBuf> = JIT_JMP_BUF.with(|b| b.get());
502 match buf_ptr {
503 Some(ptr) => unsafe { _longjmp(ptr, 1) },
504 None => {
505 let mut err = std::io::stderr().lock();
506 use std::io::Write;
507 let _ = writeln!(err, "{msg}");
508 let _ = err.flush();
509 std::process::abort();
510 }
511 }
512}
513
514struct JmpBufGuard {
522 prev: Option<*mut JitJmpBuf>,
523}
524
525impl Drop for JmpBufGuard {
526 fn drop(&mut self) {
527 JIT_JMP_BUF.with(|b| b.set(self.prev));
528 }
529}
530
531pub(crate) fn invoke_with_catch<F: FnOnce()>(f: F) {
542 use std::mem::MaybeUninit;
543 let mut buf: MaybeUninit<JitJmpBuf> = MaybeUninit::uninit();
544 let buf_ptr = buf.as_mut_ptr();
545 let prev: Option<*mut JitJmpBuf> = JIT_JMP_BUF.with(|b| b.replace(Some(buf_ptr)));
550 let _guard = JmpBufGuard { prev };
551 #[cfg(not(windows))]
552 let jmpval = unsafe { _setjmp(buf_ptr) };
553 #[cfg(windows)]
555 let jmpval = unsafe { _setjmp(buf_ptr, std::ptr::null_mut()) };
556 if jmpval == 0 {
557 f();
558 } else {
559 let msg = JIT_VIOLATION_MSG
564 .with(|m| m.borrow_mut().take())
565 .unwrap_or_else(|| "JIT predicate violation (no message)".into());
566 std::panic::resume_unwind(Box::new(msg));
570 }
571}
572
573extern "C" fn jit_is_positive_fail(value: u64, name_ptr: u64, name_len: u64) -> u64 {
577 let name = if name_ptr != 0 {
582 unsafe {
583 std::str::from_utf8_unchecked(std::slice::from_raw_parts(
584 name_ptr as *const u8,
585 name_len as usize,
586 ))
587 }
588 } else {
589 "value"
590 };
591 jit_violation_longjmp(format!(
592 "is_positive({name}): value must be > 0, got {value}"
593 ));
594}
595
596extern "C" fn jit_in_range_fail(value: u64, lo: u64, hi: u64) -> u64 {
599 jit_violation_longjmp(format!("in_range: value {value} outside [{lo}, {hi}]"));
600}
601
602extern "C" fn jit_div_zero_fail(kind: u64) -> u64 {
606 jit_violation_longjmp(
607 if kind == 0 {
608 "attempt to divide by zero"
609 } else {
610 "attempt to calculate the remainder with a divisor of zero"
611 }
612 .to_string(),
613 );
614}
615
616extern "C" fn jit_f64_mod(a_bits: u64, b_bits: u64) -> u64 {
620 let (a, b) = (f64::from_bits(a_bits), f64::from_bits(b_bits));
621 (if b != 0.0 { a % b } else { 0.0 }).to_bits()
622}
623
624extern "C" fn jit_is_one_of_fail(value: u64, set_ptr: u64, set_len: u64) -> u64 {
632 let msg = if set_ptr != 0 {
633 let set = unsafe { std::slice::from_raw_parts(set_ptr as *const u64, set_len as usize) };
634 format!("is_one_of: value {value} not in allowed set {set:?}")
635 } else {
636 format!("is_one_of: value {value} not in allowed set [..]")
637 };
638 jit_violation_longjmp(msg);
639}
640
641extern "C" fn jit_weighted_pick(
646 input: u64,
647 values_ptr: u64,
648 biases_ptr: u64,
649 primaries_ptr: u64,
650 aliases_ptr: u64,
651 n: u64,
652) -> u64 {
653 guarded(|| {
654 let n = n as usize;
655 let slot = (input as usize) % n;
656 let bias_test = ((input >> 32) as f64) / (u32::MAX as f64);
657 unsafe {
658 let biases = std::slice::from_raw_parts(biases_ptr as *const f64, n);
659 let primaries = std::slice::from_raw_parts(primaries_ptr as *const u64, n);
660 let aliases = std::slice::from_raw_parts(aliases_ptr as *const u64, n);
661 let values = std::slice::from_raw_parts(values_ptr as *const u64, n);
662 let index = if bias_test < biases[slot] {
663 primaries[slot]
664 } else {
665 aliases[slot]
666 };
667 values[index as usize]
668 }
669 })
670}
671
672fn guarded<T>(body: impl FnOnce() -> T) -> T {
677 match std::panic::catch_unwind(std::panic::AssertUnwindSafe(body)) {
678 Ok(v) => v,
679 Err(payload) => {
680 let msg = payload
681 .downcast_ref::<String>()
682 .cloned()
683 .or_else(|| payload.downcast_ref::<&str>().map(|s| s.to_string()))
684 .unwrap_or_else(|| "panic in a compiled helper".to_string());
685 jit_violation_longjmp(msg)
686 }
687 }
688}
689
690#[derive(Clone)]
693pub struct SlotKitRef(pub std::sync::Arc<crate::ast::CompiledSlotKit>);
694
695impl SlotKitRef {
696 fn new(kit: crate::ast::CompiledSlotKit) -> Self {
697 SlotKitRef(std::sync::Arc::new(kit))
698 }
699}
700
701impl std::fmt::Debug for SlotKitRef {
702 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
703 write!(
704 f,
705 "SlotKitRef({:p}, {} scratch)",
706 std::sync::Arc::as_ptr(&self.0),
707 self.0.scratch.len()
708 )
709 }
710}
711
712impl PartialEq for SlotKitRef {
713 fn eq(&self, other: &Self) -> bool {
714 std::sync::Arc::ptr_eq(&self.0, &other.0)
715 }
716}
717
718extern "C" fn jit_slot_call(
730 kit: *const crate::ast::CompiledSlotKit,
731 inputs: *const u64,
732 n_in: u64,
733 outputs: *mut u64,
734 n_out: u64,
735 scratch: *mut crate::ast::ScratchBuf,
736 base: u64,
737 n_scratch: u64,
738) {
739 guarded(|| unsafe {
740 let kit = &*kit;
741 let ins = std::slice::from_raw_parts(inputs, n_in as usize);
742 let outs = std::slice::from_raw_parts_mut(outputs, n_out as usize);
743 let sc = std::slice::from_raw_parts_mut(scratch.add(base as usize), n_scratch as usize);
744 (kit.op)(ins, outs, sc)
745 })
746}
747
748impl JitOp {
749 pub(crate) fn slot_kit(&self) -> Option<&SlotKitRef> {
751 match self {
752 JitOp::SlotCall { kit, .. } | JitOp::Convert { kit, .. } => Some(kit),
753 _ => None,
754 }
755 }
756
757 pub(crate) fn scratch_elems(&self) -> &[crate::ast::ScratchElem] {
759 const STR_ENTRY: [crate::ast::ScratchElem; 1] = [crate::ast::ScratchElem::Str];
760 const F32_ENTRY: [crate::ast::ScratchElem; 1] = [crate::ast::ScratchElem::F32];
761 match self {
762 JitOp::SlotCall { kit, .. } | JitOp::Convert { kit, .. } => &kit.0.scratch,
763 JitOp::U64ToStr { .. }
764 | JitOp::I64ToStr { .. }
765 | JitOp::F64ToStr { .. }
766 | JitOp::StrConcat { .. }
767 | JitOp::JsonToStr { .. } => &STR_ENTRY,
768 JitOp::VecProduce { .. } => &F32_ENTRY,
769 _ => &[],
770 }
771 }
772
773 pub(crate) fn place_scratch(&mut self, base: usize) {
776 match self {
777 JitOp::SlotCall { scratch_base, .. }
778 | JitOp::Convert { scratch_base, .. }
779 | JitOp::U64ToStr { scratch_base }
780 | JitOp::I64ToStr { scratch_base }
781 | JitOp::F64ToStr { scratch_base }
782 | JitOp::StrConcat { scratch_base }
783 | JitOp::JsonToStr { scratch_base }
784 | JitOp::VecProduce { scratch_base, .. } => *scratch_base = base,
785 _ => {}
786 }
787 }
788}
789
790#[derive(Debug, Clone, Copy, PartialEq, Eq)]
794pub enum VecProducer {
795 Add,
797 Scale,
799 Norm,
801 HashVec,
803 XxHash3Vec,
805 RegToVec,
807}
808
809#[derive(Debug, Clone, Copy, PartialEq, Eq)]
812pub enum VecReducer {
813 Dot,
815 L2,
817 Cosine,
819 LidMle,
821}
822
823#[derive(Debug, Clone, Copy, PartialEq, Eq)]
826pub enum RegLaneRead {
827 F32,
829 I16,
831 I64,
833}
834
835#[derive(Debug, Clone, Copy, PartialEq, Eq)]
839pub enum RegProducer {
840 WithLaneF32,
842 GatherF32,
844 VecToRegF32,
846 MulI8,
848}
849
850unsafe fn vec_f32_of<'a>(ptr: u64, len: u64) -> &'a [f32] {
856 if len == 0 {
857 &[]
858 } else {
859 unsafe { std::slice::from_raw_parts(ptr as usize as *const f32, len as usize) }
860 }
861}
862
863unsafe fn write_f32_entry(
870 scratch: *mut crate::ast::ScratchBuf,
871 base: u64,
872 buffer: *mut u64,
873 out_slot: u64,
874 f: impl FnOnce(&mut Vec<f32>),
875) {
876 unsafe {
877 let entry = &mut *scratch.add(base as usize);
878 let crate::ast::ScratchBuf::F32(v) = entry else {
879 panic!("a vector lowering's scratch entry is not an f32 vector");
880 };
881 f(v);
882 *buffer.add(out_slot as usize) = v.as_ptr() as usize as u64;
883 *buffer.add(out_slot as usize + 1) = v.len() as u64;
884 }
885}
886
887macro_rules! vec_producer {
891 ($name:ident, |$out:ident, $w0:ident, $w1:ident, $w2:ident, $w3:ident| $body:expr) => {
892 extern "C" fn $name(
893 scratch: *mut crate::ast::ScratchBuf,
894 base: u64,
895 buffer: *mut u64,
896 out_slot: u64,
897 $w0: u64,
898 $w1: u64,
899 $w2: u64,
900 $w3: u64,
901 ) {
902 guarded(|| unsafe {
903 let _ = ($w2, $w3);
904 write_f32_entry(scratch, base, buffer, out_slot, |$out| $body)
905 })
906 }
907 };
908}
909
910vec_producer!(jit_vec_add, |out, a_ptr, a_len, b_ptr, b_len| {
911 let (a, b) = (vec_f32_of(a_ptr, a_len), vec_f32_of(b_ptr, b_len));
912 crate::numeric::vector::check_lens("vec_add", a.len(), b.len());
913 crate::numeric::vector::add_f32_into(a, b, out)
914});
915vec_producer!(jit_vec_scale, |out, a_ptr, a_len, k_bits, _z| {
916 let a = vec_f32_of(a_ptr, a_len);
917 crate::numeric::vector::scale_f32_into(a, f64::from_bits(k_bits) as f32, out)
918});
919vec_producer!(jit_vec_norm, |out, a_ptr, a_len, _y, _z| {
920 crate::numeric::vector::norm_f32_into(vec_f32_of(a_ptr, a_len), out)
921});
922vec_producer!(jit_hash_vec, |out, seed, dim, _y, _z| {
923 crate::numeric::vector::hash_vec_into(seed, dim, out)
924});
925vec_producer!(jit_xxhash3_vec, |out, seed, dim, _y, _z| {
926 crate::numeric::vector::xxhash3_vec_into(seed, dim, out)
927});
928vec_producer!(jit_reg_to_vec_f32, |out, lo, hi, _y, _z| {
929 out.clear();
930 out.extend_from_slice(&crate::ast::Bits128([lo, hi]).lanes_f32())
931});
932
933macro_rules! vec_reducer {
936 ($name:ident, |$w0:ident, $w1:ident, $w2:ident, $w3:ident| $body:expr) => {
937 extern "C" fn $name($w0: u64, $w1: u64, $w2: u64, $w3: u64) -> u64 {
938 guarded(|| unsafe {
939 let _ = ($w2, $w3);
940 let r: f64 = $body;
941 r.to_bits()
942 })
943 }
944 };
945}
946
947vec_reducer!(jit_vec_dot, |a_ptr, a_len, b_ptr, b_len| {
948 let (a, b) = (vec_f32_of(a_ptr, a_len), vec_f32_of(b_ptr, b_len));
949 crate::numeric::vector::check_lens("vec_dot", a.len(), b.len());
950 crate::numeric::vector::dot_f32(a, b) as f64
951});
952vec_reducer!(jit_vec_l2, |a_ptr, a_len, b_ptr, b_len| {
953 let (a, b) = (vec_f32_of(a_ptr, a_len), vec_f32_of(b_ptr, b_len));
954 crate::numeric::vector::check_lens("vec_l2", a.len(), b.len());
955 (crate::numeric::vector::l2sq_f32(a, b) as f64).sqrt()
956});
957vec_reducer!(jit_vec_cosine, |a_ptr, a_len, b_ptr, b_len| {
958 let (a, b) = (vec_f32_of(a_ptr, a_len), vec_f32_of(b_ptr, b_len));
959 crate::numeric::vector::check_lens("vec_cosine", a.len(), b.len());
960 crate::numeric::vector::cosine_f32(a, b)
961});
962vec_reducer!(jit_lid_mle, |d_ptr, d_len, k_bits, _z| {
963 crate::numeric::vector::lid_mle_of(vec_f32_of(d_ptr, d_len), f64::from_bits(k_bits))
964});
965
966extern "C" fn jit_reg_lane_f32(lo: u64, hi: u64, i: u64) -> u64 {
968 guarded(|| crate::numeric::register::lane_f32(crate::ast::Bits128([lo, hi]), i).to_bits())
969}
970
971extern "C" fn jit_reg_lane_i16(lo: u64, hi: u64, i: u64) -> u64 {
973 guarded(|| crate::numeric::register::lane_i16(crate::ast::Bits128([lo, hi]), i) as i64 as u64)
974}
975
976extern "C" fn jit_reg_lane_i64(lo: u64, hi: u64, i: u64) -> u64 {
978 guarded(|| crate::numeric::register::lane_i64(crate::ast::Bits128([lo, hi]), i) as u64)
979}
980
981macro_rules! reg_producer {
984 ($name:ident, |$w0:ident, $w1:ident, $w2:ident, $w3:ident| $body:expr) => {
985 extern "C" fn $name(
986 buffer: *mut u64,
987 out_slot: u64,
988 $w0: u64,
989 $w1: u64,
990 $w2: u64,
991 $w3: u64,
992 ) {
993 guarded(|| unsafe {
994 let _ = ($w2, $w3);
995 let r: crate::ast::Bits128 = $body;
996 *buffer.add(out_slot as usize) = r.0[0];
997 *buffer.add(out_slot as usize + 1) = r.0[1];
998 })
999 }
1000 };
1001}
1002
1003reg_producer!(jit_reg_with_lane_f32, |lo, hi, i, v_bits| {
1004 crate::numeric::register::with_lane_f32(
1005 crate::ast::Bits128([lo, hi]),
1006 i,
1007 f64::from_bits(v_bits),
1008 )
1009});
1010reg_producer!(jit_reg_gather_f32, |v_ptr, v_len, offset, _z| {
1011 crate::numeric::register::gather_f32(vec_f32_of(v_ptr, v_len), offset)
1012});
1013reg_producer!(jit_vec_to_reg_f32, |v_ptr, v_len, _y, _z| {
1014 crate::numeric::register::to_reg_f32(vec_f32_of(v_ptr, v_len))
1015});
1016reg_producer!(jit_reg_mul_i8, |a_lo, a_hi, b_lo, b_hi| {
1017 crate::numeric::register::mul_i8(
1018 crate::ast::Bits128([a_lo, a_hi]),
1019 crate::ast::Bits128([b_lo, b_hi]),
1020 )
1021});
1022
1023unsafe fn write_str_entry(
1032 scratch: *mut crate::ast::ScratchBuf,
1033 base: u64,
1034 buffer: *mut u64,
1035 out_slot: u64,
1036 f: impl FnOnce(&mut Vec<u8>),
1037) {
1038 unsafe {
1039 let entry = &mut *scratch.add(base as usize);
1040 let crate::ast::ScratchBuf::Str(v) = entry else {
1041 panic!("a string lowering's scratch entry is not a string");
1042 };
1043 v.clear();
1044 f(v);
1045 *buffer.add(out_slot as usize) = v.as_ptr() as usize as u64;
1046 *buffer.add(out_slot as usize + 1) = v.len() as u64;
1047 }
1048}
1049
1050extern "C" fn jit_u64_to_str(
1052 scratch: *mut crate::ast::ScratchBuf,
1053 base: u64,
1054 buffer: *mut u64,
1055 out_slot: u64,
1056 value: u64,
1057) {
1058 use std::io::Write;
1059 guarded(|| unsafe {
1060 write_str_entry(scratch, base, buffer, out_slot, |v| {
1061 write!(v, "{value}").expect("a vector accepts every write")
1062 })
1063 })
1064}
1065
1066extern "C" fn jit_i64_to_str(
1068 scratch: *mut crate::ast::ScratchBuf,
1069 base: u64,
1070 buffer: *mut u64,
1071 out_slot: u64,
1072 value: u64,
1073) {
1074 use std::io::Write;
1075 guarded(|| unsafe {
1076 write_str_entry(scratch, base, buffer, out_slot, |v| {
1077 write!(v, "{}", value as i64).expect("a vector accepts every write")
1078 })
1079 })
1080}
1081
1082extern "C" fn jit_f64_to_str(
1085 scratch: *mut crate::ast::ScratchBuf,
1086 base: u64,
1087 buffer: *mut u64,
1088 out_slot: u64,
1089 bits: u64,
1090) {
1091 use std::io::Write;
1092 guarded(|| unsafe {
1093 write_str_entry(scratch, base, buffer, out_slot, |v| {
1094 write!(v, "{}", f64::from_bits(bits)).expect("a vector accepts every write")
1095 })
1096 })
1097}
1098
1099extern "C" fn jit_str_concat(
1103 scratch: *mut crate::ast::ScratchBuf,
1104 base: u64,
1105 buffer: *mut u64,
1106 out_slot: u64,
1107 pairs: *const u64,
1108 n: u64,
1109) {
1110 guarded(|| unsafe {
1111 let words = std::slice::from_raw_parts(pairs, 2 * n as usize);
1112 write_str_entry(scratch, base, buffer, out_slot, |v| {
1113 for pair in words.as_chunks::<2>().0 {
1114 let bytes =
1118 std::slice::from_raw_parts(pair[0] as usize as *const u8, pair[1] as usize);
1119 v.extend_from_slice(bytes);
1120 }
1121 })
1122 })
1123}
1124
1125extern "C" fn jit_json_to_str(
1128 scratch: *mut crate::ast::ScratchBuf,
1129 base: u64,
1130 buffer: *mut u64,
1131 out_slot: u64,
1132 ptr: u64,
1133 len: u64,
1134) {
1135 guarded(|| unsafe {
1136 let pair = [ptr, len];
1137 let value = crate::derive_support::ref_value(&pair);
1138 let json = match value {
1139 crate::ast::Value::Json(j) => j.as_ref(),
1140 other => panic!("expected Json wire, got {other:?}"),
1141 };
1142 write_str_entry(scratch, base, buffer, out_slot, |v| {
1143 serde_json::to_writer(v, json).expect("a vector accepts every write")
1144 })
1145 })
1146}
1147
1148#[derive(Debug, Clone, Copy, PartialEq)]
1155enum Scalar {
1156 Unsigned(u32),
1157 Signed(u32),
1158 Bool,
1159 F32,
1160 F64,
1161}
1162
1163impl Scalar {
1164 fn of(t: crate::ast::PortType) -> Option<Self> {
1165 use crate::ast::PortType as P;
1166 Some(match t {
1167 P::U8 => Self::Unsigned(8),
1168 P::U16 => Self::Unsigned(16),
1169 P::U32 => Self::Unsigned(32),
1170 P::U64 => Self::Unsigned(64),
1171 P::I8 => Self::Signed(8),
1172 P::I16 => Self::Signed(16),
1173 P::I32 => Self::Signed(32),
1174 P::I64 => Self::Signed(64),
1175 P::Bool => Self::Bool,
1176 P::F32 => Self::F32,
1177 P::F64 => Self::F64,
1178 _ => return None,
1179 })
1180 }
1181
1182 fn int_range(self) -> Option<(i128, i128)> {
1184 match self {
1185 Self::Unsigned(b) => Some((0, (1i128 << b) - 1)),
1186 Self::Signed(b) => Some((-(1i128 << (b - 1)), (1i128 << (b - 1)) - 1)),
1187 Self::Bool => Some((0, 1)),
1188 Self::F32 | Self::F64 => None,
1189 }
1190 }
1191}
1192
1193fn conversion_op(node: &dyn PolydatNode) -> Option<JitOp> {
1204 use crate::ast::Slot;
1205 let meta = node.meta();
1206 let [Slot::Wire(input)] = meta.ins.as_slice() else {
1207 return None;
1208 };
1209 let [output] = meta.outs.as_slice() else {
1210 return None;
1211 };
1212 let (from, to) = (input.typ, output.typ);
1213 Scalar::of(from)?;
1214 Scalar::of(to)?;
1215 let canonical = crate::compile::assembly::boundary_adapter(from, to)?;
1216 if canonical.meta().name != meta.name {
1217 return None;
1218 }
1219 let op = node.compiled_u64()?;
1220 Some(JitOp::Convert {
1221 from,
1222 to,
1223 kit: SlotKitRef::new(crate::ast::CompiledSlotKit {
1224 scratch: Vec::new(),
1225 op: Box::new(move |inputs, outputs, _| op(inputs, outputs)),
1226 }),
1227 scratch_base: 0,
1228 })
1229}
1230
1231pub fn classify_node_typed(node: &dyn PolydatNode, wire_types: &[crate::ast::PortType]) -> JitOp {
1240 use crate::ast::PortType as PT;
1241 let is_ref = |t: &crate::ast::PortType| t.slot_color() == crate::ast::SlotColor::Ref2;
1242 let vec_produce = |kind: VecProducer| JitOp::VecProduce {
1243 kind,
1244 scratch_base: 0,
1245 };
1246 let ref_copy = |ty: crate::ast::PortType| {
1247 crate::compile::assembly::ref_copy_kit(ty)
1248 .map(|kit| JitOp::SlotCall {
1249 kit: SlotKitRef::new(kit),
1250 scratch_base: 0,
1251 })
1252 .unwrap_or(JitOp::Fallback)
1253 };
1254 let meta = node.meta();
1255 let named = match meta.name.as_str() {
1256 "__u64_to_string" => JitOp::U64ToStr { scratch_base: 0 },
1259 "__i64_to_string" => JitOp::I64ToStr { scratch_base: 0 },
1260 "__f64_to_string" => JitOp::F64ToStr { scratch_base: 0 },
1261 "json_to_str" if wire_types == [crate::ast::PortType::Json] => {
1262 JitOp::JsonToStr { scratch_base: 0 }
1263 }
1264 "str_concat"
1265 if !wire_types.is_empty()
1266 && wire_types.iter().all(|t| *t == crate::ast::PortType::Str) =>
1267 {
1268 JitOp::StrConcat { scratch_base: 0 }
1269 }
1270 "vec_add" if wire_types == [PT::VecF32, PT::VecF32] => vec_produce(VecProducer::Add),
1275 "vec_scale" if wire_types == [PT::VecF32, PT::F64] => vec_produce(VecProducer::Scale),
1276 "vec_norm" if wire_types == [PT::VecF32] => vec_produce(VecProducer::Norm),
1277 "hash_vec" if wire_types == [PT::U64, PT::U64] => vec_produce(VecProducer::HashVec),
1278 "xxhash3_vec" if wire_types == [PT::U64, PT::U64] => vec_produce(VecProducer::XxHash3Vec),
1279 "reg_to_vec_f32" if wire_types == [PT::RegF32x4] => vec_produce(VecProducer::RegToVec),
1280 "vec_dot" if wire_types == [PT::VecF32, PT::VecF32] => JitOp::VecReduce(VecReducer::Dot),
1281 "vec_l2" if wire_types == [PT::VecF32, PT::VecF32] => JitOp::VecReduce(VecReducer::L2),
1282 "vec_cosine" if wire_types == [PT::VecF32, PT::VecF32] => {
1283 JitOp::VecReduce(VecReducer::Cosine)
1284 }
1285 "lid_mle" if wire_types == [PT::VecF32, PT::F64] => JitOp::VecReduce(VecReducer::LidMle),
1286 "reg_lane_f32" if wire_types == [PT::RegF32x4, PT::U64] => JitOp::RegLane(RegLaneRead::F32),
1290 "reg_lane_i16" if wire_types == [PT::RegI16x8, PT::U64] => JitOp::RegLane(RegLaneRead::I16),
1291 "reg_lane_i64" if wire_types == [PT::RegI64x2, PT::U64] => JitOp::RegLane(RegLaneRead::I64),
1292 "reg_with_lane_f32" if wire_types == [PT::RegF32x4, PT::U64, PT::F64] => {
1293 JitOp::RegProduce(RegProducer::WithLaneF32)
1294 }
1295 "reg_gather_f32" if wire_types == [PT::VecF32, PT::U64] => {
1296 JitOp::RegProduce(RegProducer::GatherF32)
1297 }
1298 "vec_to_reg_f32" if wire_types == [PT::VecF32] => {
1299 JitOp::RegProduce(RegProducer::VecToRegF32)
1300 }
1301 "reg_mul_i8" if wire_types == [PT::RegI8x16, PT::RegI8x16] => {
1302 JitOp::RegProduce(RegProducer::MulI8)
1303 }
1304 "reg_dot_f32" if wire_types == [PT::RegF32x4, PT::RegF32x4] => JitOp::RegDotF32,
1305 n if n.starts_with("__port_") || n == "default_or" => match meta.outs.first() {
1311 Some(o) if is_ref(&o.typ) => return ref_copy(o.typ),
1312 _ => JitOp::Identity,
1313 },
1314 "select" | "select_u64" if wire_types.iter().skip(1).any(|t| t.slot_width() != 1) => {
1317 JitOp::Fallback
1318 }
1319 _ if wire_types.iter().any(is_ref) || meta.outs.iter().any(|o| is_ref(&o.typ)) => {
1320 JitOp::Fallback
1321 }
1322 _ => classify_node(node),
1323 };
1324 if !matches!(named, JitOp::Fallback) {
1325 return named;
1326 }
1327 if let Some(kit) = node.compiled_slot(
1328 wire_types,
1329 crate::compile::select::Engine::Native(crate::compile::select::Provenance::Auto),
1330 ) {
1331 return JitOp::SlotCall {
1332 kit: SlotKitRef::new(kit),
1333 scratch_base: 0,
1334 };
1335 }
1336 if let Some(op) = node.compiled_u64() {
1337 return JitOp::SlotCall {
1338 kit: SlotKitRef::new(crate::ast::CompiledSlotKit {
1339 scratch: Vec::new(),
1340 op: Box::new(move |inputs, outputs, _| op(inputs, outputs)),
1341 }),
1342 scratch_base: 0,
1343 };
1344 }
1345 JitOp::Fallback
1346}
1347
1348#[derive(Debug, Clone, PartialEq)]
1355pub enum JitOp {
1356 Identity,
1359 AddConst(u64),
1361 MulConst(u64),
1363 DivConst(u64),
1365 ModConst(u64),
1367 ClampConst(u64, u64),
1369 Interleave,
1371 MixedRadixConst(Vec<u64>),
1373 Hash,
1375 SplitMix64,
1377 ShuffleConst(u64, u64, u64),
1379
1380 UnitInterval,
1383 F64ToU64,
1385 RoundToU64,
1388 FloorToU64,
1390 CeilToU64,
1392 ClampF64Const(u64, u64), LerpConst(u64, u64), ScaleRangeConst(u64, u64), QuantizeConst(u64), DiscretizeConst(u64, u64), LutSampleConst(u64, u64), WeightedPickConst(u64, u64, u64, u64, u64), MathUnary(u8),
1414 MathBinary(u8),
1417
1418 U64Add2,
1421 U64Sub2,
1423 U64Mul2,
1425 U64Div2,
1427 U64Mod2,
1429 U64And,
1431 U64Or,
1433 U64Xor,
1435 U64Shl,
1437 U64Shr,
1439 U64Not,
1441
1442 ToF64,
1445
1446 F64Add,
1448 F64Sub,
1450 F64Mul,
1452 F64Div,
1454 F64Mod,
1456 U64DivWire,
1459 U64ModWire,
1462
1463 SlotCall {
1471 kit: SlotKitRef,
1474 scratch_base: usize,
1477 },
1478
1479 Convert {
1487 from: crate::ast::PortType,
1489 to: crate::ast::PortType,
1491 kit: SlotKitRef,
1493 scratch_base: usize,
1495 },
1496
1497 U64ToStr {
1503 scratch_base: usize,
1505 },
1506 I64ToStr {
1508 scratch_base: usize,
1510 },
1511 F64ToStr {
1513 scratch_base: usize,
1515 },
1516 StrConcat {
1519 scratch_base: usize,
1521 },
1522 JsonToStr {
1524 scratch_base: usize,
1526 },
1527
1528 VecProduce {
1532 kind: VecProducer,
1534 scratch_base: usize,
1536 },
1537 VecReduce(VecReducer),
1539 RegLane(RegLaneRead),
1542 RegProduce(RegProducer),
1544 RegDotF32,
1548 RegShuffleConst([u8; 16]),
1551
1552 IsPositiveCheck {
1559 name_ptr: u64,
1561 name_len: u64,
1563 },
1564 InRangeCheck(u64, u64),
1568 IsOneOfCheck {
1576 allowed: Vec<u64>,
1578 set_ptr: u64,
1580 set_len: u64,
1582 },
1583
1584 RegBinOp(u8, u8),
1592 RegCopy,
1596 RegSplat(u8),
1600
1601 U64Cmp(ir::condcodes::IntCC),
1604 F64Cmp(ir::condcodes::FloatCC),
1606 SelectU64,
1608 SelectF64,
1610
1611 I64ToF64,
1614 ToBool,
1616 ConstU64(u64),
1618 ConstF64(u64),
1620
1621 HashRangeConst(u64),
1624 HashIntervalConst(u64, u64),
1626 InvLerpConst(u64, u64),
1628 RemapConst(u64, u64, u64, u64),
1630
1631 EpochOffsetConst(u64),
1634 EpochScaleConst(u64),
1636 ThreadId,
1638 CurrentEpochMillis,
1640
1641 Perlin1dConst(u64, u64),
1644 Perlin2dConst(u64, u64),
1646 Simplex2dConst(u64, u64),
1648 FractalNoise1dConst(u64, u64, u64),
1650 FractalNoise2dConst(u64, u64, u64),
1652
1653 VariadicSum,
1656 VariadicProduct,
1658 VariadicMin,
1660 VariadicMax,
1662 CheckedAdd,
1664 CheckedSub,
1666 CheckedMul,
1668 CeilToMultiple,
1670 MultiplesAtLeast,
1672
1673 FairCoin,
1676 BlendConst(u64),
1678 LfsrStepConst(u64),
1681 PcgConst(u64, u64),
1683 PcgStreamConst(u64),
1685 CycleWalkConst(u64, u64, u64),
1687 UnfairCoinConst(u64),
1689 CoinFlipConst(u64),
1692 ChanceConst(u64),
1694 NOfConst(u64, u64),
1696
1697 Fallback,
1700}
1701
1702pub fn classify_node(node: &dyn PolydatNode) -> JitOp {
1709 if let Some(op) = conversion_op(node) {
1710 return op;
1711 }
1712 let name = node.meta().name.as_str();
1713 let consts = node.jit_constants();
1714
1715 match name {
1716 "identity" => JitOp::Identity,
1717 "hash" | "splitmix64" | "scatter" => JitOp::SplitMix64,
1718 "fair_coin" => JitOp::FairCoin,
1719 "unfair_coin" => {
1720 if let Some(&p) = consts.first() {
1721 JitOp::UnfairCoinConst(p)
1722 } else {
1723 JitOp::Fallback
1724 }
1725 }
1726 "chance" => {
1727 if let Some(&p) = consts.first() {
1728 JitOp::ChanceConst(p)
1729 } else {
1730 JitOp::Fallback
1731 }
1732 }
1733 "xxhash3" | "xxh3" => JitOp::Hash,
1734 "hash_range" => {
1735 if let Some(&c) = consts.first() {
1736 JitOp::HashRangeConst(c)
1737 } else {
1738 JitOp::Fallback
1739 }
1740 }
1741 "hash_interval" => {
1742 if consts.len() >= 2 {
1743 JitOp::HashIntervalConst(consts[0], consts[1])
1744 } else {
1745 JitOp::Fallback
1746 }
1747 }
1748 "add" => {
1749 if let Some(&c) = consts.first() {
1750 JitOp::AddConst(c)
1751 } else {
1752 JitOp::Fallback
1753 }
1754 }
1755 "mul" => {
1756 if let Some(&c) = consts.first() {
1757 JitOp::MulConst(c)
1758 } else {
1759 JitOp::Fallback
1760 }
1761 }
1762 "div" => {
1763 if let Some(&c) = consts.first() {
1764 JitOp::DivConst(c)
1765 } else {
1766 JitOp::Fallback
1767 }
1768 }
1769 "mod" => {
1770 if let Some(&c) = consts.first() {
1771 JitOp::ModConst(c)
1772 } else {
1773 JitOp::Fallback
1774 }
1775 }
1776 "clamp" => {
1777 if consts.len() >= 2 && consts[0] <= consts[1] {
1783 JitOp::ClampConst(consts[0], consts[1])
1784 } else {
1785 JitOp::Fallback
1786 }
1787 }
1788 "interleave" => JitOp::Interleave,
1789 "mixed_radix" => {
1790 if consts.is_empty() {
1791 JitOp::Fallback
1792 } else {
1793 JitOp::MixedRadixConst(consts)
1794 }
1795 }
1796 "shuffle" => {
1797 if consts.len() >= 3 {
1798 JitOp::ShuffleConst(consts[0], consts[1], consts[2])
1799 } else {
1800 JitOp::Fallback
1801 }
1802 }
1803 "unit_interval" => JitOp::UnitInterval,
1805 "f64_to_u64" => JitOp::F64ToU64,
1806 "round_to_u64" => JitOp::RoundToU64,
1807 "floor_to_u64" => JitOp::FloorToU64,
1808 "ceil_to_u64" => JitOp::CeilToU64,
1809 "clamp_f64" => {
1810 if consts.len() >= 2 {
1811 JitOp::ClampF64Const(consts[0], consts[1])
1812 } else {
1813 JitOp::Fallback
1814 }
1815 }
1816 "lerp" => {
1817 if consts.len() >= 2 {
1818 JitOp::LerpConst(consts[0], consts[1])
1819 } else {
1820 JitOp::Fallback
1821 }
1822 }
1823 "scale_range" => {
1824 if consts.len() >= 2 {
1825 JitOp::ScaleRangeConst(consts[0], consts[1])
1826 } else {
1827 JitOp::Fallback
1828 }
1829 }
1830 "quantize" => {
1831 if let Some(&c) = consts.first() {
1832 JitOp::QuantizeConst(c)
1833 } else {
1834 JitOp::Fallback
1835 }
1836 }
1837 "discretize" => {
1838 if consts.len() >= 2 {
1839 JitOp::DiscretizeConst(consts[0], consts[1])
1840 } else {
1841 JitOp::Fallback
1842 }
1843 }
1844 "lut_sample" | "dist_normal" | "icd_normal" | "dist_exponential" | "icd_exponential"
1845 | "dist_uniform" | "dist_pareto" | "dist_zipf" | "dist_empirical" => {
1846 if consts.len() >= 2 {
1847 JitOp::LutSampleConst(consts[0], consts[1])
1848 } else {
1849 JitOp::Fallback
1850 }
1851 }
1852 "sin" => JitOp::MathUnary(0),
1854 "cos" => JitOp::MathUnary(1),
1855 "tan" => JitOp::MathUnary(2),
1856 "asin" => JitOp::MathUnary(3),
1857 "acos" => JitOp::MathUnary(4),
1858 "atan" => JitOp::MathUnary(5),
1859 "sqrt" => JitOp::MathUnary(6),
1860 "abs_f64" => JitOp::MathUnary(7),
1861 "ln" => JitOp::MathUnary(8),
1862 "exp" => JitOp::MathUnary(9),
1863 "floor_base10" => JitOp::MathUnary(10),
1864 "ceiling_base10" => JitOp::MathUnary(11),
1865 "closest_base10" => JitOp::MathUnary(12),
1866 "floor_decade" => JitOp::MathUnary(13),
1867 "ceiling_decade" => JitOp::MathUnary(14),
1868 "closest_decade" => JitOp::MathUnary(15),
1869 "floor_binomial" => JitOp::MathUnary(16),
1870 "ceiling_binomial" => JitOp::MathUnary(17),
1871 "closest_binomial" => JitOp::MathUnary(18),
1872 "floor_fibonacci" => JitOp::MathUnary(19),
1873 "ceiling_fibonacci" => JitOp::MathUnary(20),
1874 "closest_fibonacci" => JitOp::MathUnary(21),
1875 "atan2" => JitOp::MathBinary(0),
1876 "pow" => JitOp::MathBinary(1),
1877 "round_nearest" => JitOp::MathBinary(2),
1878 "round_floor" => JitOp::MathBinary(3),
1879 "round_ceiling" => JitOp::MathBinary(4),
1880 "to_f64" => JitOp::ToF64,
1881 "u64_add" => JitOp::U64Add2,
1883 "u64_sub" => JitOp::U64Sub2,
1884 "u64_mul" => JitOp::U64Mul2,
1885 "u64_div" => JitOp::U64Div2,
1886 "u64_mod" => JitOp::U64Mod2,
1887 "u64_and" => JitOp::U64And,
1888 "u64_or" => JitOp::U64Or,
1889 "u64_xor" => JitOp::U64Xor,
1890 "u64_shl" => JitOp::U64Shl,
1891 "u64_shr" => JitOp::U64Shr,
1892 "u64_not" => JitOp::U64Not,
1893
1894 "reg_add_i8" => JitOp::RegBinOp(0, 0),
1896 "reg_sub_i8" => JitOp::RegBinOp(0, 1),
1897 "reg_shuffle_bytes" => {
1902 let mut mask = [0u8; 16];
1903 if consts.len() == 16 && consts.iter().all(|&m| m < 16) {
1904 for (m, &c) in mask.iter_mut().zip(consts.iter()) {
1905 *m = c as u8;
1906 }
1907 JitOp::RegShuffleConst(mask)
1908 } else {
1909 JitOp::Fallback
1910 }
1911 }
1912 "reg_add_i16" => JitOp::RegBinOp(1, 0),
1913 "reg_sub_i16" => JitOp::RegBinOp(1, 1),
1914 "reg_mul_i16" => JitOp::RegBinOp(1, 2),
1915 "reg_add_i32" => JitOp::RegBinOp(2, 0),
1916 "reg_sub_i32" => JitOp::RegBinOp(2, 1),
1917 "reg_mul_i32" => JitOp::RegBinOp(2, 2),
1918 "reg_add_i64" => JitOp::RegBinOp(3, 0),
1919 "reg_sub_i64" => JitOp::RegBinOp(3, 1),
1920 "reg_mul_i64" => JitOp::RegBinOp(3, 2),
1921 "reg_add_f32" => JitOp::RegBinOp(4, 0),
1922 "reg_sub_f32" => JitOp::RegBinOp(4, 1),
1923 "reg_mul_f32" => JitOp::RegBinOp(4, 2),
1924 "reg_add_f64" => JitOp::RegBinOp(5, 0),
1925 "reg_sub_f64" => JitOp::RegBinOp(5, 1),
1926 "reg_mul_f64" => JitOp::RegBinOp(5, 2),
1927 "__reg_view_raw" | "__reg_view_i8x16" | "__reg_view_i16x8" | "__reg_view_i32x4"
1928 | "__reg_view_i64x2" | "__reg_view_f16x8" | "__reg_view_f32x4" | "__reg_view_f64x2" => {
1929 JitOp::RegCopy
1930 }
1931 "reg_splat_i8" => JitOp::RegSplat(0),
1932 "reg_splat_i16" => JitOp::RegSplat(1),
1933 "reg_splat_i32" => JitOp::RegSplat(2),
1934 "reg_splat_i64" => JitOp::RegSplat(3),
1935 "reg_splat_f32" => JitOp::RegSplat(4),
1936 "reg_splat_f64" => JitOp::RegSplat(5),
1937
1938 "f64_add" => JitOp::F64Add,
1939 "f64_sub" => JitOp::F64Sub,
1940 "f64_mul" => JitOp::F64Mul,
1941 "f64_div" => JitOp::F64Div,
1942 "f64_mod" => JitOp::F64Mod,
1943
1944 "u64_eq" => JitOp::U64Cmp(ir::condcodes::IntCC::Equal),
1946 "u64_ne" => JitOp::U64Cmp(ir::condcodes::IntCC::NotEqual),
1947 "u64_lt" => JitOp::U64Cmp(ir::condcodes::IntCC::UnsignedLessThan),
1948 "u64_le" => JitOp::U64Cmp(ir::condcodes::IntCC::UnsignedLessThanOrEqual),
1949 "u64_gt" => JitOp::U64Cmp(ir::condcodes::IntCC::UnsignedGreaterThan),
1950 "u64_ge" => JitOp::U64Cmp(ir::condcodes::IntCC::UnsignedGreaterThanOrEqual),
1951 "f64_eq" => JitOp::F64Cmp(ir::condcodes::FloatCC::Equal),
1952 "f64_ne" => JitOp::F64Cmp(ir::condcodes::FloatCC::NotEqual),
1953 "f64_lt" => JitOp::F64Cmp(ir::condcodes::FloatCC::LessThan),
1954 "f64_le" => JitOp::F64Cmp(ir::condcodes::FloatCC::LessThanOrEqual),
1955 "f64_gt" => JitOp::F64Cmp(ir::condcodes::FloatCC::GreaterThan),
1956 "f64_ge" => JitOp::F64Cmp(ir::condcodes::FloatCC::GreaterThanOrEqual),
1957 "select_u64" | "select" => JitOp::SelectU64,
1958 "select_f64" => JitOp::SelectF64,
1959
1960 "div_wire" => JitOp::U64DivWire,
1962 "mod_wire" => JitOp::U64ModWire,
1963 "ceil_to_multiple" => JitOp::CeilToMultiple,
1964 "multiples_at_least" => JitOp::MultiplesAtLeast,
1965 "checked_add" => JitOp::CheckedAdd,
1966 "checked_sub" => JitOp::CheckedSub,
1967 "checked_mul" => JitOp::CheckedMul,
1968
1969 "sum" => JitOp::VariadicSum,
1971 "product" => JitOp::VariadicProduct,
1972 "min" => JitOp::VariadicMin,
1973 "max" => JitOp::VariadicMax,
1974
1975 "blend" => {
1977 if let Some(&c) = consts.first() {
1981 JitOp::BlendConst(c)
1982 } else {
1983 JitOp::Fallback
1984 }
1985 }
1986 "lfsr_step" => {
1987 if let Some(&fb) = consts.first() {
1988 JitOp::LfsrStepConst(fb)
1989 } else {
1990 JitOp::Fallback
1991 }
1992 }
1993 "pcg" => {
1994 if consts.len() >= 2 {
1995 JitOp::PcgConst(consts[0], consts[1])
1996 } else {
1997 JitOp::Fallback
1998 }
1999 }
2000 "pcg_stream" => {
2001 if let Some(&seed) = consts.first() {
2002 JitOp::PcgStreamConst(seed)
2003 } else {
2004 JitOp::Fallback
2005 }
2006 }
2007 "n_of" => {
2008 if consts.len() >= 2 {
2009 JitOp::NOfConst(consts[0], consts[1])
2010 } else {
2011 JitOp::Fallback
2012 }
2013 }
2014
2015 "cycle_walk" => {
2016 if consts.len() >= 3 {
2017 JitOp::CycleWalkConst(consts[0], consts[1], consts[2])
2018 } else {
2019 JitOp::Fallback
2020 }
2021 }
2022 "coin_flip" => {
2023 if let Some(&threshold) = consts.first() {
2028 JitOp::CoinFlipConst(threshold)
2029 } else {
2030 JitOp::Fallback
2031 }
2032 }
2033 "default_or" => JitOp::Identity,
2036 "const_u64" | "const_bool" | "session_start_millis" => {
2037 if let Some(&c) = consts.first() {
2038 JitOp::ConstU64(c)
2039 } else {
2040 JitOp::Fallback
2041 }
2042 }
2043 "const_f64" => {
2044 if let Some(&c) = consts.first() {
2045 JitOp::ConstF64(c)
2046 } else {
2047 JitOp::Fallback
2048 }
2049 }
2050 "inv_lerp" => {
2051 if consts.len() >= 2 {
2052 JitOp::InvLerpConst(consts[0], consts[1])
2053 } else {
2054 JitOp::Fallback
2055 }
2056 }
2057 "remap" => {
2058 if consts.len() >= 4 {
2059 JitOp::RemapConst(consts[0], consts[1], consts[2], consts[3])
2060 } else {
2061 JitOp::Fallback
2062 }
2063 }
2064 "epoch_offset" => {
2065 if let Some(&c) = consts.first() {
2066 JitOp::EpochOffsetConst(c)
2067 } else {
2068 JitOp::Fallback
2069 }
2070 }
2071 "epoch_scale" => {
2072 if let Some(&c) = consts.first() {
2073 JitOp::EpochScaleConst(c)
2074 } else {
2075 JitOp::Fallback
2076 }
2077 }
2078 "thread_id" => JitOp::ThreadId,
2079 "current_epoch_millis" => JitOp::CurrentEpochMillis,
2080 "perlin_1d" => {
2081 if consts.len() >= 2 {
2082 JitOp::Perlin1dConst(consts[0], consts[1])
2083 } else {
2084 JitOp::Fallback
2085 }
2086 }
2087 "perlin_2d" => {
2088 if consts.len() >= 2 {
2089 JitOp::Perlin2dConst(consts[0], consts[1])
2090 } else {
2091 JitOp::Fallback
2092 }
2093 }
2094 "simplex_2d" => {
2095 if consts.len() >= 2 {
2096 JitOp::Simplex2dConst(consts[0], consts[1])
2097 } else {
2098 JitOp::Fallback
2099 }
2100 }
2101 "fractal_noise_1d" => {
2102 if consts.len() >= 3 {
2103 JitOp::FractalNoise1dConst(consts[0], consts[1], consts[2])
2104 } else {
2105 JitOp::Fallback
2106 }
2107 }
2108 "fractal_noise_2d" => {
2109 if consts.len() >= 3 {
2110 JitOp::FractalNoise2dConst(consts[0], consts[1], consts[2])
2111 } else {
2112 JitOp::Fallback
2113 }
2114 }
2115
2116 "trunc_u64" => JitOp::F64ToU64,
2120 "round_u64" => JitOp::RoundToU64,
2123
2124 "weighted_pick" => {
2125 if consts.len() >= 5 {
2126 JitOp::WeightedPickConst(consts[0], consts[1], consts[2], consts[3], consts[4])
2127 } else {
2128 JitOp::Fallback
2129 }
2130 }
2131
2132 "is_positive" => {
2138 let name = node.meta().ins.iter().find_map(|slot| match slot {
2139 crate::ast::Slot::Const {
2140 name,
2141 value: crate::ast::ConstValue::Str(v),
2142 } if name == "name" => Some(v),
2143 _ => None,
2144 });
2145 match name {
2146 Some(v) => JitOp::IsPositiveCheck {
2147 name_ptr: v.as_ptr() as u64,
2148 name_len: v.len() as u64,
2149 },
2150 None => JitOp::IsPositiveCheck {
2151 name_ptr: 0,
2152 name_len: 0,
2153 },
2154 }
2155 }
2156 "in_range" => {
2157 if consts.len() >= 2 {
2158 JitOp::InRangeCheck(consts[0], consts[1])
2159 } else {
2160 JitOp::Fallback
2161 }
2162 }
2163 "is_one_of" => {
2164 if consts.is_empty() {
2165 JitOp::Fallback
2166 } else {
2167 let set = node.meta().ins.iter().find_map(|slot| match slot {
2168 crate::ast::Slot::Const {
2169 name,
2170 value: crate::ast::ConstValue::VecU64(v),
2171 } if name == "allowed" => Some(v),
2172 _ => None,
2173 });
2174 let (set_ptr, set_len) = match set {
2175 Some(v) => (v.as_ptr() as u64, v.len() as u64),
2176 None => (0, 0),
2177 };
2178 JitOp::IsOneOfCheck {
2179 allowed: consts,
2180 set_ptr,
2181 set_len,
2182 }
2183 }
2184 }
2185 _ => JitOp::Fallback,
2190 }
2191}
2192
2193#[doc(hidden)]
2201pub fn compile_jit_raw(
2202 coord_count: usize,
2203 total_slots: usize,
2204 steps: Vec<(JitOp, Vec<usize>, Vec<usize>)>,
2205 output_map: HashMap<String, usize>,
2206 nodes: Vec<Box<dyn PolydatNode>>,
2207) -> Result<JitKernelRaw, String> {
2208 let alone = vec![false; steps.len()];
2209 compile_jit_raw_with(
2210 coord_count,
2211 total_slots,
2212 steps,
2213 output_map,
2214 nodes,
2215 crate::compile::externs::Externs::default(),
2216 super::kernels::ScratchPlan::default(),
2217 Vec::new(),
2218 alone,
2219 )
2220}
2221
2222fn pure_units(
2227 steps: &[(JitOp, Vec<usize>, Vec<usize>)],
2228 total_slots: usize,
2229 alone: &[bool],
2230 volatile: &[usize],
2231) -> crate::compile::fusion_units::UnitPlan {
2232 let mut producer = vec![usize::MAX; total_slots + 1];
2233 for (i, (_, _, outs)) in steps.iter().enumerate() {
2234 for &s in outs {
2235 if s < producer.len() {
2236 producer[s] = i;
2237 }
2238 }
2239 }
2240 let preds: Vec<Vec<usize>> = steps
2241 .iter()
2242 .map(|(_, ins, _)| {
2243 let mut p: Vec<usize> = ins
2244 .iter()
2245 .filter_map(|&s| producer.get(s).copied().filter(|&p| p != usize::MAX))
2246 .collect();
2247 p.sort_unstable();
2248 p.dedup();
2249 p
2250 })
2251 .collect();
2252 let fusible: Vec<bool> = (0..steps.len())
2253 .map(|i| !alone.get(i).copied().unwrap_or(false))
2254 .collect();
2255 let mut class = vec![0u64; steps.len()];
2256 for &v in volatile {
2257 if v < class.len() {
2258 class[v] = 1;
2259 }
2260 }
2261 let rank: Vec<usize> = (0..steps.len()).collect();
2262 crate::compile::fusion_units::plan_units(&preds, &fusible, &class, &rank, &|_| false)
2263}
2264
2265fn unit_dependents(
2268 input_dependents: Vec<Vec<usize>>,
2269 plan: &crate::compile::fusion_units::UnitPlan,
2270) -> Vec<Vec<usize>> {
2271 input_dependents
2272 .into_iter()
2273 .map(|steps| {
2274 let mut units: Vec<usize> = steps.iter().map(|&s| plan.unit_of[s]).collect();
2275 units.sort_unstable();
2276 units.dedup();
2277 units
2278 })
2279 .collect()
2280}
2281
2282#[allow(clippy::too_many_arguments)]
2285pub(crate) fn compile_jit_raw_with(
2286 coord_count: usize,
2287 total_slots: usize,
2288 steps: Vec<(JitOp, Vec<usize>, Vec<usize>)>,
2289 output_map: HashMap<String, usize>,
2290 nodes: Vec<Box<dyn PolydatNode>>,
2291 externs: crate::compile::externs::Externs,
2292 scratch: super::kernels::ScratchPlan,
2293 volatile: Vec<usize>,
2294 alone: Vec<bool>,
2295) -> Result<JitKernelRaw, String> {
2296 let plan = pure_units(&steps, total_slots, &alone, &volatile);
2299 let (_, entry, code) = compile_jit_impl(&steps, Some(&plan.units), Some(total_slots))?;
2300 let cones =
2301 super::kernels::ConePlan::new(&steps, total_slots, &plan, output_map.values().copied());
2302 let mut core = JitCore::new(
2303 total_slots,
2304 coord_count,
2305 output_map,
2306 code,
2307 nodes,
2308 scratch,
2309 volatile,
2310 entry,
2311 cones,
2312 );
2313 core.set_externs(externs);
2314 core.engine =
2315 crate::compile::select::Engine::PureNative(crate::compile::select::Provenance::Raw);
2316 Ok(JitKernelRaw { core })
2317}
2318
2319pub(crate) type JitSegmentCode = (NativeFn, super::kernels::JitCode);
2323
2324pub(crate) fn compile_jit_entry(
2328 steps: &[(JitOp, Vec<usize>, Vec<usize>)],
2329 tracker: Option<usize>,
2330) -> Result<JitSegmentCode, String> {
2331 let (raw_fn, _, code) = compile_jit_impl(steps, None, tracker)?;
2332 Ok((raw_fn, code))
2333}
2334
2335#[allow(clippy::too_many_arguments)]
2337pub(crate) fn compile_jit_push_pull(
2338 coord_count: usize,
2339 total_slots: usize,
2340 steps: Vec<(JitOp, Vec<usize>, Vec<usize>)>,
2341 output_map: HashMap<String, usize>,
2342 nodes: Vec<Box<dyn PolydatNode>>,
2343 input_dependents: Vec<Vec<usize>>,
2344 externs: crate::compile::externs::Externs,
2345 scratch: super::kernels::ScratchPlan,
2346 volatile: Vec<usize>,
2347 alone: Vec<bool>,
2348) -> Result<JitKernelPushPull, String> {
2349 let buffer_len = total_slots;
2350 let plan = pure_units(&steps, total_slots, &alone, &volatile);
2351 let (_, entry, code) = compile_jit_impl(&steps, Some(&plan.units), Some(total_slots))?;
2352 let step_outs: Vec<&[usize]> = steps.iter().map(|(_, _, o)| o.as_slice()).collect();
2353 let slot_provenance =
2354 crate::compile::slot_provenance(coord_count, buffer_len, &step_outs, &input_dependents);
2355 let cones =
2356 super::kernels::ConePlan::new(&steps, total_slots, &plan, output_map.values().copied());
2357 let input_dependents = unit_dependents(input_dependents, &plan);
2358 let mut core = JitCore::new(
2359 total_slots,
2360 coord_count,
2361 output_map,
2362 code,
2363 nodes,
2364 scratch,
2365 volatile,
2366 entry,
2367 cones,
2368 );
2369 core.set_externs(externs);
2370 Ok(JitKernelPushPull {
2371 core,
2372 input_dependents,
2373 slot_provenance,
2374 changed_mask: crate::kernel::ProvMask::all_below(coord_count),
2375 force_run: false,
2376 })
2377}
2378
2379pub type NativeFn = unsafe fn(*const u64, *mut u64, *mut crate::ast::ScratchBuf);
2383pub type NativeDispatchFn =
2387 unsafe fn(*const u64, *mut u64, *mut crate::ast::ScratchBuf, *const u32, u64, *mut u8);
2388
2389type JitCompiled = (NativeFn, NativeDispatchFn, super::kernels::JitCode);
2395
2396type JitEntry = (NativeFn, NativeDispatchFn, bool);
2399
2400type JitFunctionSpec<'a> = (
2403 &'a [(JitOp, Vec<usize>, Vec<usize>)],
2404 Option<&'a [Vec<usize>]>,
2405);
2406
2407fn compile_jit_impl(
2411 steps: &[(JitOp, Vec<usize>, Vec<usize>)],
2412 dispatch: Option<&[Vec<usize>]>,
2413 tracker: Option<usize>,
2414) -> Result<JitCompiled, String> {
2415 let (entries, code) = compile_jit_module(&[(steps, dispatch)], tracker)?;
2416 let (straight_fn, dispatch_fn, _) = entries[0];
2417 Ok((straight_fn, dispatch_fn, code))
2418}
2419
2420pub(crate) type JitStep = (JitOp, Vec<usize>, Vec<usize>);
2422
2423pub(crate) fn compile_jit_entries(
2427 batches: &[&[JitStep]],
2428 tracker: Option<usize>,
2429) -> Result<(Vec<(NativeFn, bool)>, super::kernels::JitCode), String> {
2430 let specs: Vec<JitFunctionSpec> = batches.iter().map(|&b| (b, None)).collect();
2431 let (entries, code) = compile_jit_module(&specs, tracker)?;
2432 Ok((
2433 entries
2434 .into_iter()
2435 .map(|(f, _, fallible)| (f, fallible))
2436 .collect(),
2437 code,
2438 ))
2439}
2440
2441fn compile_jit_module(
2443 functions: &[JitFunctionSpec],
2444 tracker: Option<usize>,
2445) -> Result<(Vec<JitEntry>, super::kernels::JitCode), String> {
2446 let mut flag_builder = settings::builder();
2447 flag_builder.set("opt_level", "speed").unwrap();
2448 flag_builder.set("unwind_info", "true").unwrap();
2453 flag_builder.set("preserve_frame_pointers", "true").unwrap();
2454 let isa = super::host_isa::build_host_isa(flag_builder)?;
2455
2456 let mut jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
2457
2458 jit_builder.symbol("jit_xxh3_hash", jit_xxh3_hash as *const u8);
2460 jit_builder.symbol("jit_interleave", jit_interleave as *const u8);
2461 jit_builder.symbol("jit_shuffle", jit_shuffle as *const u8);
2462 jit_builder.symbol("jit_lut_sample", jit_lut_sample as *const u8);
2463 jit_builder.symbol("jit_weighted_pick", jit_weighted_pick as *const u8);
2464 jit_builder.symbol("jit_pcg", jit_pcg as *const u8);
2465 jit_builder.symbol("jit_pcg_stream", jit_pcg_stream as *const u8);
2466 jit_builder.symbol("jit_n_of", jit_n_of as *const u8);
2467 jit_builder.symbol("jit_cycle_walk", jit_cycle_walk as *const u8);
2468 jit_builder.symbol("jit_perlin_1d", jit_perlin_1d as *const u8);
2469 jit_builder.symbol("jit_perlin_2d", jit_perlin_2d as *const u8);
2470 jit_builder.symbol("jit_simplex_2d", jit_simplex_2d as *const u8);
2471 jit_builder.symbol("jit_fractal_noise_1d", jit_fractal_noise_1d as *const u8);
2472 jit_builder.symbol("jit_fractal_noise_2d", jit_fractal_noise_2d as *const u8);
2473 jit_builder.symbol("jit_thread_id", jit_thread_id as *const u8);
2474 jit_builder.symbol(
2475 "jit_current_epoch_millis",
2476 jit_current_epoch_millis as *const u8,
2477 );
2478 jit_builder.symbol("jit_is_positive_fail", jit_is_positive_fail as *const u8);
2482 jit_builder.symbol("jit_in_range_fail", jit_in_range_fail as *const u8);
2483 jit_builder.symbol("jit_is_one_of_fail", jit_is_one_of_fail as *const u8);
2484 jit_builder.symbol("jit_slot_call", jit_slot_call as *const u8);
2487 jit_builder.symbol("jit_u64_to_str", jit_u64_to_str as *const u8);
2488 jit_builder.symbol("jit_i64_to_str", jit_i64_to_str as *const u8);
2489 jit_builder.symbol("jit_f64_to_str", jit_f64_to_str as *const u8);
2490 jit_builder.symbol("jit_str_concat", jit_str_concat as *const u8);
2491 jit_builder.symbol("jit_json_to_str", jit_json_to_str as *const u8);
2492 jit_builder.symbol("jit_vec_add", jit_vec_add as *const u8);
2493 jit_builder.symbol("jit_vec_scale", jit_vec_scale as *const u8);
2494 jit_builder.symbol("jit_vec_norm", jit_vec_norm as *const u8);
2495 jit_builder.symbol("jit_hash_vec", jit_hash_vec as *const u8);
2496 jit_builder.symbol("jit_xxhash3_vec", jit_xxhash3_vec as *const u8);
2497 jit_builder.symbol("jit_reg_to_vec_f32", jit_reg_to_vec_f32 as *const u8);
2498 jit_builder.symbol("jit_vec_dot", jit_vec_dot as *const u8);
2499 jit_builder.symbol("jit_vec_l2", jit_vec_l2 as *const u8);
2500 jit_builder.symbol("jit_vec_cosine", jit_vec_cosine as *const u8);
2501 jit_builder.symbol("jit_lid_mle", jit_lid_mle as *const u8);
2502 jit_builder.symbol("jit_reg_lane_f32", jit_reg_lane_f32 as *const u8);
2503 jit_builder.symbol("jit_reg_lane_i16", jit_reg_lane_i16 as *const u8);
2504 jit_builder.symbol("jit_reg_lane_i64", jit_reg_lane_i64 as *const u8);
2505 jit_builder.symbol("jit_reg_with_lane_f32", jit_reg_with_lane_f32 as *const u8);
2506 jit_builder.symbol("jit_reg_gather_f32", jit_reg_gather_f32 as *const u8);
2507 jit_builder.symbol("jit_vec_to_reg_f32", jit_vec_to_reg_f32 as *const u8);
2508 jit_builder.symbol("jit_reg_mul_i8", jit_reg_mul_i8 as *const u8);
2509 jit_builder.symbol("jit_sin", jit_sin as *const u8);
2511 jit_builder.symbol("jit_cos", jit_cos as *const u8);
2512 jit_builder.symbol("jit_tan", jit_tan as *const u8);
2513 jit_builder.symbol("jit_asin", jit_asin as *const u8);
2514 jit_builder.symbol("jit_acos", jit_acos as *const u8);
2515 jit_builder.symbol("jit_atan", jit_atan as *const u8);
2516 jit_builder.symbol("jit_sqrt", jit_sqrt as *const u8);
2517 jit_builder.symbol("jit_abs_f64", jit_abs_f64 as *const u8);
2518 jit_builder.symbol("jit_ln", jit_ln as *const u8);
2519 jit_builder.symbol("jit_exp", jit_exp as *const u8);
2520 jit_builder.symbol("jit_floor_base10", jit_floor_base10 as *const u8);
2521 jit_builder.symbol("jit_ceiling_base10", jit_ceiling_base10 as *const u8);
2522 jit_builder.symbol("jit_closest_base10", jit_closest_base10 as *const u8);
2523 jit_builder.symbol("jit_floor_decade", jit_floor_decade as *const u8);
2524 jit_builder.symbol("jit_ceiling_decade", jit_ceiling_decade as *const u8);
2525 jit_builder.symbol("jit_closest_decade", jit_closest_decade as *const u8);
2526 jit_builder.symbol("jit_floor_binomial", jit_floor_binomial as *const u8);
2527 jit_builder.symbol("jit_ceiling_binomial", jit_ceiling_binomial as *const u8);
2528 jit_builder.symbol("jit_closest_binomial", jit_closest_binomial as *const u8);
2529 jit_builder.symbol("jit_floor_fibonacci", jit_floor_fibonacci as *const u8);
2530 jit_builder.symbol("jit_ceiling_fibonacci", jit_ceiling_fibonacci as *const u8);
2531 jit_builder.symbol("jit_closest_fibonacci", jit_closest_fibonacci as *const u8);
2532 jit_builder.symbol("jit_atan2", jit_atan2 as *const u8);
2533 jit_builder.symbol("jit_pow", jit_pow as *const u8);
2534 jit_builder.symbol("jit_round_nearest", jit_round_nearest as *const u8);
2535 jit_builder.symbol("jit_round_floor", jit_round_floor as *const u8);
2536 jit_builder.symbol("jit_round_ceiling", jit_round_ceiling as *const u8);
2537 jit_builder.symbol("jit_f64_mod", jit_f64_mod as *const u8);
2538 jit_builder.symbol("jit_div_zero_fail", jit_div_zero_fail as *const u8);
2539
2540 let mut module = JITModule::new(jit_builder);
2541
2542 let hash_func_id = {
2544 let mut sig = module.make_signature();
2545 sig.params.push(AbiParam::new(types::I64));
2546 sig.returns.push(AbiParam::new(types::I64));
2547 module
2548 .declare_function("jit_xxh3_hash", Linkage::Import, &sig)
2549 .map_err(|e| format!("declare hash: {e}"))?
2550 };
2551
2552 let interleave_func_id = {
2554 let mut sig = module.make_signature();
2555 sig.params.push(AbiParam::new(types::I64));
2556 sig.params.push(AbiParam::new(types::I64));
2557 sig.returns.push(AbiParam::new(types::I64));
2558 module
2559 .declare_function("jit_interleave", Linkage::Import, &sig)
2560 .map_err(|e| format!("declare interleave: {e}"))?
2561 };
2562
2563 let shuffle_func_id = {
2565 let mut sig = module.make_signature();
2566 for _ in 0..4 {
2567 sig.params.push(AbiParam::new(types::I64));
2568 }
2569 sig.returns.push(AbiParam::new(types::I64));
2570 module
2571 .declare_function("jit_shuffle", Linkage::Import, &sig)
2572 .map_err(|e| format!("declare shuffle: {e}"))?
2573 };
2574
2575 let lut_sample_func_id = {
2577 let mut sig = module.make_signature();
2578 for _ in 0..3 {
2579 sig.params.push(AbiParam::new(types::I64));
2580 }
2581 sig.returns.push(AbiParam::new(types::I64));
2582 module
2583 .declare_function("jit_lut_sample", Linkage::Import, &sig)
2584 .map_err(|e| format!("declare lut_sample: {e}"))?
2585 };
2586
2587 let weighted_pick_func_id = {
2589 let mut sig = module.make_signature();
2590 for _ in 0..6 {
2591 sig.params.push(AbiParam::new(types::I64));
2592 }
2593 sig.returns.push(AbiParam::new(types::I64));
2594 module
2595 .declare_function("jit_weighted_pick", Linkage::Import, &sig)
2596 .map_err(|e| format!("declare weighted_pick: {e}"))?
2597 };
2598
2599 let pcg_func_id = {
2600 let mut sig = module.make_signature();
2601 for _ in 0..3 {
2602 sig.params.push(AbiParam::new(types::I64));
2603 }
2604 sig.returns.push(AbiParam::new(types::I64));
2605 module
2606 .declare_function("jit_pcg", Linkage::Import, &sig)
2607 .map_err(|e| format!("declare pcg: {e}"))?
2608 };
2609 let pcg_stream_func_id = {
2610 let mut sig = module.make_signature();
2611 for _ in 0..3 {
2612 sig.params.push(AbiParam::new(types::I64));
2613 }
2614 sig.returns.push(AbiParam::new(types::I64));
2615 module
2616 .declare_function("jit_pcg_stream", Linkage::Import, &sig)
2617 .map_err(|e| format!("declare pcg_stream: {e}"))?
2618 };
2619 let n_of_func_id = {
2620 let mut sig = module.make_signature();
2621 for _ in 0..3 {
2622 sig.params.push(AbiParam::new(types::I64));
2623 }
2624 sig.returns.push(AbiParam::new(types::I64));
2625 module
2626 .declare_function("jit_n_of", Linkage::Import, &sig)
2627 .map_err(|e| format!("declare n_of: {e}"))?
2628 };
2629 let cycle_walk_func_id = {
2630 let mut sig = module.make_signature();
2631 for _ in 0..4 {
2632 sig.params.push(AbiParam::new(types::I64));
2633 }
2634 sig.returns.push(AbiParam::new(types::I64));
2635 module
2636 .declare_function("jit_cycle_walk", Linkage::Import, &sig)
2637 .map_err(|e| format!("declare cycle_walk: {e}"))?
2638 };
2639 let perlin_1d_func_id = {
2640 let mut sig = module.make_signature();
2641 for _ in 0..3 {
2642 sig.params.push(AbiParam::new(types::I64));
2643 }
2644 sig.returns.push(AbiParam::new(types::I64));
2645 module
2646 .declare_function("jit_perlin_1d", Linkage::Import, &sig)
2647 .map_err(|e| format!("declare perlin_1d: {e}"))?
2648 };
2649 let perlin_2d_func_id = {
2650 let mut sig = module.make_signature();
2651 for _ in 0..4 {
2652 sig.params.push(AbiParam::new(types::I64));
2653 }
2654 sig.returns.push(AbiParam::new(types::I64));
2655 module
2656 .declare_function("jit_perlin_2d", Linkage::Import, &sig)
2657 .map_err(|e| format!("declare perlin_2d: {e}"))?
2658 };
2659 let simplex_2d_func_id = {
2660 let mut sig = module.make_signature();
2661 for _ in 0..4 {
2662 sig.params.push(AbiParam::new(types::I64));
2663 }
2664 sig.returns.push(AbiParam::new(types::I64));
2665 module
2666 .declare_function("jit_simplex_2d", Linkage::Import, &sig)
2667 .map_err(|e| format!("declare simplex_2d: {e}"))?
2668 };
2669 let fractal_noise_1d_func_id = {
2670 let mut sig = module.make_signature();
2671 for _ in 0..4 {
2672 sig.params.push(AbiParam::new(types::I64));
2673 }
2674 sig.returns.push(AbiParam::new(types::I64));
2675 module
2676 .declare_function("jit_fractal_noise_1d", Linkage::Import, &sig)
2677 .map_err(|e| format!("declare fractal_noise_1d: {e}"))?
2678 };
2679 let fractal_noise_2d_func_id = {
2680 let mut sig = module.make_signature();
2681 for _ in 0..5 {
2682 sig.params.push(AbiParam::new(types::I64));
2683 }
2684 sig.returns.push(AbiParam::new(types::I64));
2685 module
2686 .declare_function("jit_fractal_noise_2d", Linkage::Import, &sig)
2687 .map_err(|e| format!("declare fractal_noise_2d: {e}"))?
2688 };
2689 let thread_id_func_id = {
2690 let mut sig = module.make_signature();
2691 sig.returns.push(AbiParam::new(types::I64));
2692 module
2693 .declare_function("jit_thread_id", Linkage::Import, &sig)
2694 .map_err(|e| format!("declare thread_id: {e}"))?
2695 };
2696 let current_epoch_millis_func_id = {
2697 let mut sig = module.make_signature();
2698 sig.returns.push(AbiParam::new(types::I64));
2699 module
2700 .declare_function("jit_current_epoch_millis", Linkage::Import, &sig)
2701 .map_err(|e| format!("declare current_epoch_millis: {e}"))?
2702 };
2703
2704 let math_unary_names = [
2706 "jit_sin",
2707 "jit_cos",
2708 "jit_tan",
2709 "jit_asin",
2710 "jit_acos",
2711 "jit_atan",
2712 "jit_sqrt",
2713 "jit_abs_f64",
2714 "jit_ln",
2715 "jit_exp",
2716 "jit_floor_base10",
2717 "jit_ceiling_base10",
2718 "jit_closest_base10",
2719 "jit_floor_decade",
2720 "jit_ceiling_decade",
2721 "jit_closest_decade",
2722 "jit_floor_binomial",
2723 "jit_ceiling_binomial",
2724 "jit_closest_binomial",
2725 "jit_floor_fibonacci",
2726 "jit_ceiling_fibonacci",
2727 "jit_closest_fibonacci",
2728 ];
2729 let mut math_unary_ids = Vec::new();
2730 for name in &math_unary_names {
2731 let mut sig = module.make_signature();
2732 sig.params.push(AbiParam::new(types::I64));
2733 sig.returns.push(AbiParam::new(types::I64));
2734 math_unary_ids.push(
2735 module
2736 .declare_function(name, Linkage::Import, &sig)
2737 .map_err(|e| format!("declare {name}: {e}"))?,
2738 );
2739 }
2740
2741 let is_positive_fail_id = {
2745 let mut sig = module.make_signature();
2746 sig.params.push(AbiParam::new(types::I64));
2747 sig.params.push(AbiParam::new(types::I64));
2748 sig.params.push(AbiParam::new(types::I64));
2749 sig.returns.push(AbiParam::new(types::I64));
2750 module
2751 .declare_function("jit_is_positive_fail", Linkage::Import, &sig)
2752 .map_err(|e| format!("declare is_positive_fail: {e}"))?
2753 };
2754
2755 let in_range_fail_id = {
2757 let mut sig = module.make_signature();
2758 for _ in 0..3 {
2759 sig.params.push(AbiParam::new(types::I64));
2760 }
2761 sig.returns.push(AbiParam::new(types::I64));
2762 module
2763 .declare_function("jit_in_range_fail", Linkage::Import, &sig)
2764 .map_err(|e| format!("declare in_range_fail: {e}"))?
2765 };
2766
2767 let is_one_of_fail_id = {
2770 let mut sig = module.make_signature();
2771 sig.params.push(AbiParam::new(types::I64));
2772 sig.params.push(AbiParam::new(types::I64));
2773 sig.params.push(AbiParam::new(types::I64));
2774 sig.returns.push(AbiParam::new(types::I64));
2775 module
2776 .declare_function("jit_is_one_of_fail", Linkage::Import, &sig)
2777 .map_err(|e| format!("declare is_one_of_fail: {e}"))?
2778 };
2779
2780 let math_binary_names = [
2782 "jit_atan2",
2783 "jit_pow",
2784 "jit_round_nearest",
2785 "jit_round_floor",
2786 "jit_round_ceiling",
2787 "jit_f64_mod",
2788 ];
2789 const F64_MOD_HELPER: usize = 5;
2790
2791 let div_zero_fail_id = {
2793 let mut sig = module.make_signature();
2794 sig.params.push(AbiParam::new(types::I64));
2795 sig.returns.push(AbiParam::new(types::I64));
2796 module
2797 .declare_function("jit_div_zero_fail", Linkage::Import, &sig)
2798 .map_err(|e| format!("declare div_zero_fail: {e}"))?
2799 };
2800 let mut math_binary_ids = Vec::new();
2801 for name in &math_binary_names {
2802 let mut sig = module.make_signature();
2803 sig.params.push(AbiParam::new(types::I64));
2804 sig.params.push(AbiParam::new(types::I64));
2805 sig.returns.push(AbiParam::new(types::I64));
2806 math_binary_ids.push(
2807 module
2808 .declare_function(name, Linkage::Import, &sig)
2809 .map_err(|e| format!("declare {name}: {e}"))?,
2810 );
2811 }
2812
2813 let slot_call_id = {
2816 let mut sig = module.make_signature();
2817 for _ in 0..8 {
2818 sig.params.push(AbiParam::new(types::I64));
2819 }
2820 module
2821 .declare_function("jit_slot_call", Linkage::Import, &sig)
2822 .map_err(|e| format!("declare jit_slot_call: {e}"))?
2823 };
2824
2825 let mut declare_str = |name: &str, args: usize| -> Result<cranelift_module::FuncId, String> {
2829 let mut sig = module.make_signature();
2830 for _ in 0..args {
2831 sig.params.push(AbiParam::new(types::I64));
2832 }
2833 module
2834 .declare_function(name, Linkage::Import, &sig)
2835 .map_err(|e| format!("declare {name}: {e}"))
2836 };
2837 let u64_to_str_id = declare_str("jit_u64_to_str", 5)?;
2838 let i64_to_str_id = declare_str("jit_i64_to_str", 5)?;
2839 let f64_to_str_id = declare_str("jit_f64_to_str", 5)?;
2840 let str_concat_id = declare_str("jit_str_concat", 6)?;
2841 let json_to_str_id = declare_str("jit_json_to_str", 6)?;
2842
2843 let mut declare_words =
2848 |name: &str, args: usize, returns: bool| -> Result<cranelift_module::FuncId, String> {
2849 let mut sig = module.make_signature();
2850 for _ in 0..args {
2851 sig.params.push(AbiParam::new(types::I64));
2852 }
2853 if returns {
2854 sig.returns.push(AbiParam::new(types::I64));
2855 }
2856 module
2857 .declare_function(name, Linkage::Import, &sig)
2858 .map_err(|e| format!("declare {name}: {e}"))
2859 };
2860 let vec_producer_ids = [
2861 (VecProducer::Add, declare_words("jit_vec_add", 8, false)?),
2862 (
2863 VecProducer::Scale,
2864 declare_words("jit_vec_scale", 8, false)?,
2865 ),
2866 (VecProducer::Norm, declare_words("jit_vec_norm", 8, false)?),
2867 (
2868 VecProducer::HashVec,
2869 declare_words("jit_hash_vec", 8, false)?,
2870 ),
2871 (
2872 VecProducer::XxHash3Vec,
2873 declare_words("jit_xxhash3_vec", 8, false)?,
2874 ),
2875 (
2876 VecProducer::RegToVec,
2877 declare_words("jit_reg_to_vec_f32", 8, false)?,
2878 ),
2879 ];
2880 let vec_reducer_ids = [
2881 (VecReducer::Dot, declare_words("jit_vec_dot", 4, true)?),
2882 (VecReducer::L2, declare_words("jit_vec_l2", 4, true)?),
2883 (
2884 VecReducer::Cosine,
2885 declare_words("jit_vec_cosine", 4, true)?,
2886 ),
2887 (VecReducer::LidMle, declare_words("jit_lid_mle", 4, true)?),
2888 ];
2889 let reg_lane_ids = [
2890 (
2891 RegLaneRead::F32,
2892 declare_words("jit_reg_lane_f32", 3, true)?,
2893 ),
2894 (
2895 RegLaneRead::I16,
2896 declare_words("jit_reg_lane_i16", 3, true)?,
2897 ),
2898 (
2899 RegLaneRead::I64,
2900 declare_words("jit_reg_lane_i64", 3, true)?,
2901 ),
2902 ];
2903 let reg_producer_ids = [
2904 (
2905 RegProducer::WithLaneF32,
2906 declare_words("jit_reg_with_lane_f32", 6, false)?,
2907 ),
2908 (
2909 RegProducer::GatherF32,
2910 declare_words("jit_reg_gather_f32", 6, false)?,
2911 ),
2912 (
2913 RegProducer::VecToRegF32,
2914 declare_words("jit_vec_to_reg_f32", 6, false)?,
2915 ),
2916 (
2917 RegProducer::MulI8,
2918 declare_words("jit_reg_mul_i8", 6, false)?,
2919 ),
2920 ];
2921
2922 let mut defined: Vec<(cranelift_module::FuncId, bool)> = Vec::with_capacity(functions.len());
2927 for (function_idx, &(steps, dispatch)) in functions.iter().enumerate() {
2928 let mut sig = module.make_signature();
2932 sig.params.push(AbiParam::new(types::I64)); sig.params.push(AbiParam::new(types::I64)); sig.params.push(AbiParam::new(types::I64)); if dispatch.is_some() {
2936 sig.params.push(AbiParam::new(types::I64)); sig.params.push(AbiParam::new(types::I64)); sig.params.push(AbiParam::new(types::I64)); }
2940 let func_id = module
2941 .declare_function(
2942 &format!("polydat_kernel_{function_idx}"),
2943 Linkage::Local,
2944 &sig,
2945 )
2946 .map_err(|e| format!("declare kernel: {e}"))?;
2947
2948 let mut ctx = module.make_context();
2949 ctx.func.signature = sig;
2950
2951 let mut fb_ctx = FunctionBuilderContext::new();
2952 {
2953 let mut builder = FunctionBuilder::new(&mut ctx.func, &mut fb_ctx);
2954 let block = builder.create_block();
2955 builder.append_block_params_for_function_params(block);
2956 builder.switch_to_block(block);
2957 builder.seal_block(block);
2958
2959 let _coords_ptr = builder.block_params(block)[0];
2960 let buffer_ptr = builder.block_params(block)[1];
2961 let scratch_ptr = builder.block_params(block)[2];
2962
2963 let hash_func_ref = module.declare_func_in_func(hash_func_id, builder.func);
2965 let interleave_func_ref = module.declare_func_in_func(interleave_func_id, builder.func);
2966 let shuffle_func_ref = module.declare_func_in_func(shuffle_func_id, builder.func);
2967 let lut_sample_func_ref = module.declare_func_in_func(lut_sample_func_id, builder.func);
2968 let weighted_pick_func_ref =
2969 module.declare_func_in_func(weighted_pick_func_id, builder.func);
2970 let is_positive_fail_ref =
2971 module.declare_func_in_func(is_positive_fail_id, builder.func);
2972 let in_range_fail_ref = module.declare_func_in_func(in_range_fail_id, builder.func);
2973 let div_zero_fail_ref = module.declare_func_in_func(div_zero_fail_id, builder.func);
2974 let is_one_of_fail_ref = module.declare_func_in_func(is_one_of_fail_id, builder.func);
2975 let slot_call_ref = module.declare_func_in_func(slot_call_id, builder.func);
2976 let u64_to_str_ref = module.declare_func_in_func(u64_to_str_id, builder.func);
2977 let i64_to_str_ref = module.declare_func_in_func(i64_to_str_id, builder.func);
2978 let f64_to_str_ref = module.declare_func_in_func(f64_to_str_id, builder.func);
2979 let str_concat_ref = module.declare_func_in_func(str_concat_id, builder.func);
2980 let json_to_str_ref = module.declare_func_in_func(json_to_str_id, builder.func);
2981 let vec_producer_refs: Vec<(VecProducer, ir::FuncRef)> = vec_producer_ids
2982 .iter()
2983 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
2984 .collect();
2985 let vec_reducer_refs: Vec<(VecReducer, ir::FuncRef)> = vec_reducer_ids
2986 .iter()
2987 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
2988 .collect();
2989 let reg_lane_refs: Vec<(RegLaneRead, ir::FuncRef)> = reg_lane_ids
2990 .iter()
2991 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
2992 .collect();
2993 let reg_producer_refs: Vec<(RegProducer, ir::FuncRef)> = reg_producer_ids
2994 .iter()
2995 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
2996 .collect();
2997 let pcg_func_ref = module.declare_func_in_func(pcg_func_id, builder.func);
2998 let pcg_stream_func_ref = module.declare_func_in_func(pcg_stream_func_id, builder.func);
2999 let n_of_func_ref = module.declare_func_in_func(n_of_func_id, builder.func);
3000 let cycle_walk_func_ref = module.declare_func_in_func(cycle_walk_func_id, builder.func);
3001 let perlin_1d_func_ref = module.declare_func_in_func(perlin_1d_func_id, builder.func);
3002 let perlin_2d_func_ref = module.declare_func_in_func(perlin_2d_func_id, builder.func);
3003 let simplex_2d_func_ref = module.declare_func_in_func(simplex_2d_func_id, builder.func);
3004 let fractal_noise_1d_func_ref =
3005 module.declare_func_in_func(fractal_noise_1d_func_id, builder.func);
3006 let fractal_noise_2d_func_ref =
3007 module.declare_func_in_func(fractal_noise_2d_func_id, builder.func);
3008 let thread_id_func_ref = module.declare_func_in_func(thread_id_func_id, builder.func);
3009 let current_epoch_millis_func_ref =
3010 module.declare_func_in_func(current_epoch_millis_func_id, builder.func);
3011 let math_unary_refs: Vec<_> = math_unary_ids
3012 .iter()
3013 .map(|id| module.declare_func_in_func(*id, builder.func))
3014 .collect();
3015 let math_binary_refs: Vec<_> = math_binary_ids
3016 .iter()
3017 .map(|id| module.declare_func_in_func(*id, builder.func))
3018 .collect();
3019 let everything: [Vec<usize>; 1] = [(0..steps.len()).collect()];
3027 let schedule: &[Vec<usize>] = dispatch.unwrap_or(&everything);
3028 let dispatcher = dispatch.map(|units| {
3029 let list_ptr = builder.block_params(block)[3];
3030 let list_len = builder.block_params(block)[4];
3031 let clean_ptr = builder.block_params(block)[5];
3032 let at = builder.create_sized_stack_slot(ir::StackSlotData::new(
3033 ir::StackSlotKind::ExplicitSlot,
3034 8,
3035 3,
3036 ));
3037 let zero = builder.ins().iconst(types::I64, 0);
3038 builder.ins().stack_store(zero, at, 0);
3039 let head = builder.create_block();
3040 let fetch = builder.create_block();
3041 let dispatch_unit = builder.create_block();
3042 let skip = builder.create_block();
3043 let exit = builder.create_block();
3044 let unit_blocks: Vec<ir::Block> =
3045 units.iter().map(|_| builder.create_block()).collect();
3046 builder.ins().jump(head, &[]);
3047 builder.switch_to_block(head);
3049 let i = builder.ins().stack_load(types::I64, at, 0);
3050 let done = builder.ins().icmp(
3051 ir::condcodes::IntCC::UnsignedGreaterThanOrEqual,
3052 i,
3053 list_len,
3054 );
3055 builder.ins().brif(done, exit, &[], fetch, &[]);
3056 builder.switch_to_block(fetch);
3059 builder.seal_block(fetch);
3060 let offset = builder.ins().ishl_imm(i, 2);
3061 let addr = builder.ins().iadd(list_ptr, offset);
3062 let unit = builder
3063 .ins()
3064 .load(types::I32, ir::MemFlags::trusted(), addr, 0);
3065 let unit_wide = builder.ins().uextend(types::I64, unit);
3066 let flag_addr = builder.ins().iadd(clean_ptr, unit_wide);
3067 let flag = builder
3068 .ins()
3069 .load(types::I8, ir::MemFlags::trusted(), flag_addr, 0);
3070 builder.ins().brif(flag, skip, &[], dispatch_unit, &[]);
3071 builder.switch_to_block(skip);
3073 builder.seal_block(skip);
3074 let next = builder.ins().iadd_imm(i, 1);
3075 builder.ins().stack_store(next, at, 0);
3076 builder.ins().jump(head, &[]);
3077 builder.switch_to_block(dispatch_unit);
3079 builder.seal_block(dispatch_unit);
3080 let default = builder.func.dfg.block_call(exit, &[]);
3081 let targets: Vec<ir::BlockCall> = unit_blocks
3082 .iter()
3083 .map(|&b| builder.func.dfg.block_call(b, &[]))
3084 .collect();
3085 let table = builder.create_jump_table(ir::JumpTableData::new(default, &targets));
3086 builder.ins().br_table(unit, table);
3087 (at, head, exit, unit_blocks, clean_ptr)
3088 });
3089
3090 for (unit_idx, members) in schedule.iter().enumerate() {
3092 if let Some((_, _, _, unit_blocks, _)) = &dispatcher {
3093 builder.switch_to_block(unit_blocks[unit_idx]);
3094 builder.seal_block(unit_blocks[unit_idx]);
3095 }
3096 for &step_idx in members {
3097 let (jit_op, input_slots, output_slots) = &steps[step_idx];
3098 let tracker_store = tracker.map(|t| {
3102 let idx = builder.ins().iconst(types::I64, step_idx as i64);
3103 let inst = store_slot(&mut builder, buffer_ptr, t, idx);
3104 (inst, builder.func.dfg.num_insts())
3105 });
3106 match jit_op {
3107 JitOp::Identity => {
3108 for (&i, &o) in input_slots.iter().zip(output_slots.iter()) {
3111 let val = load_slot(&mut builder, buffer_ptr, i);
3112 store_slot(&mut builder, buffer_ptr, o, val);
3113 }
3114 }
3115 JitOp::AddConst(c) => {
3116 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3117 let c_val = builder.ins().iconst(types::I64, *c as i64);
3118 let result = builder.ins().iadd(val, c_val);
3119 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3120 }
3121 JitOp::MulConst(c) => {
3122 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3123 let c_val = builder.ins().iconst(types::I64, *c as i64);
3124 let result = builder.ins().imul(val, c_val);
3125 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3126 }
3127 JitOp::DivConst(c) | JitOp::ModConst(c) => {
3128 let is_div = matches!(jit_op, JitOp::DivConst(_));
3131 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3132 if *c == 0 {
3133 let kind =
3134 builder.ins().iconst(types::I64, if is_div { 0 } else { 1 });
3135 let _ = builder.ins().call(div_zero_fail_ref, &[kind]);
3136 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3137 } else {
3138 let c_val = builder.ins().iconst(types::I64, *c as i64);
3139 let result = if is_div {
3140 builder.ins().udiv(val, c_val)
3141 } else {
3142 builder.ins().urem(val, c_val)
3143 };
3144 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3145 }
3146 }
3147 JitOp::U64DivWire | JitOp::U64ModWire => {
3148 let is_div = matches!(jit_op, JitOp::U64DivWire);
3152 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3153 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3154 let zero = builder.ins().iconst(types::I64, 0);
3155 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, b, zero);
3156 let fail_block = builder.create_block();
3157 let ok_block = builder.create_block();
3158 builder.ins().brif(is_zero, fail_block, &[], ok_block, &[]);
3159 builder.switch_to_block(fail_block);
3160 builder.seal_block(fail_block);
3161 let kind = builder.ins().iconst(types::I64, if is_div { 0 } else { 1 });
3162 let _ = builder.ins().call(div_zero_fail_ref, &[kind]);
3163 builder.ins().jump(ok_block, &[]);
3164 builder.switch_to_block(ok_block);
3165 builder.seal_block(ok_block);
3166 let result = if is_div {
3167 builder.ins().udiv(a, b)
3168 } else {
3169 builder.ins().urem(a, b)
3170 };
3171 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3172 }
3173 JitOp::ClampConst(min, max) => {
3174 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3175 let min_val = builder.ins().iconst(types::I64, *min as i64);
3176 let max_val = builder.ins().iconst(types::I64, *max as i64);
3177 let clamped_lo = builder.ins().umax(val, min_val);
3178 let clamped = builder.ins().umin(clamped_lo, max_val);
3179 store_slot(&mut builder, buffer_ptr, output_slots[0], clamped);
3180 }
3181 JitOp::Interleave => {
3182 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3183 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3184 let call = builder.ins().call(interleave_func_ref, &[a, b]);
3185 let result = builder.inst_results(call)[0];
3186 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3187 }
3188 JitOp::MixedRadixConst(radixes) => {
3189 let mut remainder = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3191 for (i, &radix) in radixes.iter().enumerate() {
3192 if radix == 0 {
3193 store_slot(
3195 &mut builder,
3196 buffer_ptr,
3197 output_slots[i],
3198 remainder,
3199 );
3200 } else {
3201 let r = builder.ins().iconst(types::I64, radix as i64);
3202 let digit = builder.ins().urem(remainder, r);
3203 store_slot(&mut builder, buffer_ptr, output_slots[i], digit);
3204 remainder = builder.ins().udiv(remainder, r);
3205 }
3206 }
3207 }
3208 JitOp::Hash => {
3209 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3210 let call = builder.ins().call(hash_func_ref, &[val]);
3211 let result = builder.inst_results(call)[0];
3212 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3213 }
3214 JitOp::SplitMix64 => {
3215 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3216 let c_gamma = builder
3217 .ins()
3218 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3219 let x1 = builder.ins().iadd(x0, c_gamma);
3220 let s30 = builder.ins().ushr_imm(x1, 30);
3221 let x2 = builder.ins().bxor(x1, s30);
3222 let c_m1 = builder
3223 .ins()
3224 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3225 let x3 = builder.ins().imul(x2, c_m1);
3226 let s27 = builder.ins().ushr_imm(x3, 27);
3227 let x4 = builder.ins().bxor(x3, s27);
3228 let c_m2 = builder
3229 .ins()
3230 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3231 let x5 = builder.ins().imul(x4, c_m2);
3232 let s31 = builder.ins().ushr_imm(x5, 31);
3233 let result = builder.ins().bxor(x5, s31);
3234 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3235 }
3236 JitOp::FairCoin => {
3237 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3238 let c_gamma = builder
3239 .ins()
3240 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3241 let x1 = builder.ins().iadd(x0, c_gamma);
3242 let s30 = builder.ins().ushr_imm(x1, 30);
3243 let x2 = builder.ins().bxor(x1, s30);
3244 let c_m1 = builder
3245 .ins()
3246 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3247 let x3 = builder.ins().imul(x2, c_m1);
3248 let s27 = builder.ins().ushr_imm(x3, 27);
3249 let x4 = builder.ins().bxor(x3, s27);
3250 let c_m2 = builder
3251 .ins()
3252 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3253 let x5 = builder.ins().imul(x4, c_m2);
3254 let s31 = builder.ins().ushr_imm(x5, 31);
3255 let h = builder.ins().bxor(x5, s31);
3256 let one = builder.ins().iconst(types::I64, 1);
3257 let result = builder.ins().band(h, one);
3258 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3259 }
3260 JitOp::CoinFlipConst(threshold) => {
3261 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3262 let thr = builder.ins().iconst(types::I64, *threshold as i64);
3263 let cmp =
3264 builder
3265 .ins()
3266 .icmp(ir::condcodes::IntCC::UnsignedLessThan, x, thr);
3267 let zero = builder.ins().iconst(types::I64, 0);
3268 let one = builder.ins().iconst(types::I64, 1);
3269 let result = builder.ins().select(cmp, one, zero);
3270 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3271 }
3272 JitOp::UnfairCoinConst(p_bits) => {
3273 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3274 let c_gamma = builder
3275 .ins()
3276 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3277 let x1 = builder.ins().iadd(x0, c_gamma);
3278 let s30 = builder.ins().ushr_imm(x1, 30);
3279 let x2 = builder.ins().bxor(x1, s30);
3280 let c_m1 = builder
3281 .ins()
3282 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3283 let x3 = builder.ins().imul(x2, c_m1);
3284 let s27 = builder.ins().ushr_imm(x3, 27);
3285 let x4 = builder.ins().bxor(x3, s27);
3286 let c_m2 = builder
3287 .ins()
3288 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3289 let x5 = builder.ins().imul(x4, c_m2);
3290 let s31 = builder.ins().ushr_imm(x5, 31);
3291 let h = builder.ins().bxor(x5, s31);
3292
3293 let fval = builder.ins().fcvt_from_uint(types::F64, h);
3294 let max_f = builder.ins().f64const(u64::MAX as f64);
3295 let unit = builder.ins().fdiv(fval, max_f);
3296 let p_f = builder.ins().f64const(f64::from_bits(*p_bits));
3297 let cmp =
3298 builder
3299 .ins()
3300 .fcmp(ir::condcodes::FloatCC::LessThan, unit, p_f);
3301 let zero = builder.ins().iconst(types::I64, 0);
3302 let one = builder.ins().iconst(types::I64, 1);
3303 let result = builder.ins().select(cmp, one, zero);
3304 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3305 }
3306 JitOp::ChanceConst(p_bits) => {
3307 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3308 let c_gamma = builder
3309 .ins()
3310 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3311 let x1 = builder.ins().iadd(x0, c_gamma);
3312 let s30 = builder.ins().ushr_imm(x1, 30);
3313 let x2 = builder.ins().bxor(x1, s30);
3314 let c_m1 = builder
3315 .ins()
3316 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3317 let x3 = builder.ins().imul(x2, c_m1);
3318 let s27 = builder.ins().ushr_imm(x3, 27);
3319 let x4 = builder.ins().bxor(x3, s27);
3320 let c_m2 = builder
3321 .ins()
3322 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3323 let x5 = builder.ins().imul(x4, c_m2);
3324 let s31 = builder.ins().ushr_imm(x5, 31);
3325 let h = builder.ins().bxor(x5, s31);
3326
3327 let fval = builder.ins().fcvt_from_uint(types::F64, h);
3328 let max_f = builder.ins().f64const(u64::MAX as f64);
3329 let unit = builder.ins().fdiv(fval, max_f);
3330 let p_f = builder.ins().f64const(f64::from_bits(*p_bits));
3331 let cmp =
3332 builder
3333 .ins()
3334 .fcmp(ir::condcodes::FloatCC::LessThan, unit, p_f);
3335 let zero_bits =
3336 builder.ins().iconst(types::I64, 0.0_f64.to_bits() as i64);
3337 let one_bits =
3338 builder.ins().iconst(types::I64, 1.0_f64.to_bits() as i64);
3339 let result = builder.ins().select(cmp, one_bits, zero_bits);
3340 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3341 }
3342 JitOp::ShuffleConst(feedback, size, min) => {
3343 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3344 let fb = builder.ins().iconst(types::I64, *feedback as i64);
3345 let sz = builder.ins().iconst(types::I64, *size as i64);
3346 let mn = builder.ins().iconst(types::I64, *min as i64);
3347 let call = builder.ins().call(shuffle_func_ref, &[val, fb, sz, mn]);
3348 let result = builder.inst_results(call)[0];
3349 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3350 }
3351
3352 JitOp::UnitInterval => {
3354 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3356 let fval = builder.ins().fcvt_from_uint(types::F64, val);
3357 let max_f = builder.ins().f64const(u64::MAX as f64);
3358 let result = builder.ins().fdiv(fval, max_f);
3359 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3360 }
3361 JitOp::F64ToU64 => {
3362 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3363 let result = builder.ins().fcvt_to_uint_sat(types::I64, fval);
3364 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3365 }
3366 JitOp::RoundToU64 => {
3367 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3368 let rounded = round_half_away(&mut builder, fval);
3369 let result = builder.ins().fcvt_to_uint_sat(types::I64, rounded);
3370 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3371 }
3372 JitOp::FloorToU64 => {
3373 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3374 let floored = builder.ins().floor(fval);
3375 let result = builder.ins().fcvt_to_uint_sat(types::I64, floored);
3376 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3377 }
3378 JitOp::CeilToU64 => {
3379 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3380 let ceiled = builder.ins().ceil(fval);
3381 let result = builder.ins().fcvt_to_uint_sat(types::I64, ceiled);
3382 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3383 }
3384 JitOp::ClampF64Const(min_bits, max_bits) => {
3385 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3386 let fmin = builder.ins().f64const(f64::from_bits(*min_bits));
3387 let fmax = builder.ins().f64const(f64::from_bits(*max_bits));
3388 let clamped = clamp_ir(&mut builder, fval, fmin, fmax);
3389 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], clamped);
3390 }
3391 JitOp::LerpConst(a_bits, b_bits) => {
3392 let t = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3394 let a = builder.ins().f64const(f64::from_bits(*a_bits));
3395 let b = builder.ins().f64const(f64::from_bits(*b_bits));
3396 let diff = builder.ins().fsub(b, a);
3397 let scaled = builder.ins().fmul(t, diff);
3398 let result = builder.ins().fadd(a, scaled);
3399 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3400 }
3401 JitOp::ScaleRangeConst(min_bits, range_bits) => {
3402 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3404 let fval = builder.ins().fcvt_from_uint(types::F64, val);
3405 let max_f = builder.ins().f64const(u64::MAX as f64);
3406 let t = builder.ins().fdiv(fval, max_f);
3407 let fmin = builder.ins().f64const(f64::from_bits(*min_bits));
3408 let frange = builder.ins().f64const(f64::from_bits(*range_bits));
3409 let scaled = builder.ins().fmul(t, frange);
3410 let result = builder.ins().fadd(fmin, scaled);
3411 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3412 }
3413 JitOp::QuantizeConst(step_bits) => {
3414 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3416 let step = builder.ins().f64const(f64::from_bits(*step_bits));
3417 let divided = builder.ins().fdiv(fval, step);
3418 let rounded = round_half_away(&mut builder, divided);
3419 let result = builder.ins().fmul(rounded, step);
3420 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3421 }
3422
3423 JitOp::LutSampleConst(lut_ptr, lut_len) => {
3424 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3426 let ptr_val = builder.ins().iconst(types::I64, *lut_ptr as i64);
3427 let len_val = builder.ins().iconst(types::I64, *lut_len as i64);
3428 let call = builder
3429 .ins()
3430 .call(lut_sample_func_ref, &[input, ptr_val, len_val]);
3431 let result = builder.inst_results(call)[0];
3432 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3433 }
3434 JitOp::DiscretizeConst(range_bits, buckets) => {
3435 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3440 let range = f64::from_bits(*range_bits);
3441 let fzero = builder.ins().f64const(0.0);
3442 let frange = builder.ins().f64const(range);
3443 let fbuckets = builder.ins().f64const(*buckets as f64);
3444 let clamped = clamp_ir(&mut builder, fval, fzero, frange);
3445 let divided = builder.ins().fdiv(clamped, frange);
3446 let scaled = builder.ins().fmul(divided, fbuckets);
3447 let as_u64 = builder.ins().fcvt_to_uint_sat(types::I64, scaled);
3448 let max_bucket =
3449 builder.ins().iconst(types::I64, (*buckets - 1) as i64);
3450 let result = builder.ins().umin(as_u64, max_bucket);
3451 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3452 }
3453
3454 JitOp::WeightedPickConst(
3455 values_ptr,
3456 biases_ptr,
3457 primaries_ptr,
3458 aliases_ptr,
3459 n,
3460 ) => {
3461 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3463 let v_ptr = builder.ins().iconst(types::I64, *values_ptr as i64);
3464 let b_ptr = builder.ins().iconst(types::I64, *biases_ptr as i64);
3465 let p_ptr = builder.ins().iconst(types::I64, *primaries_ptr as i64);
3466 let a_ptr = builder.ins().iconst(types::I64, *aliases_ptr as i64);
3467 let n_val = builder.ins().iconst(types::I64, *n as i64);
3468 let call = builder.ins().call(
3469 weighted_pick_func_ref,
3470 &[input, v_ptr, b_ptr, p_ptr, a_ptr, n_val],
3471 );
3472 let result = builder.inst_results(call)[0];
3473 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3474 }
3475
3476 JitOp::MathUnary(idx) => {
3477 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3478 let func_ref = math_unary_refs[*idx as usize];
3479 let call = builder.ins().call(func_ref, &[input]);
3480 let result = builder.inst_results(call)[0];
3481 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3482 }
3483
3484 JitOp::MathBinary(idx) => {
3485 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3486 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3487 let func_ref = math_binary_refs[*idx as usize];
3488 let call = builder.ins().call(func_ref, &[a, b]);
3489 let result = builder.inst_results(call)[0];
3490 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3491 }
3492
3493 JitOp::ToF64 => {
3494 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3495 let fval = builder.ins().fcvt_from_uint(types::F64, val);
3496 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], fval);
3497 }
3498
3499 JitOp::RegBinOp(lane, arith) => {
3501 let vt = reg_lane_type(*lane);
3502 let a = load_reg128(&mut builder, buffer_ptr, input_slots[0], vt);
3503 let b = load_reg128(&mut builder, buffer_ptr, input_slots[2], vt);
3504 let is_float = matches!(*lane, 4 | 5);
3505 let r = match (arith, is_float) {
3506 (0, false) => builder.ins().iadd(a, b),
3507 (1, false) => builder.ins().isub(a, b),
3508 (2, false) => builder.ins().imul(a, b),
3509 (0, true) => builder.ins().fadd(a, b),
3510 (1, true) => builder.ins().fsub(a, b),
3511 (2, true) => builder.ins().fmul(a, b),
3512 _ => unreachable!("RegBinOp arith index out of range"),
3513 };
3514 store_reg128(&mut builder, buffer_ptr, output_slots[0], r);
3515 }
3516 JitOp::RegCopy => {
3517 let v =
3518 load_reg128(&mut builder, buffer_ptr, input_slots[0], types::I64X2);
3519 store_reg128(&mut builder, buffer_ptr, output_slots[0], v);
3520 }
3521 JitOp::RegSplat(lane) => {
3522 let vt = reg_lane_type(*lane);
3523 let scalar = match *lane {
3524 0 => {
3526 let v = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3527 builder.ins().ireduce(types::I8, v)
3528 }
3529 1 => {
3530 let v = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3531 builder.ins().ireduce(types::I16, v)
3532 }
3533 2 => {
3534 let v = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3535 builder.ins().ireduce(types::I32, v)
3536 }
3537 3 => load_slot(&mut builder, buffer_ptr, input_slots[0]),
3538 4 => {
3540 let f = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3541 builder.ins().fdemote(types::F32, f)
3542 }
3543 5 => load_slot_f64(&mut builder, buffer_ptr, input_slots[0]),
3544 _ => unreachable!("RegSplat lane index out of range"),
3545 };
3546 let v = builder.ins().splat(vt, scalar);
3547 store_reg128(&mut builder, buffer_ptr, output_slots[0], v);
3548 }
3549
3550 JitOp::U64Add2 => {
3552 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3553 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3554 let result = builder.ins().iadd(a, b);
3555 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3556 }
3557 JitOp::U64Sub2 => {
3558 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3559 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3560 let result = builder.ins().isub(a, b);
3561 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3562 }
3563 JitOp::U64Mul2 => {
3564 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3565 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3566 let result = builder.ins().imul(a, b);
3567 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3568 }
3569 JitOp::U64Div2 => {
3570 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3571 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3572 let zero = builder.ins().iconst(types::I64, 0);
3575 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, b, zero);
3576 let div_block = builder.create_block();
3577 let merge_block = builder.create_block();
3578 builder.append_block_param(merge_block, types::I64);
3579 builder
3580 .ins()
3581 .brif(is_zero, merge_block, &[zero], div_block, &[]);
3582 builder.switch_to_block(div_block);
3583 builder.seal_block(div_block);
3584 let div_result = builder.ins().udiv(a, b);
3585 builder.ins().jump(merge_block, &[div_result]);
3586 builder.switch_to_block(merge_block);
3587 builder.seal_block(merge_block);
3588 let result = builder.block_params(merge_block)[0];
3589 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3590 }
3591 JitOp::U64Mod2 => {
3592 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3593 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3594 let zero = builder.ins().iconst(types::I64, 0);
3597 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, b, zero);
3598 let rem_block = builder.create_block();
3599 let merge_block = builder.create_block();
3600 builder.append_block_param(merge_block, types::I64);
3601 builder
3602 .ins()
3603 .brif(is_zero, merge_block, &[zero], rem_block, &[]);
3604 builder.switch_to_block(rem_block);
3605 builder.seal_block(rem_block);
3606 let rem_result = builder.ins().urem(a, b);
3607 builder.ins().jump(merge_block, &[rem_result]);
3608 builder.switch_to_block(merge_block);
3609 builder.seal_block(merge_block);
3610 let result = builder.block_params(merge_block)[0];
3611 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3612 }
3613 JitOp::U64And => {
3614 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3615 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3616 let result = builder.ins().band(a, b);
3617 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3618 }
3619 JitOp::U64Or => {
3620 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3621 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3622 let result = builder.ins().bor(a, b);
3623 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3624 }
3625 JitOp::U64Xor => {
3626 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3627 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3628 let result = builder.ins().bxor(a, b);
3629 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3630 }
3631 JitOp::U64Shl => {
3632 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3633 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3634 let result = builder.ins().ishl(a, b);
3635 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3636 }
3637 JitOp::U64Shr => {
3638 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3639 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3640 let result = builder.ins().ushr(a, b);
3641 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3642 }
3643 JitOp::U64Not => {
3644 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3645 let result = builder.ins().bnot(a);
3646 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3647 }
3648
3649 JitOp::F64Add => {
3651 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3652 let b = load_slot_f64(&mut builder, buffer_ptr, input_slots[1]);
3653 let result = builder.ins().fadd(a, b);
3654 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3655 }
3656 JitOp::F64Sub => {
3657 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3658 let b = load_slot_f64(&mut builder, buffer_ptr, input_slots[1]);
3659 let result = builder.ins().fsub(a, b);
3660 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3661 }
3662 JitOp::F64Mul => {
3663 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3664 let b = load_slot_f64(&mut builder, buffer_ptr, input_slots[1]);
3665 let result = builder.ins().fmul(a, b);
3666 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3667 }
3668 JitOp::F64Div => {
3669 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3670 let b = load_slot_f64(&mut builder, buffer_ptr, input_slots[1]);
3671 let zero = builder.ins().f64const(0.0);
3673 let is_zero =
3674 builder.ins().fcmp(ir::condcodes::FloatCC::Equal, b, zero);
3675 let div_result = builder.ins().fdiv(a, b);
3676 let result = builder.ins().select(is_zero, zero, div_result);
3677 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3678 }
3679 JitOp::F64Mod => {
3680 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3682 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3683 let call = builder
3684 .ins()
3685 .call(math_binary_refs[F64_MOD_HELPER], &[a, b]);
3686 let result = builder.inst_results(call)[0];
3687 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3688 }
3689
3690 JitOp::IsPositiveCheck { name_ptr, name_len } => {
3691 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3697 let zero = builder.ins().iconst(types::I64, 0);
3698 let is_zero =
3699 builder.ins().icmp(ir::condcodes::IntCC::Equal, val, zero);
3700 let fail_block = builder.create_block();
3701 let ok_block = builder.create_block();
3702 builder.ins().brif(is_zero, fail_block, &[], ok_block, &[]);
3703
3704 builder.switch_to_block(fail_block);
3705 builder.seal_block(fail_block);
3706 let np = builder.ins().iconst(types::I64, *name_ptr as i64);
3707 let nl = builder.ins().iconst(types::I64, *name_len as i64);
3708 let _ = builder.ins().call(is_positive_fail_ref, &[val, np, nl]);
3709 builder.ins().jump(ok_block, &[]);
3713
3714 builder.switch_to_block(ok_block);
3715 builder.seal_block(ok_block);
3716 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3717 }
3718
3719 JitOp::InRangeCheck(lo, hi) => {
3720 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3724 let lo_v = builder.ins().iconst(types::I64, *lo as i64);
3725 let hi_v = builder.ins().iconst(types::I64, *hi as i64);
3726 let below = builder.ins().icmp(
3727 ir::condcodes::IntCC::UnsignedLessThan,
3728 val,
3729 lo_v,
3730 );
3731 let above = builder.ins().icmp(
3732 ir::condcodes::IntCC::UnsignedGreaterThan,
3733 val,
3734 hi_v,
3735 );
3736 let out_of_range = builder.ins().bor(below, above);
3737
3738 let fail_block = builder.create_block();
3739 let ok_block = builder.create_block();
3740 builder
3741 .ins()
3742 .brif(out_of_range, fail_block, &[], ok_block, &[]);
3743
3744 builder.switch_to_block(fail_block);
3745 builder.seal_block(fail_block);
3746 let _ = builder.ins().call(in_range_fail_ref, &[val, lo_v, hi_v]);
3747 builder.ins().jump(ok_block, &[]);
3748
3749 builder.switch_to_block(ok_block);
3750 builder.seal_block(ok_block);
3751 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3752 }
3753
3754 JitOp::IsOneOfCheck {
3755 allowed,
3756 set_ptr,
3757 set_len,
3758 } => {
3759 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3765 let mut any_match = builder.ins().iconst(types::I8, 0);
3766 for allow in allowed.iter() {
3767 let c = builder.ins().iconst(types::I64, *allow as i64);
3768 let eq = builder.ins().icmp(ir::condcodes::IntCC::Equal, val, c);
3769 any_match = builder.ins().bor(any_match, eq);
3770 }
3771 let fail_block = builder.create_block();
3772 let ok_block = builder.create_block();
3773 builder
3777 .ins()
3778 .brif(any_match, ok_block, &[], fail_block, &[]);
3779
3780 builder.switch_to_block(fail_block);
3781 builder.seal_block(fail_block);
3782 let sp = builder.ins().iconst(types::I64, *set_ptr as i64);
3783 let sl = builder.ins().iconst(types::I64, *set_len as i64);
3784 let _ = builder.ins().call(is_one_of_fail_ref, &[val, sp, sl]);
3785 builder.ins().jump(ok_block, &[]);
3786
3787 builder.switch_to_block(ok_block);
3788 builder.seal_block(ok_block);
3789 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3790 }
3791
3792 JitOp::U64Cmp(cc) => {
3793 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3794 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3795 let cmp = builder.ins().icmp(*cc, a, b);
3796 let zero = builder.ins().iconst(types::I64, 0);
3797 let one = builder.ins().iconst(types::I64, 1);
3798 let result = builder.ins().select(cmp, one, zero);
3799 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3800 }
3801 JitOp::F64Cmp(cc) => {
3802 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3803 let b = load_slot_f64(
3804 &mut builder,
3805 buffer_ptr,
3806 if input_slots.len() > 1 {
3807 input_slots[1]
3808 } else {
3809 input_slots[0]
3810 },
3811 );
3812 let cmp = builder.ins().fcmp(*cc, a, b);
3813 let zero = builder.ins().iconst(types::I64, 0);
3814 let one = builder.ins().iconst(types::I64, 1);
3815 let result = builder.ins().select(cmp, one, zero);
3816 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3817 }
3818 JitOp::SelectU64 => {
3819 let cond = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3820 let a = load_slot(
3821 &mut builder,
3822 buffer_ptr,
3823 if input_slots.len() > 1 {
3824 input_slots[1]
3825 } else {
3826 input_slots[0]
3827 },
3828 );
3829 let b = load_slot(
3830 &mut builder,
3831 buffer_ptr,
3832 if input_slots.len() > 2 {
3833 input_slots[2]
3834 } else {
3835 input_slots[0]
3836 },
3837 );
3838 let zero = builder.ins().iconst(types::I64, 0);
3839 let is_nonzero =
3840 builder
3841 .ins()
3842 .icmp(ir::condcodes::IntCC::NotEqual, cond, zero);
3843 let result = builder.ins().select(is_nonzero, a, b);
3844 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3845 }
3846 JitOp::SelectF64 => {
3847 let cond = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3848 let a = load_slot_f64(
3849 &mut builder,
3850 buffer_ptr,
3851 if input_slots.len() > 1 {
3852 input_slots[1]
3853 } else {
3854 input_slots[0]
3855 },
3856 );
3857 let b = load_slot_f64(
3858 &mut builder,
3859 buffer_ptr,
3860 if input_slots.len() > 2 {
3861 input_slots[2]
3862 } else {
3863 input_slots[0]
3864 },
3865 );
3866 let zero = builder.ins().iconst(types::I64, 0);
3867 let is_nonzero =
3868 builder
3869 .ins()
3870 .icmp(ir::condcodes::IntCC::NotEqual, cond, zero);
3871 let result = builder.ins().select(is_nonzero, a, b);
3872 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3873 }
3874
3875 JitOp::I64ToF64 => {
3876 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3877 let fval = builder.ins().fcvt_from_sint(types::F64, val);
3878 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], fval);
3879 }
3880 JitOp::ToBool => {
3881 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3882 let zero = builder.ins().iconst(types::I64, 0);
3883 let one = builder.ins().iconst(types::I64, 1);
3884 let cmp = builder
3885 .ins()
3886 .icmp(ir::condcodes::IntCC::NotEqual, val, zero);
3887 let result = builder.ins().select(cmp, one, zero);
3888 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3889 }
3890 JitOp::ConstU64(v) | JitOp::ConstF64(v) => {
3891 let result = builder.ins().iconst(types::I64, *v as i64);
3892 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3893 }
3894 JitOp::HashRangeConst(max) => {
3895 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3896 let c_gamma = builder
3897 .ins()
3898 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3899 let x1 = builder.ins().iadd(input, c_gamma);
3900 let s30 = builder.ins().ushr_imm(x1, 30);
3901 let x2 = builder.ins().bxor(x1, s30);
3902 let c_m1 = builder
3903 .ins()
3904 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3905 let x3 = builder.ins().imul(x2, c_m1);
3906 let s27 = builder.ins().ushr_imm(x3, 27);
3907 let x4 = builder.ins().bxor(x3, s27);
3908 let c_m2 = builder
3909 .ins()
3910 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3911 let x5 = builder.ins().imul(x4, c_m2);
3912 let s31 = builder.ins().ushr_imm(x5, 31);
3913 let h = builder.ins().bxor(x5, s31);
3914 if *max == 0 {
3915 let zero = builder.ins().iconst(types::I64, 0);
3916 store_slot(&mut builder, buffer_ptr, output_slots[0], zero);
3917 } else {
3918 let m = builder.ins().iconst(types::I64, *max as i64);
3919 let rem = builder.ins().urem(h, m);
3920 store_slot(&mut builder, buffer_ptr, output_slots[0], rem);
3921 }
3922 }
3923 JitOp::HashIntervalConst(min_bits, max_bits) => {
3924 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3925 let c_gamma = builder
3926 .ins()
3927 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3928 let x1 = builder.ins().iadd(input, c_gamma);
3929 let s30 = builder.ins().ushr_imm(x1, 30);
3930 let x2 = builder.ins().bxor(x1, s30);
3931 let c_m1 = builder
3932 .ins()
3933 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3934 let x3 = builder.ins().imul(x2, c_m1);
3935 let s27 = builder.ins().ushr_imm(x3, 27);
3936 let x4 = builder.ins().bxor(x3, s27);
3937 let c_m2 = builder
3938 .ins()
3939 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3940 let x5 = builder.ins().imul(x4, c_m2);
3941 let s31 = builder.ins().ushr_imm(x5, 31);
3942 let h = builder.ins().bxor(x5, s31);
3943
3944 let h_f = builder.ins().fcvt_from_uint(types::F64, h);
3945 let denom = builder.ins().f64const(u64::MAX as f64);
3946 let unit = builder.ins().fdiv(h_f, denom);
3947 let min_f = f64::from_bits(*min_bits);
3948 let max_f = f64::from_bits(*max_bits);
3949 let span = builder.ins().f64const(max_f - min_f);
3950 let min_val = builder.ins().f64const(min_f);
3951 let scaled = builder.ins().fmul(unit, span);
3952 let res_f = builder.ins().fadd(min_val, scaled);
3953 let res = builder
3954 .ins()
3955 .bitcast(types::I64, ir::MemFlags::new(), res_f);
3956 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3957 }
3958 JitOp::InvLerpConst(a_bits, b_bits) => {
3959 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3960 let in_f =
3961 builder
3962 .ins()
3963 .bitcast(types::F64, ir::MemFlags::new(), input);
3964 let a_f = f64::from_bits(*a_bits);
3965 let b_f = f64::from_bits(*b_bits);
3966 let a_val = builder.ins().f64const(a_f);
3967 let inv_span = builder.ins().f64const(1.0 / (b_f - a_f));
3971 let diff = builder.ins().fsub(in_f, a_val);
3972 let t = builder.ins().fmul(diff, inv_span);
3973 let zero = builder.ins().f64const(0.0);
3974 let one = builder.ins().f64const(1.0);
3975 let res_f = clamp_ir(&mut builder, t, zero, one);
3976 let res = builder
3977 .ins()
3978 .bitcast(types::I64, ir::MemFlags::new(), res_f);
3979 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3980 }
3981 JitOp::RemapConst(in_min_bits, in_max_bits, out_min_bits, out_max_bits) => {
3982 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3983 let in_f =
3984 builder
3985 .ins()
3986 .bitcast(types::F64, ir::MemFlags::new(), input);
3987 let in_min = f64::from_bits(*in_min_bits);
3988 let in_max = f64::from_bits(*in_max_bits);
3989 let out_min = f64::from_bits(*out_min_bits);
3990 let out_max = f64::from_bits(*out_max_bits);
3991 let in_span_val = builder.ins().f64const(in_max - in_min);
3995 let in_min_val = builder.ins().f64const(in_min);
3996 let out_min_val = builder.ins().f64const(out_min);
3997 let out_span_val = builder.ins().f64const(out_max - out_min);
3998 let diff = builder.ins().fsub(in_f, in_min_val);
3999 let t = builder.ins().fdiv(diff, in_span_val);
4000 let scaled = builder.ins().fmul(t, out_span_val);
4001 let res_f = builder.ins().fadd(out_min_val, scaled);
4002 let res = builder
4003 .ins()
4004 .bitcast(types::I64, ir::MemFlags::new(), res_f);
4005 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4006 }
4007 JitOp::EpochOffsetConst(base) => {
4008 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4009 let b = builder.ins().iconst(types::I64, *base as i64);
4010 let res = builder.ins().iadd(val, b);
4011 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4012 }
4013 JitOp::EpochScaleConst(factor) => {
4014 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4015 let f = builder.ins().iconst(types::I64, *factor as i64);
4016 let res = builder.ins().imul(val, f);
4017 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4018 }
4019 JitOp::ThreadId => {
4020 let call = builder.ins().call(thread_id_func_ref, &[]);
4021 let res = builder.inst_results(call)[0];
4022 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4023 }
4024 JitOp::CurrentEpochMillis => {
4025 let call = builder.ins().call(current_epoch_millis_func_ref, &[]);
4026 let res = builder.inst_results(call)[0];
4027 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4028 }
4029 JitOp::Perlin1dConst(perm_ptr, freq_bits) => {
4030 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4031 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
4032 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
4033 let call = builder.ins().call(perlin_1d_func_ref, &[input, p, fb]);
4034 let res = builder.inst_results(call)[0];
4035 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4036 }
4037 JitOp::Perlin2dConst(perm_ptr, freq_bits) => {
4038 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4039 let y = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4040 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
4041 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
4042 let call = builder.ins().call(perlin_2d_func_ref, &[x, y, p, fb]);
4043 let res = builder.inst_results(call)[0];
4044 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4045 }
4046 JitOp::Simplex2dConst(perm_ptr, freq_bits) => {
4047 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4048 let y = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4049 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
4050 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
4051 let call = builder.ins().call(simplex_2d_func_ref, &[x, y, p, fb]);
4052 let res = builder.inst_results(call)[0];
4053 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4054 }
4055 JitOp::FractalNoise1dConst(perm_ptr, freq_bits, octaves) => {
4056 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4057 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
4058 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
4059 let oct = builder.ins().iconst(types::I64, *octaves as i64);
4060 let call = builder
4061 .ins()
4062 .call(fractal_noise_1d_func_ref, &[input, p, fb, oct]);
4063 let res = builder.inst_results(call)[0];
4064 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4065 }
4066 JitOp::FractalNoise2dConst(perm_ptr, freq_bits, octaves) => {
4067 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4068 let y = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4069 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
4070 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
4071 let oct = builder.ins().iconst(types::I64, *octaves as i64);
4072 let call = builder
4073 .ins()
4074 .call(fractal_noise_2d_func_ref, &[x, y, p, fb, oct]);
4075 let res = builder.inst_results(call)[0];
4076 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4077 }
4078 JitOp::CycleWalkConst(range, seed, inc) => {
4079 let pos = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4080 let r = builder.ins().iconst(types::I64, *range as i64);
4081 let s = builder.ins().iconst(types::I64, *seed as i64);
4082 let i = builder.ins().iconst(types::I64, *inc as i64);
4083 let call = builder.ins().call(cycle_walk_func_ref, &[pos, r, s, i]);
4084 let res = builder.inst_results(call)[0];
4085 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4086 }
4087
4088 JitOp::VariadicSum => {
4089 if input_slots.is_empty() {
4090 let zero = builder.ins().iconst(types::I64, 0);
4091 store_slot(&mut builder, buffer_ptr, output_slots[0], zero);
4092 } else {
4093 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4094 for &slot in &input_slots[1..] {
4095 let v = load_slot(&mut builder, buffer_ptr, slot);
4096 acc = builder.ins().iadd(acc, v);
4097 }
4098 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4099 }
4100 }
4101 JitOp::VariadicProduct => {
4102 if input_slots.is_empty() {
4103 let one = builder.ins().iconst(types::I64, 1);
4104 store_slot(&mut builder, buffer_ptr, output_slots[0], one);
4105 } else {
4106 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4107 for &slot in &input_slots[1..] {
4108 let v = load_slot(&mut builder, buffer_ptr, slot);
4109 acc = builder.ins().imul(acc, v);
4110 }
4111 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4112 }
4113 }
4114 JitOp::VariadicMin => {
4115 if input_slots.is_empty() {
4116 let ident = builder.ins().iconst(types::I64, u64::MAX as i64);
4125 store_slot(&mut builder, buffer_ptr, output_slots[0], ident);
4126 } else {
4127 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4128 for &slot in &input_slots[1..] {
4129 let v = load_slot(&mut builder, buffer_ptr, slot);
4130 let cmp = builder.ins().icmp(
4131 ir::condcodes::IntCC::UnsignedLessThan,
4132 v,
4133 acc,
4134 );
4135 acc = builder.ins().select(cmp, v, acc);
4136 }
4137 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4138 }
4139 }
4140 JitOp::VariadicMax => {
4141 if input_slots.is_empty() {
4142 let zero = builder.ins().iconst(types::I64, 0);
4143 store_slot(&mut builder, buffer_ptr, output_slots[0], zero);
4144 } else {
4145 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4146 for &slot in &input_slots[1..] {
4147 let v = load_slot(&mut builder, buffer_ptr, slot);
4148 let cmp = builder.ins().icmp(
4149 ir::condcodes::IntCC::UnsignedGreaterThan,
4150 v,
4151 acc,
4152 );
4153 acc = builder.ins().select(cmp, v, acc);
4154 }
4155 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4156 }
4157 }
4158
4159 JitOp::CeilToMultiple => {
4160 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4161 let m = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4162 let zero = builder.ins().iconst(types::I64, 0);
4163 let one = builder.ins().iconst(types::I64, 1);
4164 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, m, zero);
4165 let calc_block = builder.create_block();
4166 let merge_block = builder.create_block();
4167 builder.append_block_param(merge_block, types::I64);
4168 builder
4169 .ins()
4170 .brif(is_zero, merge_block, &[val], calc_block, &[]);
4171 builder.switch_to_block(calc_block);
4172 builder.seal_block(calc_block);
4173 let div = div_ceil(&mut builder, val, m, one);
4176 let high = builder.ins().umulhi(div, m);
4177 let low = builder.ins().imul(div, m);
4178 let zero_hi = builder.ins().iconst(types::I64, 0);
4179 let overflows =
4180 builder
4181 .ins()
4182 .icmp(ir::condcodes::IntCC::NotEqual, high, zero_hi);
4183 let max = builder.ins().iconst(types::I64, -1);
4184 let mul = builder.ins().select(overflows, max, low);
4185 builder.ins().jump(merge_block, &[mul]);
4186 builder.switch_to_block(merge_block);
4187 builder.seal_block(merge_block);
4188 let result = builder.block_params(merge_block)[0];
4189 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4190 }
4191 JitOp::CheckedAdd => {
4192 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4193 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4194 let sum = builder.ins().iadd(a, b);
4195 let is_overflow =
4196 builder
4197 .ins()
4198 .icmp(ir::condcodes::IntCC::UnsignedLessThan, sum, a);
4199 let zero = builder.ins().iconst(types::I64, 0);
4200 let result = builder.ins().select(is_overflow, zero, sum);
4201 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4202 }
4203 JitOp::CheckedSub => {
4204 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4205 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4206 let is_lt =
4207 builder
4208 .ins()
4209 .icmp(ir::condcodes::IntCC::UnsignedLessThan, a, b);
4210 let diff = builder.ins().isub(a, b);
4211 let zero = builder.ins().iconst(types::I64, 0);
4212 let result = builder.ins().select(is_lt, zero, diff);
4213 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4214 }
4215 JitOp::CheckedMul => {
4216 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4217 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4218 let prod = builder.ins().imul(a, b);
4219 let zero = builder.ins().iconst(types::I64, 0);
4220 let a_is_zero =
4221 builder.ins().icmp(ir::condcodes::IntCC::Equal, a, zero);
4222 let div_block = builder.create_block();
4223 let merge_block = builder.create_block();
4224 builder.append_block_param(merge_block, types::I64);
4225 builder
4226 .ins()
4227 .brif(a_is_zero, merge_block, &[zero], div_block, &[]);
4228 builder.switch_to_block(div_block);
4229 builder.seal_block(div_block);
4230 let div = builder.ins().udiv(prod, a);
4231 let ok = builder.ins().icmp(ir::condcodes::IntCC::Equal, div, b);
4232 let mul_res = builder.ins().select(ok, prod, zero);
4233 builder.ins().jump(merge_block, &[mul_res]);
4234 builder.switch_to_block(merge_block);
4235 builder.seal_block(merge_block);
4236 let result = builder.block_params(merge_block)[0];
4237 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4238 }
4239 JitOp::MultiplesAtLeast => {
4240 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4241 let m = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4242 let zero = builder.ins().iconst(types::I64, 0);
4243 let one = builder.ins().iconst(types::I64, 1);
4244 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, m, zero);
4245 let calc_block = builder.create_block();
4246 let merge_block = builder.create_block();
4247 builder.append_block_param(merge_block, types::I64);
4248 builder
4249 .ins()
4250 .brif(is_zero, merge_block, &[zero], calc_block, &[]);
4251 builder.switch_to_block(calc_block);
4252 builder.seal_block(calc_block);
4253 let div = div_ceil(&mut builder, val, m, one);
4254 builder.ins().jump(merge_block, &[div]);
4255 builder.switch_to_block(merge_block);
4256 builder.seal_block(merge_block);
4257 let result = builder.block_params(merge_block)[0];
4258 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4259 }
4260
4261 JitOp::BlendConst(mix_bits) => {
4262 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4267 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4268 let fa = builder.ins().bitcast(types::F64, ir::MemFlags::new(), a);
4269 let fb = builder.ins().bitcast(types::F64, ir::MemFlags::new(), b);
4270 let mix_f64 = f64::from_bits(*mix_bits);
4271 let mix_val = builder.ins().f64const(mix_f64);
4272 let one = builder.ins().f64const(1.0);
4273 let one_minus_mix = builder.ins().fsub(one, mix_val);
4274 let a_part = builder.ins().fmul(fa, one_minus_mix);
4275 let b_part = builder.ins().fmul(fb, mix_val);
4276 let sum = builder.ins().fadd(a_part, b_part);
4277 let result =
4278 builder.ins().bitcast(types::I64, ir::MemFlags::new(), sum);
4279 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4280 }
4281 JitOp::LfsrStepConst(feedback) => {
4282 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4283 let feedback = builder.ins().iconst(types::I64, *feedback as i64);
4284 let one = builder.ins().iconst(types::I64, 1);
4285 let zero = builder.ins().iconst(types::I64, 0);
4286 let shifted = builder.ins().ushr(val, one);
4287 let lsb = builder.ins().band(val, one);
4288 let is_odd =
4289 builder
4290 .ins()
4291 .icmp(ir::condcodes::IntCC::NotEqual, lsb, zero);
4292 let fb_mask = builder.ins().select(is_odd, feedback, zero);
4293 let result = builder.ins().bxor(shifted, fb_mask);
4294 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4295 }
4296 JitOp::PcgConst(seed, stream) => {
4297 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4298 let s = builder.ins().iconst(types::I64, *seed as i64);
4299 let st = builder.ins().iconst(types::I64, *stream as i64);
4300 let call = builder.ins().call(pcg_func_ref, &[input, s, st]);
4301 let result = builder.inst_results(call)[0];
4302 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4303 }
4304 JitOp::PcgStreamConst(seed) => {
4305 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4306 let st = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4307 let s = builder.ins().iconst(types::I64, *seed as i64);
4308 let call = builder.ins().call(pcg_stream_func_ref, &[input, st, s]);
4309 let result = builder.inst_results(call)[0];
4310 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4311 }
4312 JitOp::NOfConst(n, m) => {
4313 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4314 let n_val = builder.ins().iconst(types::I64, *n as i64);
4315 let m_val = builder.ins().iconst(types::I64, *m as i64);
4316 let call = builder.ins().call(n_of_func_ref, &[input, n_val, m_val]);
4317 let result = builder.inst_results(call)[0];
4318 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4319 }
4320
4321 JitOp::SlotCall { kit, scratch_base } => {
4322 emit_slot_call(
4323 &mut builder,
4324 buffer_ptr,
4325 scratch_ptr,
4326 slot_call_ref,
4327 kit,
4328 *scratch_base,
4329 input_slots,
4330 output_slots,
4331 );
4332 }
4333
4334 JitOp::Convert {
4335 from,
4336 to,
4337 kit,
4338 scratch_base,
4339 } => {
4340 emit_conversion(
4341 &mut builder,
4342 buffer_ptr,
4343 input_slots[0],
4344 output_slots[0],
4345 *from,
4346 *to,
4347 |builder| {
4348 emit_slot_call(
4349 builder,
4350 buffer_ptr,
4351 scratch_ptr,
4352 slot_call_ref,
4353 kit,
4354 *scratch_base,
4355 input_slots,
4356 output_slots,
4357 )
4358 },
4359 );
4360 }
4361
4362 JitOp::U64ToStr { scratch_base }
4363 | JitOp::I64ToStr { scratch_base }
4364 | JitOp::F64ToStr { scratch_base } => {
4365 let func = match jit_op {
4368 JitOp::U64ToStr { .. } => u64_to_str_ref,
4369 JitOp::I64ToStr { .. } => i64_to_str_ref,
4370 _ => f64_to_str_ref,
4371 };
4372 let value = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4373 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4374 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4375 builder
4376 .ins()
4377 .call(func, &[scratch_ptr, base_v, buffer_ptr, out_v, value]);
4378 }
4379 JitOp::JsonToStr { scratch_base } => {
4380 let ptr = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4381 let len = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4382 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4383 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4384 builder.ins().call(
4385 json_to_str_ref,
4386 &[scratch_ptr, base_v, buffer_ptr, out_v, ptr, len],
4387 );
4388 }
4389 JitOp::StrConcat { scratch_base } => {
4390 let n_words = input_slots.len();
4393 let frame = builder.create_sized_stack_slot(ir::StackSlotData::new(
4394 ir::StackSlotKind::ExplicitSlot,
4395 (n_words.max(1) * 8) as u32,
4396 3,
4397 ));
4398 for (k, &s) in input_slots.iter().enumerate() {
4399 let v = load_slot(&mut builder, buffer_ptr, s);
4400 builder.ins().stack_store(v, frame, (k * 8) as i32);
4401 }
4402 let pairs_ptr = builder.ins().stack_addr(types::I64, frame, 0);
4403 let n_v = builder.ins().iconst(types::I64, (n_words / 2) as i64);
4404 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4405 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4406 builder.ins().call(
4407 str_concat_ref,
4408 &[scratch_ptr, base_v, buffer_ptr, out_v, pairs_ptr, n_v],
4409 );
4410 }
4411
4412 JitOp::VecProduce { kind, scratch_base } => {
4413 let func = func_of(&vec_producer_refs, *kind);
4417 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4418 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4419 let words = load_words(&mut builder, buffer_ptr, input_slots, 4);
4420 let mut args = vec![scratch_ptr, base_v, buffer_ptr, out_v];
4421 args.extend(words);
4422 builder.ins().call(func, &args);
4423 }
4424 JitOp::VecReduce(kind) => {
4425 let func = func_of(&vec_reducer_refs, *kind);
4426 let words = load_words(&mut builder, buffer_ptr, input_slots, 4);
4427 let call = builder.ins().call(func, &words);
4428 let bits = builder.inst_results(call)[0];
4429 store_slot(&mut builder, buffer_ptr, output_slots[0], bits);
4430 }
4431 JitOp::RegLane(kind) => {
4432 let func = func_of(®_lane_refs, *kind);
4433 let words = load_words(&mut builder, buffer_ptr, input_slots, 3);
4434 let call = builder.ins().call(func, &words);
4435 let word = builder.inst_results(call)[0];
4436 store_slot(&mut builder, buffer_ptr, output_slots[0], word);
4437 }
4438 JitOp::RegProduce(kind) => {
4439 let func = func_of(®_producer_refs, *kind);
4440 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4441 let words = load_words(&mut builder, buffer_ptr, input_slots, 4);
4442 let mut args = vec![buffer_ptr, out_v];
4443 args.extend(words);
4444 builder.ins().call(func, &args);
4445 }
4446 JitOp::RegDotF32 => {
4447 let a =
4451 load_reg128(&mut builder, buffer_ptr, input_slots[0], types::F32X4);
4452 let b =
4453 load_reg128(&mut builder, buffer_ptr, input_slots[2], types::F32X4);
4454 let p = builder.ins().fmul(a, b);
4455 let p0 = builder.ins().extractlane(p, 0);
4456 let p1 = builder.ins().extractlane(p, 1);
4457 let p2 = builder.ins().extractlane(p, 2);
4458 let p3 = builder.ins().extractlane(p, 3);
4459 let s01 = builder.ins().fadd(p0, p1);
4460 let s23 = builder.ins().fadd(p2, p3);
4461 let s = builder.ins().fadd(s01, s23);
4462 let wide = builder.ins().fpromote(types::F64, s);
4463 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], wide);
4464 }
4465 JitOp::RegShuffleConst(mask) => {
4466 let x =
4470 load_reg128(&mut builder, buffer_ptr, input_slots[0], types::I8X16);
4471 let imm = builder
4472 .func
4473 .dfg
4474 .immediates
4475 .push(ir::ConstantData::from(&mask[..]));
4476 let r = builder.ins().shuffle(x, x, imm);
4477 store_reg128(&mut builder, buffer_ptr, output_slots[0], r);
4478 }
4479
4480 JitOp::Fallback => {
4481 }
4484 }
4485 if let Some((inst, mark)) = tracker_store {
4486 let calls = (mark..builder.func.dfg.num_insts()).any(|i| {
4487 builder.func.dfg.insts[ir::Inst::from_u32(i as u32)]
4488 .opcode()
4489 .is_call()
4490 });
4491 if !calls {
4492 builder.func.layout.remove_inst(inst);
4493 }
4494 }
4495 }
4496 if let Some((at, head, _, _, clean_ptr)) = &dispatcher {
4499 let one = builder.ins().iconst(types::I8, 1);
4500 builder
4501 .ins()
4502 .store(ir::MemFlags::trusted(), one, *clean_ptr, unit_idx as i32);
4503 let i = builder.ins().stack_load(types::I64, *at, 0);
4504 let next = builder.ins().iadd_imm(i, 1);
4505 builder.ins().stack_store(next, *at, 0);
4506 builder.ins().jump(*head, &[]);
4507 }
4508 }
4509
4510 if let Some((_, head, exit, _, _)) = dispatcher {
4511 builder.seal_block(head);
4514 builder.switch_to_block(exit);
4515 builder.seal_block(exit);
4516 }
4517 builder.ins().return_(&[]);
4518 builder.finalize();
4519 }
4520 let fallible = ctx.func.layout.blocks().any(|block| {
4523 ctx.func
4524 .layout
4525 .block_insts(block)
4526 .any(|inst| ctx.func.dfg.insts[inst].opcode().is_call())
4527 });
4528
4529 module
4530 .define_function(func_id, &mut ctx)
4531 .map_err(|e| format!("define function: {e}"))?;
4532 module.clear_context(&mut ctx);
4533 defined.push((func_id, fallible));
4534 }
4535 module
4536 .finalize_definitions()
4537 .map_err(|e| format!("finalize: {e}"))?;
4538
4539 let entries: Vec<JitEntry> = defined
4540 .iter()
4541 .map(|&(func_id, fallible)| {
4542 let code_ptr = module.get_finalized_function(func_id);
4543 let straight_fn: NativeFn = unsafe { mem::transmute(code_ptr) };
4544 let dispatch_fn: NativeDispatchFn = unsafe { mem::transmute(code_ptr) };
4545 (straight_fn, dispatch_fn, fallible)
4546 })
4547 .collect();
4548 let kits: Vec<SlotKitRef> = functions
4550 .iter()
4551 .flat_map(|(steps, _)| steps.iter())
4552 .filter_map(|(op, _, _)| op.slot_kit().cloned())
4553 .collect();
4554 let any_fallible = defined.iter().any(|&(_, f)| f);
4555 let code = super::kernels::JitCode::new(module, kits, any_fallible);
4556 Ok((entries, code))
4557}
4558
4559#[allow(clippy::too_many_arguments)]
4568fn emit_slot_call(
4569 builder: &mut FunctionBuilder,
4570 buffer_ptr: ir::Value,
4571 scratch_ptr: ir::Value,
4572 slot_call_ref: ir::FuncRef,
4573 kit: &SlotKitRef,
4574 scratch_base: usize,
4575 input_slots: &[usize],
4576 output_slots: &[usize],
4577) {
4578 let n_in = input_slots.len();
4579 let n_out = output_slots.len();
4580 let frame = |builder: &mut FunctionBuilder, n: usize| {
4581 builder.create_sized_stack_slot(ir::StackSlotData::new(
4582 ir::StackSlotKind::ExplicitSlot,
4583 (n.max(1) * 8) as u32,
4584 3,
4585 ))
4586 };
4587 let in_frame = frame(builder, n_in);
4588 let out_frame = frame(builder, n_out);
4589 for (k, &s) in input_slots.iter().enumerate() {
4590 let v = load_slot(builder, buffer_ptr, s);
4591 builder.ins().stack_store(v, in_frame, (k * 8) as i32);
4592 }
4593 let kit_ptr = builder
4594 .ins()
4595 .iconst(types::I64, std::sync::Arc::as_ptr(&kit.0) as usize as i64);
4596 let in_ptr = builder.ins().stack_addr(types::I64, in_frame, 0);
4597 let n_in_v = builder.ins().iconst(types::I64, n_in as i64);
4598 let out_ptr = builder.ins().stack_addr(types::I64, out_frame, 0);
4599 let n_out_v = builder.ins().iconst(types::I64, n_out as i64);
4600 let base_v = builder.ins().iconst(types::I64, scratch_base as i64);
4601 let n_sc_v = builder.ins().iconst(types::I64, kit.0.scratch.len() as i64);
4602 builder.ins().call(
4603 slot_call_ref,
4604 &[
4605 kit_ptr,
4606 in_ptr,
4607 n_in_v,
4608 out_ptr,
4609 n_out_v,
4610 scratch_ptr,
4611 base_v,
4612 n_sc_v,
4613 ],
4614 );
4615 for (k, &s) in output_slots.iter().enumerate() {
4616 let v = builder
4617 .ins()
4618 .stack_load(types::I64, out_frame, (k * 8) as i32);
4619 store_slot(builder, buffer_ptr, s, v);
4620 }
4621}
4622
4623fn emit_conversion(
4636 builder: &mut FunctionBuilder,
4637 buffer_ptr: ir::Value,
4638 in_slot: usize,
4639 out_slot: usize,
4640 from: crate::ast::PortType,
4641 to: crate::ast::PortType,
4642 slow: impl FnOnce(&mut FunctionBuilder),
4643) {
4644 use Scalar::{Bool, F32, F64, Signed, Unsigned};
4645 use ir::condcodes::{FloatCC, IntCC};
4646 let (Some(src), Some(dst)) = (Scalar::of(from), Scalar::of(to)) else {
4647 slow(builder);
4648 return;
4649 };
4650 let raw = load_slot(builder, buffer_ptr, in_slot);
4651
4652 let float_of = |builder: &mut FunctionBuilder| -> ir::Value {
4654 match src {
4655 F32 => {
4656 let bits = builder.ins().ireduce(types::I32, raw);
4657 let x = builder.ins().bitcast(types::F32, ir::MemFlags::new(), bits);
4658 builder.ins().fpromote(types::F64, x)
4659 }
4660 _ => builder.ins().bitcast(types::F64, ir::MemFlags::new(), raw),
4661 }
4662 };
4663 let store_float = |builder: &mut FunctionBuilder, x: ir::Value| {
4666 let word = if dst == F32 {
4667 let bits = builder.ins().bitcast(types::I32, ir::MemFlags::new(), x);
4668 builder.ins().uextend(types::I64, bits)
4669 } else {
4670 builder.ins().bitcast(types::I64, ir::MemFlags::new(), x)
4671 };
4672 store_slot(builder, buffer_ptr, out_slot, word);
4673 };
4674
4675 match (src, dst) {
4676 (_, Bool) => {
4679 let truth = match src {
4680 F32 | F64 => {
4681 let x = float_of(builder);
4686 let zero = builder.ins().f64const(0.0);
4687 let nonzero = builder.ins().fcmp(FloatCC::NotEqual, x, zero);
4688 let ordered = builder.ins().fcmp(FloatCC::Ordered, x, x);
4689 builder.ins().band(nonzero, ordered)
4690 }
4691 _ => builder.ins().icmp_imm(IntCC::NotEqual, raw, 0),
4692 };
4693 let word = builder.ins().uextend(types::I64, truth);
4694 store_slot(builder, buffer_ptr, out_slot, word);
4695 }
4696 (s, d) if s.int_range().is_some() && d.int_range().is_some() => {
4699 let (smin, smax) = s.int_range().expect("an integer");
4700 let (dmin, dmax) = d.int_range().expect("an integer");
4701 let mut fits = Vec::new();
4702 if dmin > smin {
4703 fits.push(builder.ins().icmp_imm(
4705 IntCC::SignedGreaterThanOrEqual,
4706 raw,
4707 dmin as i64,
4708 ));
4709 }
4710 if dmax < smax {
4711 let cc = if matches!(s, Signed(_)) {
4712 IntCC::SignedLessThanOrEqual
4713 } else {
4714 IntCC::UnsignedLessThanOrEqual
4715 };
4716 fits.push(builder.ins().icmp_imm(cc, raw, dmax as u64 as i64));
4717 }
4718 branch_on(
4719 builder,
4720 fits,
4721 |b| {
4722 store_slot(b, buffer_ptr, out_slot, raw);
4723 },
4724 slow,
4725 );
4726 }
4727 (s, F32 | F64) if s.int_range().is_some() => {
4729 let ty = if dst == F32 { types::F32 } else { types::F64 };
4730 let x = if matches!(s, Signed(_)) {
4731 builder.ins().fcvt_from_sint(ty, raw)
4732 } else {
4733 builder.ins().fcvt_from_uint(ty, raw)
4734 };
4735 store_float(builder, x);
4736 }
4737 (F32, F64) => {
4738 let x = float_of(builder);
4739 store_float(builder, x);
4740 }
4741 (F64, F32) => {
4742 let x = float_of(builder);
4743 let narrow = builder.ins().fdemote(types::F32, x);
4744 store_float(builder, narrow);
4745 }
4746 (F32 | F64, d) => {
4750 let (lo, hi) = match d {
4751 Unsigned(b) => (0.0, 2f64.powi(b as i32)),
4752 Signed(b) => (-(2f64.powi(b as i32 - 1)), 2f64.powi(b as i32 - 1)),
4753 _ => {
4754 slow(builder);
4755 return;
4756 }
4757 };
4758 let x = float_of(builder);
4759 let lo_v = builder.ins().f64const(lo);
4760 let hi_v = builder.ins().f64const(hi);
4761 let above = builder.ins().fcmp(FloatCC::GreaterThanOrEqual, x, lo_v);
4762 let below = builder.ins().fcmp(FloatCC::LessThan, x, hi_v);
4763 branch_on(
4764 builder,
4765 vec![above, below],
4766 |b| {
4767 let word = if matches!(d, Signed(_)) {
4768 b.ins().fcvt_to_sint_sat(types::I64, x)
4769 } else {
4770 b.ins().fcvt_to_uint_sat(types::I64, x)
4771 };
4772 store_slot(b, buffer_ptr, out_slot, word);
4773 },
4774 slow,
4775 );
4776 }
4777 _ => slow(builder),
4778 }
4779}
4780
4781fn branch_on(
4784 builder: &mut FunctionBuilder,
4785 conds: Vec<ir::Value>,
4786 fast: impl FnOnce(&mut FunctionBuilder),
4787 slow: impl FnOnce(&mut FunctionBuilder),
4788) {
4789 let mut conds = conds.into_iter();
4790 let Some(first) = conds.next() else {
4791 fast(builder);
4792 return;
4793 };
4794 let mut ok = first;
4795 for c in conds {
4796 ok = builder.ins().band(ok, c);
4797 }
4798 let fast_block = builder.create_block();
4799 let slow_block = builder.create_block();
4800 let done = builder.create_block();
4801 builder.ins().brif(ok, fast_block, &[], slow_block, &[]);
4802 builder.switch_to_block(fast_block);
4803 builder.seal_block(fast_block);
4804 fast(builder);
4805 builder.ins().jump(done, &[]);
4806 builder.switch_to_block(slow_block);
4807 builder.seal_block(slow_block);
4808 slow(builder);
4809 builder.ins().jump(done, &[]);
4810 builder.switch_to_block(done);
4811 builder.seal_block(done);
4812}
4813
4814fn load_slot(builder: &mut FunctionBuilder, buffer_ptr: ir::Value, slot: usize) -> ir::Value {
4815 let offset = (slot * 8) as i32;
4816 builder
4817 .ins()
4818 .load(types::I64, ir::MemFlags::trusted(), buffer_ptr, offset)
4819}
4820
4821fn store_slot(
4823 builder: &mut FunctionBuilder,
4824 buffer_ptr: ir::Value,
4825 slot: usize,
4826 value: ir::Value,
4827) -> ir::Inst {
4828 let offset = (slot * 8) as i32;
4829 builder
4830 .ins()
4831 .store(ir::MemFlags::trusted(), value, buffer_ptr, offset)
4832}
4833
4834fn reg_lane_type(lane: u8) -> ir::Type {
4837 match lane {
4838 0 => types::I8X16,
4839 1 => types::I16X8,
4840 2 => types::I32X4,
4841 3 => types::I64X2,
4842 4 => types::F32X4,
4843 5 => types::F64X2,
4844 _ => unreachable!("register lane index out of range"),
4845 }
4846}
4847
4848fn load_reg128(
4853 builder: &mut FunctionBuilder,
4854 buffer_ptr: ir::Value,
4855 first_slot: usize,
4856 vt: ir::Type,
4857) -> ir::Value {
4858 let offset = (first_slot * 8) as i32;
4859 builder
4860 .ins()
4861 .load(vt, ir::MemFlags::new(), buffer_ptr, offset)
4862}
4863
4864fn store_reg128(
4866 builder: &mut FunctionBuilder,
4867 buffer_ptr: ir::Value,
4868 first_slot: usize,
4869 value: ir::Value,
4870) {
4871 let offset = (first_slot * 8) as i32;
4872 builder
4873 .ins()
4874 .store(ir::MemFlags::new(), value, buffer_ptr, offset);
4875}
4876
4877fn round_half_away(builder: &mut FunctionBuilder, x: ir::Value) -> ir::Value {
4882 let t = builder.ins().trunc(x);
4883 let frac = builder.ins().fsub(x, t);
4884 let mag = builder.ins().fabs(frac);
4885 let half = builder.ins().f64const(0.5);
4886 let reaches = builder
4887 .ins()
4888 .fcmp(ir::condcodes::FloatCC::GreaterThanOrEqual, mag, half);
4889 let one = builder.ins().f64const(1.0);
4890 let step = builder.ins().fcopysign(one, x);
4891 let up = builder.ins().fadd(t, step);
4892 builder.ins().select(reaches, up, t)
4893}
4894
4895fn clamp_ir(
4900 builder: &mut FunctionBuilder,
4901 x: ir::Value,
4902 lo: ir::Value,
4903 hi: ir::Value,
4904) -> ir::Value {
4905 let below = builder.ins().fcmp(ir::condcodes::FloatCC::LessThan, x, lo);
4906 let above = builder
4907 .ins()
4908 .fcmp(ir::condcodes::FloatCC::GreaterThan, x, hi);
4909 let capped = builder.ins().select(above, hi, x);
4910 builder.ins().select(below, lo, capped)
4911}
4912
4913fn div_ceil(
4916 builder: &mut FunctionBuilder,
4917 val: ir::Value,
4918 m: ir::Value,
4919 one: ir::Value,
4920) -> ir::Value {
4921 let q = builder.ins().udiv(val, m);
4922 let r = builder.ins().urem(val, m);
4923 let zero = builder.ins().iconst(types::I64, 0);
4924 let inexact = builder.ins().icmp(ir::condcodes::IntCC::NotEqual, r, zero);
4925 let q1 = builder.ins().iadd(q, one);
4926 builder.ins().select(inexact, q1, q)
4927}
4928
4929fn func_of<K: PartialEq + Copy>(refs: &[(K, ir::FuncRef)], key: K) -> ir::FuncRef {
4931 refs.iter()
4932 .find(|(k, _)| *k == key)
4933 .map(|(_, r)| *r)
4934 .expect("every helper of the group is declared")
4935}
4936
4937fn load_words(
4940 builder: &mut FunctionBuilder,
4941 buffer_ptr: ir::Value,
4942 input_slots: &[usize],
4943 n: usize,
4944) -> Vec<ir::Value> {
4945 (0..n)
4946 .map(|k| match input_slots.get(k) {
4947 Some(&s) => load_slot(builder, buffer_ptr, s),
4948 None => builder.ins().iconst(types::I64, 0),
4949 })
4950 .collect()
4951}
4952
4953fn load_slot_f64(builder: &mut FunctionBuilder, buffer_ptr: ir::Value, slot: usize) -> ir::Value {
4955 let i64_val = load_slot(builder, buffer_ptr, slot);
4956 builder
4957 .ins()
4958 .bitcast(types::F64, ir::MemFlags::new(), i64_val)
4959}
4960
4961fn store_slot_f64(
4963 builder: &mut FunctionBuilder,
4964 buffer_ptr: ir::Value,
4965 slot: usize,
4966 value: ir::Value,
4967) {
4968 let i64_val = builder
4969 .ins()
4970 .bitcast(types::I64, ir::MemFlags::new(), value);
4971 store_slot(builder, buffer_ptr, slot, i64_val);
4972}
4973
4974#[cfg(test)]
4977mod tests {
4978 use super::*;
4979
4980 #[test]
4985 fn every_scalar_conversion_lowers_by_its_types() {
4986 let mut missing = Vec::new();
4987 for &from in crate::ast::PortType::ALL {
4988 for &to in crate::ast::PortType::ALL {
4989 if from == to || Scalar::of(from).is_none() || Scalar::of(to).is_none() {
4990 continue;
4991 }
4992 let Some(node) = crate::compile::assembly::boundary_adapter(from, to) else {
4993 continue;
4994 };
4995 match classify_node(node.as_ref()) {
4996 JitOp::Convert { from: f, to: t, .. } if f == from && t == to => {}
4997 other => missing.push(format!(
4998 "{} ({from:?} -> {to:?}) classified as {other:?}",
4999 node.meta().name
5000 )),
5001 }
5002 }
5003 }
5004 assert!(missing.is_empty(), "{}", missing.join("\n"));
5005 }
5006
5007 #[test]
5008 fn jit_identity() {
5009 let steps = vec![(JitOp::Identity, vec![0], vec![1])];
5010 let mut output_map = HashMap::new();
5011 output_map.insert("out".into(), 1);
5012 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5013 kernel.eval(&[42]);
5014 assert_eq!(kernel.get("out"), 42);
5015 }
5016
5017 #[test]
5018 fn jit_add_const() {
5019 let steps = vec![(JitOp::AddConst(100), vec![0], vec![1])];
5020 let mut output_map = HashMap::new();
5021 output_map.insert("out".into(), 1);
5022 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5023 kernel.eval(&[5]);
5024 assert_eq!(kernel.get("out"), 105);
5025 }
5026
5027 #[test]
5028 fn jit_mul_const() {
5029 let steps = vec![(JitOp::MulConst(7), vec![0], vec![1])];
5030 let mut output_map = HashMap::new();
5031 output_map.insert("out".into(), 1);
5032 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5033 kernel.eval(&[6]);
5034 assert_eq!(kernel.get("out"), 42);
5035 }
5036
5037 #[test]
5038 fn jit_mod_const() {
5039 let steps = vec![(JitOp::ModConst(100), vec![0], vec![1])];
5040 let mut output_map = HashMap::new();
5041 output_map.insert("out".into(), 1);
5042 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5043 kernel.eval(&[542]);
5044 assert_eq!(kernel.get("out"), 42);
5045 }
5046
5047 #[test]
5048 fn jit_hash() {
5049 let steps = vec![(JitOp::Hash, vec![0], vec![1])];
5050 let mut output_map = HashMap::new();
5051 output_map.insert("out".into(), 1);
5052 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5053
5054 kernel.eval(&[42]);
5055 let v1 = kernel.get("out");
5056
5057 let expected = xxhash_rust::xxh3::xxh3_64(&42u64.to_le_bytes());
5059 assert_eq!(v1, expected);
5060 }
5061
5062 #[test]
5063 fn jit_hash_deterministic() {
5064 let steps = vec![(JitOp::Hash, vec![0], vec![1])];
5065 let mut output_map = HashMap::new();
5066 output_map.insert("out".into(), 1);
5067 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5068
5069 kernel.eval(&[42]);
5070 let v1 = kernel.get("out");
5071 kernel.eval(&[42]);
5072 let v2 = kernel.get("out");
5073 assert_eq!(v1, v2);
5074 }
5075
5076 #[test]
5077 fn jit_chain_hash_mod() {
5078 let steps = vec![
5080 (JitOp::Hash, vec![0], vec![1]), (JitOp::ModConst(1_000_000), vec![1], vec![2]), ];
5083 let mut output_map = HashMap::new();
5084 output_map.insert("user_id".into(), 2);
5085 let mut kernel = compile_jit_raw(1, 3, steps, output_map, Vec::new()).unwrap();
5086
5087 kernel.eval(&[42]);
5088 let uid = kernel.get("user_id");
5089 assert!(uid < 1_000_000, "got {uid}");
5090 }
5091
5092 #[test]
5093 fn jit_clamp_const() {
5094 let steps = vec![(JitOp::ClampConst(10, 50), vec![0], vec![1])];
5095 let mut output_map = HashMap::new();
5096 output_map.insert("out".into(), 1);
5097 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5098
5099 kernel.eval(&[5]);
5100 assert_eq!(kernel.get("out"), 10); kernel.eval(&[30]);
5103 assert_eq!(kernel.get("out"), 30); kernel.eval(&[100]);
5106 assert_eq!(kernel.get("out"), 50); }
5108
5109 #[test]
5110 fn jit_interleave() {
5111 let steps = vec![(JitOp::Interleave, vec![0, 1], vec![2])];
5112 let mut output_map = HashMap::new();
5113 output_map.insert("out".into(), 2);
5114 let mut kernel = compile_jit_raw(2, 3, steps, output_map, Vec::new()).unwrap();
5115
5116 kernel.eval(&[0b101, 0b010]);
5117 assert_eq!(kernel.get("out"), 0b01_10_01);
5119 }
5120
5121 #[test]
5122 fn jit_mixed_radix() {
5123 let steps = vec![(
5125 JitOp::MixedRadixConst(vec![100, 1000, 0]),
5126 vec![0],
5127 vec![1, 2, 3],
5128 )];
5129 let mut output_map = HashMap::new();
5130 output_map.insert("d0".into(), 1);
5131 output_map.insert("d1".into(), 2);
5132 output_map.insert("d2".into(), 3);
5133 let mut kernel = compile_jit_raw(1, 4, steps, output_map, Vec::new()).unwrap();
5134
5135 kernel.eval(&[4_201_337]);
5137 assert_eq!(kernel.get("d0"), 37);
5138 assert_eq!(kernel.get("d1"), 13);
5139 assert_eq!(kernel.get("d2"), 42);
5140 }
5141
5142 #[test]
5143 fn jit_unit_interval() {
5144 let steps = vec![(JitOp::UnitInterval, vec![0], vec![1])];
5145 let mut output_map = HashMap::new();
5146 output_map.insert("out".into(), 1);
5147 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5148
5149 kernel.eval(&[0]);
5150 let v = f64::from_bits(kernel.get("out"));
5151 assert!((v - 0.0).abs() < 1e-10);
5152
5153 kernel.eval(&[u64::MAX]);
5154 let v = f64::from_bits(kernel.get("out"));
5155 assert!((v - 1.0).abs() < 1e-10);
5156 }
5157
5158 #[test]
5159 fn jit_f64_to_u64() {
5160 let steps = vec![(JitOp::F64ToU64, vec![0], vec![1])];
5162 let mut output_map = HashMap::new();
5163 output_map.insert("out".into(), 1);
5164 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5165
5166 kernel.eval(&[3.7f64.to_bits()]);
5167 assert_eq!(kernel.get("out"), 3); }
5169
5170 #[test]
5171 fn jit_round_to_u64() {
5172 let steps = vec![(JitOp::RoundToU64, vec![0], vec![1])];
5173 let mut output_map = HashMap::new();
5174 output_map.insert("out".into(), 1);
5175 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5176
5177 kernel.eval(&[3.7f64.to_bits()]);
5178 assert_eq!(kernel.get("out"), 4);
5179
5180 kernel.eval(&[3.2f64.to_bits()]);
5181 assert_eq!(kernel.get("out"), 3);
5182 }
5183
5184 #[test]
5185 fn jit_clamp_f64() {
5186 let steps = vec![(
5187 JitOp::ClampF64Const(0.0f64.to_bits(), 1.0f64.to_bits()),
5188 vec![0],
5189 vec![1],
5190 )];
5191 let mut output_map = HashMap::new();
5192 output_map.insert("out".into(), 1);
5193 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5194
5195 kernel.eval(&[(-0.5f64).to_bits()]);
5196 assert_eq!(f64::from_bits(kernel.get("out")), 0.0);
5197
5198 kernel.eval(&[0.5f64.to_bits()]);
5199 assert_eq!(f64::from_bits(kernel.get("out")), 0.5);
5200
5201 kernel.eval(&[1.5f64.to_bits()]);
5202 assert_eq!(f64::from_bits(kernel.get("out")), 1.0);
5203 }
5204
5205 #[test]
5206 fn jit_lerp() {
5207 let steps = vec![(
5208 JitOp::LerpConst(10.0f64.to_bits(), 20.0f64.to_bits()),
5209 vec![0],
5210 vec![1],
5211 )];
5212 let mut output_map = HashMap::new();
5213 output_map.insert("out".into(), 1);
5214 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5215
5216 kernel.eval(&[0.0f64.to_bits()]);
5217 assert_eq!(f64::from_bits(kernel.get("out")), 10.0);
5218
5219 kernel.eval(&[1.0f64.to_bits()]);
5220 assert_eq!(f64::from_bits(kernel.get("out")), 20.0);
5221
5222 kernel.eval(&[0.5f64.to_bits()]);
5223 assert_eq!(f64::from_bits(kernel.get("out")), 15.0);
5224 }
5225
5226 #[test]
5227 fn jit_scale_range() {
5228 let steps = vec![(
5229 JitOp::ScaleRangeConst(10.0f64.to_bits(), 10.0f64.to_bits()),
5230 vec![0],
5231 vec![1],
5232 )];
5233 let mut output_map = HashMap::new();
5234 output_map.insert("out".into(), 1);
5235 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5236
5237 kernel.eval(&[0]);
5238 let v = f64::from_bits(kernel.get("out"));
5239 assert!((v - 10.0).abs() < 0.001);
5240
5241 kernel.eval(&[u64::MAX]);
5242 let v = f64::from_bits(kernel.get("out"));
5243 assert!((v - 20.0).abs() < 0.001);
5244 }
5245
5246 #[test]
5247 fn jit_quantize() {
5248 let steps = vec![(JitOp::QuantizeConst(10.0f64.to_bits()), vec![0], vec![1])];
5249 let mut output_map = HashMap::new();
5250 output_map.insert("out".into(), 1);
5251 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5252
5253 kernel.eval(&[13.0f64.to_bits()]);
5254 assert_eq!(f64::from_bits(kernel.get("out")), 10.0);
5255
5256 kernel.eval(&[17.0f64.to_bits()]);
5257 assert_eq!(f64::from_bits(kernel.get("out")), 20.0);
5258 }
5259
5260 #[test]
5261 fn jit_discretize() {
5262 let steps = vec![(
5263 JitOp::DiscretizeConst(100.0f64.to_bits(), 10),
5264 vec![0],
5265 vec![1],
5266 )];
5267 let mut output_map = HashMap::new();
5268 output_map.insert("out".into(), 1);
5269 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5270
5271 kernel.eval(&[0.0f64.to_bits()]);
5272 assert_eq!(kernel.get("out"), 0);
5273
5274 kernel.eval(&[55.0f64.to_bits()]);
5275 assert_eq!(kernel.get("out"), 5);
5276
5277 kernel.eval(&[99.0f64.to_bits()]);
5278 assert_eq!(kernel.get("out"), 9);
5279
5280 kernel.eval(&[200.0f64.to_bits()]);
5282 assert_eq!(kernel.get("out"), 9);
5283 }
5284
5285 #[test]
5286 fn jit_chain_unit_interval_lerp() {
5287 let steps = vec![
5289 (JitOp::UnitInterval, vec![0], vec![1]),
5290 (
5291 JitOp::LerpConst(100.0f64.to_bits(), 200.0f64.to_bits()),
5292 vec![1],
5293 vec![2],
5294 ),
5295 ];
5296 let mut output_map = HashMap::new();
5297 output_map.insert("out".into(), 2);
5298 let mut kernel = compile_jit_raw(1, 3, steps, output_map, Vec::new()).unwrap();
5299
5300 kernel.eval(&[0]);
5301 let v = f64::from_bits(kernel.get("out"));
5302 assert!((v - 100.0).abs() < 0.001);
5303
5304 kernel.eval(&[u64::MAX]);
5305 let v = f64::from_bits(kernel.get("out"));
5306 assert!((v - 200.0).abs() < 0.001);
5307 }
5308
5309 #[test]
5310 fn jit_multi_step_chain() {
5311 let steps = vec![
5313 (JitOp::AddConst(10), vec![0], vec![1]),
5314 (JitOp::MulConst(3), vec![1], vec![2]),
5315 (JitOp::ModConst(100), vec![2], vec![3]),
5316 ];
5317 let mut output_map = HashMap::new();
5318 output_map.insert("out".into(), 3);
5319 let mut kernel = compile_jit_raw(1, 4, steps, output_map, Vec::new()).unwrap();
5320
5321 kernel.eval(&[5]);
5322 assert_eq!(kernel.get("out"), 45);
5324 }
5325
5326 #[test]
5329 fn jit_is_positive_check_passes_positive() {
5330 let steps = vec![(
5331 JitOp::IsPositiveCheck {
5332 name_ptr: 0,
5333 name_len: 0,
5334 },
5335 vec![0],
5336 vec![1],
5337 )];
5338 let mut output_map = HashMap::new();
5339 output_map.insert("out".into(), 1);
5340 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5341 kernel.eval(&[42]);
5342 assert_eq!(kernel.get("out"), 42);
5343 kernel.eval(&[u64::MAX]);
5346 assert_eq!(kernel.get("out"), u64::MAX);
5347 }
5348
5349 #[test]
5350 fn jit_in_range_check_passes_interior() {
5351 let steps = vec![(JitOp::InRangeCheck(10, 100), vec![0], vec![1])];
5352 let mut output_map = HashMap::new();
5353 output_map.insert("out".into(), 1);
5354 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5355 kernel.eval(&[50]);
5356 assert_eq!(kernel.get("out"), 50);
5357 kernel.eval(&[10]);
5359 assert_eq!(kernel.get("out"), 10);
5360 kernel.eval(&[100]);
5361 assert_eq!(kernel.get("out"), 100);
5362 }
5363
5364 #[test]
5369 fn jit_is_one_of_check_passes_allowed_values() {
5370 let steps = vec![(
5371 JitOp::IsOneOfCheck {
5372 allowed: vec![1, 2, 3, 5, 8],
5373 set_ptr: 0,
5374 set_len: 0,
5375 },
5376 vec![0],
5377 vec![1],
5378 )];
5379 let mut output_map = HashMap::new();
5380 output_map.insert("out".into(), 1);
5381 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5382 for v in [1u64, 2, 3, 5, 8] {
5384 kernel.eval(&[v]);
5385 assert_eq!(kernel.get("out"), v);
5386 }
5387 }
5388
5389 #[test]
5390 fn jit_is_one_of_check_accepts_single_element_allow_list() {
5391 let steps = vec![(
5394 JitOp::IsOneOfCheck {
5395 allowed: vec![42],
5396 set_ptr: 0,
5397 set_len: 0,
5398 },
5399 vec![0],
5400 vec![1],
5401 )];
5402 let mut output_map = HashMap::new();
5403 output_map.insert("out".into(), 1);
5404 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5405 kernel.eval(&[42]);
5406 assert_eq!(kernel.get("out"), 42);
5407 }
5408
5409 fn extract_panic_msg(payload: Box<dyn std::any::Any + Send + 'static>) -> String {
5419 payload
5420 .downcast_ref::<String>()
5421 .cloned()
5422 .or_else(|| payload.downcast_ref::<&str>().map(|s| s.to_string()))
5423 .unwrap_or_else(|| "(non-string panic)".into())
5424 }
5425
5426 #[test]
5427 fn jit_is_positive_violation_is_catchable() {
5428 let steps = vec![(
5429 JitOp::IsPositiveCheck {
5430 name_ptr: 0,
5431 name_len: 0,
5432 },
5433 vec![0],
5434 vec![1],
5435 )];
5436 let mut output_map = HashMap::new();
5437 output_map.insert("out".into(), 1);
5438 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5439 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[0])))
5440 .expect_err("JIT violation should panic");
5441 assert!(extract_panic_msg(err).contains("must be > 0"));
5442 }
5443
5444 #[test]
5445 fn jit_in_range_violation_is_catchable() {
5446 let steps = vec![(JitOp::InRangeCheck(10, 100), vec![0], vec![1])];
5447 let mut output_map = HashMap::new();
5448 output_map.insert("out".into(), 1);
5449 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5450
5451 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[5])))
5452 .expect_err("below-range should panic");
5453 assert!(extract_panic_msg(err).contains("outside [10, 100]"));
5454
5455 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[500])))
5456 .expect_err("above-range should panic");
5457 assert!(extract_panic_msg(err).contains("outside [10, 100]"));
5458 }
5459
5460 #[test]
5461 fn jit_is_one_of_violation_is_catchable() {
5462 let steps = vec![(
5463 JitOp::IsOneOfCheck {
5464 allowed: vec![1, 3, 5],
5465 set_ptr: 0,
5466 set_len: 0,
5467 },
5468 vec![0],
5469 vec![1],
5470 )];
5471 let mut output_map = HashMap::new();
5472 output_map.insert("out".into(), 1);
5473 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5474 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[2])))
5475 .expect_err("disallowed value should panic");
5476 assert!(extract_panic_msg(err).contains("not in allowed set"));
5477 }
5478
5479 #[test]
5480 fn invoke_with_catch_restores_slot_after_foreign_panic() {
5481 let caught = std::panic::catch_unwind(|| {
5489 invoke_with_catch(|| panic!("foreign panic"));
5490 });
5491 assert!(caught.is_err(), "foreign panic should propagate out");
5492
5493 let steps = vec![(
5495 JitOp::IsPositiveCheck {
5496 name_ptr: 0,
5497 name_len: 0,
5498 },
5499 vec![0],
5500 vec![1],
5501 )];
5502 let mut output_map = HashMap::new();
5503 output_map.insert("out".into(), 1);
5504 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5505 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[0])))
5506 .expect_err("JIT violation should panic cleanly after foreign panic");
5507 assert!(extract_panic_msg(err).contains("must be > 0"));
5508
5509 kernel.eval(&[42]);
5511 assert_eq!(kernel.get("out"), 42);
5512 }
5513
5514 #[test]
5515 fn jit_kernel_survives_multiple_violations() {
5516 let steps = vec![(
5520 JitOp::IsPositiveCheck {
5521 name_ptr: 0,
5522 name_len: 0,
5523 },
5524 vec![0],
5525 vec![1],
5526 )];
5527 let mut output_map = HashMap::new();
5528 output_map.insert("out".into(), 1);
5529 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5530
5531 for _ in 0..3 {
5532 let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[0])))
5533 .expect_err("violation should still panic");
5534 }
5535 kernel.eval(&[42]);
5537 assert_eq!(kernel.get("out"), 42);
5538 }
5539}