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" => {
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::coordinates_only(coord_count),
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 inputs_read: Vec<Vec<usize>> = steps
2254 .iter()
2255 .map(|(_, ins, _)| {
2256 ins.iter()
2257 .copied()
2258 .filter(|&s| producer.get(s).is_some_and(|&p| p == usize::MAX))
2259 .collect()
2260 })
2261 .collect();
2262 let fusible: Vec<bool> = (0..steps.len())
2263 .map(|i| !alone.get(i).copied().unwrap_or(false))
2264 .collect();
2265 let mut class = vec![0u64; steps.len()];
2266 for &v in volatile {
2267 if v < class.len() {
2268 class[v] = 1;
2269 }
2270 }
2271 let rank: Vec<usize> = (0..steps.len()).collect();
2272 crate::compile::fusion_units::plan_units(&preds, &inputs_read, &fusible, &class, &rank, &|_| {
2273 false
2274 })
2275}
2276
2277fn unit_dependents(
2280 input_dependents: Vec<Vec<usize>>,
2281 plan: &crate::compile::fusion_units::UnitPlan,
2282) -> Vec<Vec<usize>> {
2283 input_dependents
2284 .into_iter()
2285 .map(|steps| {
2286 let mut units: Vec<usize> = steps.iter().map(|&s| plan.unit_of[s]).collect();
2287 units.sort_unstable();
2288 units.dedup();
2289 units
2290 })
2291 .collect()
2292}
2293
2294#[allow(clippy::too_many_arguments)]
2297pub(crate) fn compile_jit_raw_with(
2298 coord_count: usize,
2299 total_slots: usize,
2300 steps: Vec<(JitOp, Vec<usize>, Vec<usize>)>,
2301 output_map: HashMap<String, usize>,
2302 nodes: Vec<Box<dyn PolydatNode>>,
2303 externs: crate::compile::externs::Externs,
2304 scratch: super::kernels::ScratchPlan,
2305 volatile: Vec<usize>,
2306 alone: Vec<bool>,
2307) -> Result<JitKernelRaw, String> {
2308 let plan = pure_units(&steps, total_slots, &alone, &volatile);
2311 let (_, entry, code) = compile_jit_impl(&steps, Some(&plan.units), Some(total_slots))?;
2312 let cones =
2313 super::kernels::ConePlan::new(&steps, total_slots, &plan, output_map.values().copied());
2314 let mut core = JitCore::new(
2315 total_slots,
2316 coord_count,
2317 output_map,
2318 code,
2319 nodes,
2320 scratch,
2321 volatile,
2322 entry,
2323 cones,
2324 );
2325 core.set_externs(externs);
2326 core.engine =
2327 crate::compile::select::Engine::PureNative(crate::compile::select::Provenance::Raw);
2328 Ok(JitKernelRaw { core })
2329}
2330
2331pub(crate) type JitSegmentCode = (NativeFn, super::kernels::JitCode);
2335
2336pub(crate) fn compile_jit_entry(
2340 steps: &[(JitOp, Vec<usize>, Vec<usize>)],
2341 tracker: Option<usize>,
2342) -> Result<JitSegmentCode, String> {
2343 let (raw_fn, _, code) = compile_jit_impl(steps, None, tracker)?;
2344 Ok((raw_fn, code))
2345}
2346
2347#[allow(clippy::too_many_arguments)]
2349pub(crate) fn compile_jit_push_pull(
2350 coord_count: usize,
2351 total_slots: usize,
2352 steps: Vec<(JitOp, Vec<usize>, Vec<usize>)>,
2353 output_map: HashMap<String, usize>,
2354 nodes: Vec<Box<dyn PolydatNode>>,
2355 input_dependents: Vec<Vec<usize>>,
2356 externs: crate::compile::externs::Externs,
2357 scratch: super::kernels::ScratchPlan,
2358 volatile: Vec<usize>,
2359 alone: Vec<bool>,
2360) -> Result<JitKernelPushPull, String> {
2361 let buffer_len = total_slots;
2362 let plan = pure_units(&steps, total_slots, &alone, &volatile);
2363 let (_, entry, code) = compile_jit_impl(&steps, Some(&plan.units), Some(total_slots))?;
2364 let step_outs: Vec<&[usize]> = steps.iter().map(|(_, _, o)| o.as_slice()).collect();
2365 let slot_provenance =
2366 crate::compile::slot_provenance(coord_count, buffer_len, &step_outs, &input_dependents);
2367 let cones =
2368 super::kernels::ConePlan::new(&steps, total_slots, &plan, output_map.values().copied());
2369 let input_dependents = unit_dependents(input_dependents, &plan);
2370 let mut core = JitCore::new(
2371 total_slots,
2372 coord_count,
2373 output_map,
2374 code,
2375 nodes,
2376 scratch,
2377 volatile,
2378 entry,
2379 cones,
2380 );
2381 core.set_externs(externs);
2382 Ok(JitKernelPushPull {
2383 core,
2384 input_dependents,
2385 slot_provenance,
2386 changed_mask: crate::kernel::ProvMask::all_below(coord_count),
2387 force_run: false,
2388 })
2389}
2390
2391pub type NativeFn = unsafe fn(*const u64, *mut u64, *mut crate::ast::ScratchBuf);
2395pub type NativeDispatchFn =
2399 unsafe fn(*const u64, *mut u64, *mut crate::ast::ScratchBuf, *const u32, u64, *mut u8);
2400
2401type JitCompiled = (NativeFn, NativeDispatchFn, super::kernels::JitCode);
2407
2408type JitEntry = (NativeFn, NativeDispatchFn, bool);
2411
2412type JitFunctionSpec<'a> = (
2415 &'a [(JitOp, Vec<usize>, Vec<usize>)],
2416 Option<&'a [Vec<usize>]>,
2417);
2418
2419fn compile_jit_impl(
2423 steps: &[(JitOp, Vec<usize>, Vec<usize>)],
2424 dispatch: Option<&[Vec<usize>]>,
2425 tracker: Option<usize>,
2426) -> Result<JitCompiled, String> {
2427 let (entries, code) = compile_jit_module(&[(steps, dispatch)], tracker)?;
2428 let (straight_fn, dispatch_fn, _) = entries[0];
2429 Ok((straight_fn, dispatch_fn, code))
2430}
2431
2432pub(crate) type JitStep = (JitOp, Vec<usize>, Vec<usize>);
2434
2435pub(crate) fn compile_jit_entries(
2439 batches: &[&[JitStep]],
2440 tracker: Option<usize>,
2441) -> Result<(Vec<(NativeFn, bool)>, super::kernels::JitCode), String> {
2442 let specs: Vec<JitFunctionSpec> = batches.iter().map(|&b| (b, None)).collect();
2443 let (entries, code) = compile_jit_module(&specs, tracker)?;
2444 Ok((
2445 entries
2446 .into_iter()
2447 .map(|(f, _, fallible)| (f, fallible))
2448 .collect(),
2449 code,
2450 ))
2451}
2452
2453fn compile_jit_module(
2455 functions: &[JitFunctionSpec],
2456 tracker: Option<usize>,
2457) -> Result<(Vec<JitEntry>, super::kernels::JitCode), String> {
2458 let mut flag_builder = settings::builder();
2459 flag_builder.set("opt_level", "speed").unwrap();
2460 flag_builder.set("unwind_info", "true").unwrap();
2465 flag_builder.set("preserve_frame_pointers", "true").unwrap();
2466 let isa = super::host_isa::build_host_isa(flag_builder)?;
2467
2468 let mut jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
2469
2470 jit_builder.symbol("jit_xxh3_hash", jit_xxh3_hash as *const u8);
2472 jit_builder.symbol("jit_interleave", jit_interleave as *const u8);
2473 jit_builder.symbol("jit_shuffle", jit_shuffle as *const u8);
2474 jit_builder.symbol("jit_lut_sample", jit_lut_sample as *const u8);
2475 jit_builder.symbol("jit_weighted_pick", jit_weighted_pick as *const u8);
2476 jit_builder.symbol("jit_pcg", jit_pcg as *const u8);
2477 jit_builder.symbol("jit_pcg_stream", jit_pcg_stream as *const u8);
2478 jit_builder.symbol("jit_n_of", jit_n_of as *const u8);
2479 jit_builder.symbol("jit_cycle_walk", jit_cycle_walk as *const u8);
2480 jit_builder.symbol("jit_perlin_1d", jit_perlin_1d as *const u8);
2481 jit_builder.symbol("jit_perlin_2d", jit_perlin_2d as *const u8);
2482 jit_builder.symbol("jit_simplex_2d", jit_simplex_2d as *const u8);
2483 jit_builder.symbol("jit_fractal_noise_1d", jit_fractal_noise_1d as *const u8);
2484 jit_builder.symbol("jit_fractal_noise_2d", jit_fractal_noise_2d as *const u8);
2485 jit_builder.symbol("jit_thread_id", jit_thread_id as *const u8);
2486 jit_builder.symbol(
2487 "jit_current_epoch_millis",
2488 jit_current_epoch_millis as *const u8,
2489 );
2490 jit_builder.symbol("jit_is_positive_fail", jit_is_positive_fail as *const u8);
2494 jit_builder.symbol("jit_in_range_fail", jit_in_range_fail as *const u8);
2495 jit_builder.symbol("jit_is_one_of_fail", jit_is_one_of_fail as *const u8);
2496 jit_builder.symbol("jit_slot_call", jit_slot_call as *const u8);
2499 jit_builder.symbol("jit_u64_to_str", jit_u64_to_str as *const u8);
2500 jit_builder.symbol("jit_i64_to_str", jit_i64_to_str as *const u8);
2501 jit_builder.symbol("jit_f64_to_str", jit_f64_to_str as *const u8);
2502 jit_builder.symbol("jit_str_concat", jit_str_concat as *const u8);
2503 jit_builder.symbol("jit_json_to_str", jit_json_to_str as *const u8);
2504 jit_builder.symbol("jit_vec_add", jit_vec_add as *const u8);
2505 jit_builder.symbol("jit_vec_scale", jit_vec_scale as *const u8);
2506 jit_builder.symbol("jit_vec_norm", jit_vec_norm as *const u8);
2507 jit_builder.symbol("jit_hash_vec", jit_hash_vec as *const u8);
2508 jit_builder.symbol("jit_xxhash3_vec", jit_xxhash3_vec as *const u8);
2509 jit_builder.symbol("jit_reg_to_vec_f32", jit_reg_to_vec_f32 as *const u8);
2510 jit_builder.symbol("jit_vec_dot", jit_vec_dot as *const u8);
2511 jit_builder.symbol("jit_vec_l2", jit_vec_l2 as *const u8);
2512 jit_builder.symbol("jit_vec_cosine", jit_vec_cosine as *const u8);
2513 jit_builder.symbol("jit_lid_mle", jit_lid_mle as *const u8);
2514 jit_builder.symbol("jit_reg_lane_f32", jit_reg_lane_f32 as *const u8);
2515 jit_builder.symbol("jit_reg_lane_i16", jit_reg_lane_i16 as *const u8);
2516 jit_builder.symbol("jit_reg_lane_i64", jit_reg_lane_i64 as *const u8);
2517 jit_builder.symbol("jit_reg_with_lane_f32", jit_reg_with_lane_f32 as *const u8);
2518 jit_builder.symbol("jit_reg_gather_f32", jit_reg_gather_f32 as *const u8);
2519 jit_builder.symbol("jit_vec_to_reg_f32", jit_vec_to_reg_f32 as *const u8);
2520 jit_builder.symbol("jit_reg_mul_i8", jit_reg_mul_i8 as *const u8);
2521 jit_builder.symbol("jit_sin", jit_sin as *const u8);
2523 jit_builder.symbol("jit_cos", jit_cos as *const u8);
2524 jit_builder.symbol("jit_tan", jit_tan as *const u8);
2525 jit_builder.symbol("jit_asin", jit_asin as *const u8);
2526 jit_builder.symbol("jit_acos", jit_acos as *const u8);
2527 jit_builder.symbol("jit_atan", jit_atan as *const u8);
2528 jit_builder.symbol("jit_sqrt", jit_sqrt as *const u8);
2529 jit_builder.symbol("jit_abs_f64", jit_abs_f64 as *const u8);
2530 jit_builder.symbol("jit_ln", jit_ln as *const u8);
2531 jit_builder.symbol("jit_exp", jit_exp as *const u8);
2532 jit_builder.symbol("jit_floor_base10", jit_floor_base10 as *const u8);
2533 jit_builder.symbol("jit_ceiling_base10", jit_ceiling_base10 as *const u8);
2534 jit_builder.symbol("jit_closest_base10", jit_closest_base10 as *const u8);
2535 jit_builder.symbol("jit_floor_decade", jit_floor_decade as *const u8);
2536 jit_builder.symbol("jit_ceiling_decade", jit_ceiling_decade as *const u8);
2537 jit_builder.symbol("jit_closest_decade", jit_closest_decade as *const u8);
2538 jit_builder.symbol("jit_floor_binomial", jit_floor_binomial as *const u8);
2539 jit_builder.symbol("jit_ceiling_binomial", jit_ceiling_binomial as *const u8);
2540 jit_builder.symbol("jit_closest_binomial", jit_closest_binomial as *const u8);
2541 jit_builder.symbol("jit_floor_fibonacci", jit_floor_fibonacci as *const u8);
2542 jit_builder.symbol("jit_ceiling_fibonacci", jit_ceiling_fibonacci as *const u8);
2543 jit_builder.symbol("jit_closest_fibonacci", jit_closest_fibonacci as *const u8);
2544 jit_builder.symbol("jit_atan2", jit_atan2 as *const u8);
2545 jit_builder.symbol("jit_pow", jit_pow as *const u8);
2546 jit_builder.symbol("jit_round_nearest", jit_round_nearest as *const u8);
2547 jit_builder.symbol("jit_round_floor", jit_round_floor as *const u8);
2548 jit_builder.symbol("jit_round_ceiling", jit_round_ceiling as *const u8);
2549 jit_builder.symbol("jit_f64_mod", jit_f64_mod as *const u8);
2550 jit_builder.symbol("jit_div_zero_fail", jit_div_zero_fail as *const u8);
2551
2552 let mut module = JITModule::new(jit_builder);
2553
2554 let hash_func_id = {
2556 let mut sig = module.make_signature();
2557 sig.params.push(AbiParam::new(types::I64));
2558 sig.returns.push(AbiParam::new(types::I64));
2559 module
2560 .declare_function("jit_xxh3_hash", Linkage::Import, &sig)
2561 .map_err(|e| format!("declare hash: {e}"))?
2562 };
2563
2564 let interleave_func_id = {
2566 let mut sig = module.make_signature();
2567 sig.params.push(AbiParam::new(types::I64));
2568 sig.params.push(AbiParam::new(types::I64));
2569 sig.returns.push(AbiParam::new(types::I64));
2570 module
2571 .declare_function("jit_interleave", Linkage::Import, &sig)
2572 .map_err(|e| format!("declare interleave: {e}"))?
2573 };
2574
2575 let shuffle_func_id = {
2577 let mut sig = module.make_signature();
2578 for _ in 0..4 {
2579 sig.params.push(AbiParam::new(types::I64));
2580 }
2581 sig.returns.push(AbiParam::new(types::I64));
2582 module
2583 .declare_function("jit_shuffle", Linkage::Import, &sig)
2584 .map_err(|e| format!("declare shuffle: {e}"))?
2585 };
2586
2587 let lut_sample_func_id = {
2589 let mut sig = module.make_signature();
2590 for _ in 0..3 {
2591 sig.params.push(AbiParam::new(types::I64));
2592 }
2593 sig.returns.push(AbiParam::new(types::I64));
2594 module
2595 .declare_function("jit_lut_sample", Linkage::Import, &sig)
2596 .map_err(|e| format!("declare lut_sample: {e}"))?
2597 };
2598
2599 let weighted_pick_func_id = {
2601 let mut sig = module.make_signature();
2602 for _ in 0..6 {
2603 sig.params.push(AbiParam::new(types::I64));
2604 }
2605 sig.returns.push(AbiParam::new(types::I64));
2606 module
2607 .declare_function("jit_weighted_pick", Linkage::Import, &sig)
2608 .map_err(|e| format!("declare weighted_pick: {e}"))?
2609 };
2610
2611 let pcg_func_id = {
2612 let mut sig = module.make_signature();
2613 for _ in 0..3 {
2614 sig.params.push(AbiParam::new(types::I64));
2615 }
2616 sig.returns.push(AbiParam::new(types::I64));
2617 module
2618 .declare_function("jit_pcg", Linkage::Import, &sig)
2619 .map_err(|e| format!("declare pcg: {e}"))?
2620 };
2621 let pcg_stream_func_id = {
2622 let mut sig = module.make_signature();
2623 for _ in 0..3 {
2624 sig.params.push(AbiParam::new(types::I64));
2625 }
2626 sig.returns.push(AbiParam::new(types::I64));
2627 module
2628 .declare_function("jit_pcg_stream", Linkage::Import, &sig)
2629 .map_err(|e| format!("declare pcg_stream: {e}"))?
2630 };
2631 let n_of_func_id = {
2632 let mut sig = module.make_signature();
2633 for _ in 0..3 {
2634 sig.params.push(AbiParam::new(types::I64));
2635 }
2636 sig.returns.push(AbiParam::new(types::I64));
2637 module
2638 .declare_function("jit_n_of", Linkage::Import, &sig)
2639 .map_err(|e| format!("declare n_of: {e}"))?
2640 };
2641 let cycle_walk_func_id = {
2642 let mut sig = module.make_signature();
2643 for _ in 0..4 {
2644 sig.params.push(AbiParam::new(types::I64));
2645 }
2646 sig.returns.push(AbiParam::new(types::I64));
2647 module
2648 .declare_function("jit_cycle_walk", Linkage::Import, &sig)
2649 .map_err(|e| format!("declare cycle_walk: {e}"))?
2650 };
2651 let perlin_1d_func_id = {
2652 let mut sig = module.make_signature();
2653 for _ in 0..3 {
2654 sig.params.push(AbiParam::new(types::I64));
2655 }
2656 sig.returns.push(AbiParam::new(types::I64));
2657 module
2658 .declare_function("jit_perlin_1d", Linkage::Import, &sig)
2659 .map_err(|e| format!("declare perlin_1d: {e}"))?
2660 };
2661 let perlin_2d_func_id = {
2662 let mut sig = module.make_signature();
2663 for _ in 0..4 {
2664 sig.params.push(AbiParam::new(types::I64));
2665 }
2666 sig.returns.push(AbiParam::new(types::I64));
2667 module
2668 .declare_function("jit_perlin_2d", Linkage::Import, &sig)
2669 .map_err(|e| format!("declare perlin_2d: {e}"))?
2670 };
2671 let simplex_2d_func_id = {
2672 let mut sig = module.make_signature();
2673 for _ in 0..4 {
2674 sig.params.push(AbiParam::new(types::I64));
2675 }
2676 sig.returns.push(AbiParam::new(types::I64));
2677 module
2678 .declare_function("jit_simplex_2d", Linkage::Import, &sig)
2679 .map_err(|e| format!("declare simplex_2d: {e}"))?
2680 };
2681 let fractal_noise_1d_func_id = {
2682 let mut sig = module.make_signature();
2683 for _ in 0..4 {
2684 sig.params.push(AbiParam::new(types::I64));
2685 }
2686 sig.returns.push(AbiParam::new(types::I64));
2687 module
2688 .declare_function("jit_fractal_noise_1d", Linkage::Import, &sig)
2689 .map_err(|e| format!("declare fractal_noise_1d: {e}"))?
2690 };
2691 let fractal_noise_2d_func_id = {
2692 let mut sig = module.make_signature();
2693 for _ in 0..5 {
2694 sig.params.push(AbiParam::new(types::I64));
2695 }
2696 sig.returns.push(AbiParam::new(types::I64));
2697 module
2698 .declare_function("jit_fractal_noise_2d", Linkage::Import, &sig)
2699 .map_err(|e| format!("declare fractal_noise_2d: {e}"))?
2700 };
2701 let thread_id_func_id = {
2702 let mut sig = module.make_signature();
2703 sig.returns.push(AbiParam::new(types::I64));
2704 module
2705 .declare_function("jit_thread_id", Linkage::Import, &sig)
2706 .map_err(|e| format!("declare thread_id: {e}"))?
2707 };
2708 let current_epoch_millis_func_id = {
2709 let mut sig = module.make_signature();
2710 sig.returns.push(AbiParam::new(types::I64));
2711 module
2712 .declare_function("jit_current_epoch_millis", Linkage::Import, &sig)
2713 .map_err(|e| format!("declare current_epoch_millis: {e}"))?
2714 };
2715
2716 let math_unary_names = [
2718 "jit_sin",
2719 "jit_cos",
2720 "jit_tan",
2721 "jit_asin",
2722 "jit_acos",
2723 "jit_atan",
2724 "jit_sqrt",
2725 "jit_abs_f64",
2726 "jit_ln",
2727 "jit_exp",
2728 "jit_floor_base10",
2729 "jit_ceiling_base10",
2730 "jit_closest_base10",
2731 "jit_floor_decade",
2732 "jit_ceiling_decade",
2733 "jit_closest_decade",
2734 "jit_floor_binomial",
2735 "jit_ceiling_binomial",
2736 "jit_closest_binomial",
2737 "jit_floor_fibonacci",
2738 "jit_ceiling_fibonacci",
2739 "jit_closest_fibonacci",
2740 ];
2741 let mut math_unary_ids = Vec::new();
2742 for name in &math_unary_names {
2743 let mut sig = module.make_signature();
2744 sig.params.push(AbiParam::new(types::I64));
2745 sig.returns.push(AbiParam::new(types::I64));
2746 math_unary_ids.push(
2747 module
2748 .declare_function(name, Linkage::Import, &sig)
2749 .map_err(|e| format!("declare {name}: {e}"))?,
2750 );
2751 }
2752
2753 let is_positive_fail_id = {
2757 let mut sig = module.make_signature();
2758 sig.params.push(AbiParam::new(types::I64));
2759 sig.params.push(AbiParam::new(types::I64));
2760 sig.params.push(AbiParam::new(types::I64));
2761 sig.returns.push(AbiParam::new(types::I64));
2762 module
2763 .declare_function("jit_is_positive_fail", Linkage::Import, &sig)
2764 .map_err(|e| format!("declare is_positive_fail: {e}"))?
2765 };
2766
2767 let in_range_fail_id = {
2769 let mut sig = module.make_signature();
2770 for _ in 0..3 {
2771 sig.params.push(AbiParam::new(types::I64));
2772 }
2773 sig.returns.push(AbiParam::new(types::I64));
2774 module
2775 .declare_function("jit_in_range_fail", Linkage::Import, &sig)
2776 .map_err(|e| format!("declare in_range_fail: {e}"))?
2777 };
2778
2779 let is_one_of_fail_id = {
2782 let mut sig = module.make_signature();
2783 sig.params.push(AbiParam::new(types::I64));
2784 sig.params.push(AbiParam::new(types::I64));
2785 sig.params.push(AbiParam::new(types::I64));
2786 sig.returns.push(AbiParam::new(types::I64));
2787 module
2788 .declare_function("jit_is_one_of_fail", Linkage::Import, &sig)
2789 .map_err(|e| format!("declare is_one_of_fail: {e}"))?
2790 };
2791
2792 let math_binary_names = [
2794 "jit_atan2",
2795 "jit_pow",
2796 "jit_round_nearest",
2797 "jit_round_floor",
2798 "jit_round_ceiling",
2799 "jit_f64_mod",
2800 ];
2801 const F64_MOD_HELPER: usize = 5;
2802
2803 let div_zero_fail_id = {
2805 let mut sig = module.make_signature();
2806 sig.params.push(AbiParam::new(types::I64));
2807 sig.returns.push(AbiParam::new(types::I64));
2808 module
2809 .declare_function("jit_div_zero_fail", Linkage::Import, &sig)
2810 .map_err(|e| format!("declare div_zero_fail: {e}"))?
2811 };
2812 let mut math_binary_ids = Vec::new();
2813 for name in &math_binary_names {
2814 let mut sig = module.make_signature();
2815 sig.params.push(AbiParam::new(types::I64));
2816 sig.params.push(AbiParam::new(types::I64));
2817 sig.returns.push(AbiParam::new(types::I64));
2818 math_binary_ids.push(
2819 module
2820 .declare_function(name, Linkage::Import, &sig)
2821 .map_err(|e| format!("declare {name}: {e}"))?,
2822 );
2823 }
2824
2825 let slot_call_id = {
2828 let mut sig = module.make_signature();
2829 for _ in 0..8 {
2830 sig.params.push(AbiParam::new(types::I64));
2831 }
2832 module
2833 .declare_function("jit_slot_call", Linkage::Import, &sig)
2834 .map_err(|e| format!("declare jit_slot_call: {e}"))?
2835 };
2836
2837 let mut declare_str = |name: &str, args: usize| -> Result<cranelift_module::FuncId, String> {
2841 let mut sig = module.make_signature();
2842 for _ in 0..args {
2843 sig.params.push(AbiParam::new(types::I64));
2844 }
2845 module
2846 .declare_function(name, Linkage::Import, &sig)
2847 .map_err(|e| format!("declare {name}: {e}"))
2848 };
2849 let u64_to_str_id = declare_str("jit_u64_to_str", 5)?;
2850 let i64_to_str_id = declare_str("jit_i64_to_str", 5)?;
2851 let f64_to_str_id = declare_str("jit_f64_to_str", 5)?;
2852 let str_concat_id = declare_str("jit_str_concat", 6)?;
2853 let json_to_str_id = declare_str("jit_json_to_str", 6)?;
2854
2855 let mut declare_words =
2860 |name: &str, args: usize, returns: bool| -> Result<cranelift_module::FuncId, String> {
2861 let mut sig = module.make_signature();
2862 for _ in 0..args {
2863 sig.params.push(AbiParam::new(types::I64));
2864 }
2865 if returns {
2866 sig.returns.push(AbiParam::new(types::I64));
2867 }
2868 module
2869 .declare_function(name, Linkage::Import, &sig)
2870 .map_err(|e| format!("declare {name}: {e}"))
2871 };
2872 let vec_producer_ids = [
2873 (VecProducer::Add, declare_words("jit_vec_add", 8, false)?),
2874 (
2875 VecProducer::Scale,
2876 declare_words("jit_vec_scale", 8, false)?,
2877 ),
2878 (VecProducer::Norm, declare_words("jit_vec_norm", 8, false)?),
2879 (
2880 VecProducer::HashVec,
2881 declare_words("jit_hash_vec", 8, false)?,
2882 ),
2883 (
2884 VecProducer::XxHash3Vec,
2885 declare_words("jit_xxhash3_vec", 8, false)?,
2886 ),
2887 (
2888 VecProducer::RegToVec,
2889 declare_words("jit_reg_to_vec_f32", 8, false)?,
2890 ),
2891 ];
2892 let vec_reducer_ids = [
2893 (VecReducer::Dot, declare_words("jit_vec_dot", 4, true)?),
2894 (VecReducer::L2, declare_words("jit_vec_l2", 4, true)?),
2895 (
2896 VecReducer::Cosine,
2897 declare_words("jit_vec_cosine", 4, true)?,
2898 ),
2899 (VecReducer::LidMle, declare_words("jit_lid_mle", 4, true)?),
2900 ];
2901 let reg_lane_ids = [
2902 (
2903 RegLaneRead::F32,
2904 declare_words("jit_reg_lane_f32", 3, true)?,
2905 ),
2906 (
2907 RegLaneRead::I16,
2908 declare_words("jit_reg_lane_i16", 3, true)?,
2909 ),
2910 (
2911 RegLaneRead::I64,
2912 declare_words("jit_reg_lane_i64", 3, true)?,
2913 ),
2914 ];
2915 let reg_producer_ids = [
2916 (
2917 RegProducer::WithLaneF32,
2918 declare_words("jit_reg_with_lane_f32", 6, false)?,
2919 ),
2920 (
2921 RegProducer::GatherF32,
2922 declare_words("jit_reg_gather_f32", 6, false)?,
2923 ),
2924 (
2925 RegProducer::VecToRegF32,
2926 declare_words("jit_vec_to_reg_f32", 6, false)?,
2927 ),
2928 (
2929 RegProducer::MulI8,
2930 declare_words("jit_reg_mul_i8", 6, false)?,
2931 ),
2932 ];
2933
2934 let mut defined: Vec<(cranelift_module::FuncId, bool)> = Vec::with_capacity(functions.len());
2939 for (function_idx, &(steps, dispatch)) in functions.iter().enumerate() {
2940 let mut sig = module.make_signature();
2944 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() {
2948 sig.params.push(AbiParam::new(types::I64)); sig.params.push(AbiParam::new(types::I64)); sig.params.push(AbiParam::new(types::I64)); }
2952 let func_id = module
2953 .declare_function(
2954 &format!("polydat_kernel_{function_idx}"),
2955 Linkage::Local,
2956 &sig,
2957 )
2958 .map_err(|e| format!("declare kernel: {e}"))?;
2959
2960 let mut ctx = module.make_context();
2961 ctx.func.signature = sig;
2962
2963 let mut fb_ctx = FunctionBuilderContext::new();
2964 {
2965 let mut builder = FunctionBuilder::new(&mut ctx.func, &mut fb_ctx);
2966 let block = builder.create_block();
2967 builder.append_block_params_for_function_params(block);
2968 builder.switch_to_block(block);
2969 builder.seal_block(block);
2970
2971 let _coords_ptr = builder.block_params(block)[0];
2972 let buffer_ptr = builder.block_params(block)[1];
2973 let scratch_ptr = builder.block_params(block)[2];
2974
2975 let hash_func_ref = module.declare_func_in_func(hash_func_id, builder.func);
2977 let interleave_func_ref = module.declare_func_in_func(interleave_func_id, builder.func);
2978 let shuffle_func_ref = module.declare_func_in_func(shuffle_func_id, builder.func);
2979 let lut_sample_func_ref = module.declare_func_in_func(lut_sample_func_id, builder.func);
2980 let weighted_pick_func_ref =
2981 module.declare_func_in_func(weighted_pick_func_id, builder.func);
2982 let is_positive_fail_ref =
2983 module.declare_func_in_func(is_positive_fail_id, builder.func);
2984 let in_range_fail_ref = module.declare_func_in_func(in_range_fail_id, builder.func);
2985 let div_zero_fail_ref = module.declare_func_in_func(div_zero_fail_id, builder.func);
2986 let is_one_of_fail_ref = module.declare_func_in_func(is_one_of_fail_id, builder.func);
2987 let slot_call_ref = module.declare_func_in_func(slot_call_id, builder.func);
2988 let u64_to_str_ref = module.declare_func_in_func(u64_to_str_id, builder.func);
2989 let i64_to_str_ref = module.declare_func_in_func(i64_to_str_id, builder.func);
2990 let f64_to_str_ref = module.declare_func_in_func(f64_to_str_id, builder.func);
2991 let str_concat_ref = module.declare_func_in_func(str_concat_id, builder.func);
2992 let json_to_str_ref = module.declare_func_in_func(json_to_str_id, builder.func);
2993 let vec_producer_refs: Vec<(VecProducer, ir::FuncRef)> = vec_producer_ids
2994 .iter()
2995 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
2996 .collect();
2997 let vec_reducer_refs: Vec<(VecReducer, ir::FuncRef)> = vec_reducer_ids
2998 .iter()
2999 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
3000 .collect();
3001 let reg_lane_refs: Vec<(RegLaneRead, ir::FuncRef)> = reg_lane_ids
3002 .iter()
3003 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
3004 .collect();
3005 let reg_producer_refs: Vec<(RegProducer, ir::FuncRef)> = reg_producer_ids
3006 .iter()
3007 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
3008 .collect();
3009 let pcg_func_ref = module.declare_func_in_func(pcg_func_id, builder.func);
3010 let pcg_stream_func_ref = module.declare_func_in_func(pcg_stream_func_id, builder.func);
3011 let n_of_func_ref = module.declare_func_in_func(n_of_func_id, builder.func);
3012 let cycle_walk_func_ref = module.declare_func_in_func(cycle_walk_func_id, builder.func);
3013 let perlin_1d_func_ref = module.declare_func_in_func(perlin_1d_func_id, builder.func);
3014 let perlin_2d_func_ref = module.declare_func_in_func(perlin_2d_func_id, builder.func);
3015 let simplex_2d_func_ref = module.declare_func_in_func(simplex_2d_func_id, builder.func);
3016 let fractal_noise_1d_func_ref =
3017 module.declare_func_in_func(fractal_noise_1d_func_id, builder.func);
3018 let fractal_noise_2d_func_ref =
3019 module.declare_func_in_func(fractal_noise_2d_func_id, builder.func);
3020 let thread_id_func_ref = module.declare_func_in_func(thread_id_func_id, builder.func);
3021 let current_epoch_millis_func_ref =
3022 module.declare_func_in_func(current_epoch_millis_func_id, builder.func);
3023 let math_unary_refs: Vec<_> = math_unary_ids
3024 .iter()
3025 .map(|id| module.declare_func_in_func(*id, builder.func))
3026 .collect();
3027 let math_binary_refs: Vec<_> = math_binary_ids
3028 .iter()
3029 .map(|id| module.declare_func_in_func(*id, builder.func))
3030 .collect();
3031 let everything: [Vec<usize>; 1] = [(0..steps.len()).collect()];
3039 let schedule: &[Vec<usize>] = dispatch.unwrap_or(&everything);
3040 let dispatcher = dispatch.map(|units| {
3041 let list_ptr = builder.block_params(block)[3];
3042 let list_len = builder.block_params(block)[4];
3043 let clean_ptr = builder.block_params(block)[5];
3044 let at = builder.create_sized_stack_slot(ir::StackSlotData::new(
3045 ir::StackSlotKind::ExplicitSlot,
3046 8,
3047 3,
3048 ));
3049 let zero = builder.ins().iconst(types::I64, 0);
3050 builder.ins().stack_store(zero, at, 0);
3051 let head = builder.create_block();
3052 let fetch = builder.create_block();
3053 let dispatch_unit = builder.create_block();
3054 let skip = builder.create_block();
3055 let exit = builder.create_block();
3056 let unit_blocks: Vec<ir::Block> =
3057 units.iter().map(|_| builder.create_block()).collect();
3058 builder.ins().jump(head, &[]);
3059 builder.switch_to_block(head);
3061 let i = builder.ins().stack_load(types::I64, at, 0);
3062 let done = builder.ins().icmp(
3063 ir::condcodes::IntCC::UnsignedGreaterThanOrEqual,
3064 i,
3065 list_len,
3066 );
3067 builder.ins().brif(done, exit, &[], fetch, &[]);
3068 builder.switch_to_block(fetch);
3071 builder.seal_block(fetch);
3072 let offset = builder.ins().ishl_imm(i, 2);
3073 let addr = builder.ins().iadd(list_ptr, offset);
3074 let unit = builder
3075 .ins()
3076 .load(types::I32, ir::MemFlags::trusted(), addr, 0);
3077 let unit_wide = builder.ins().uextend(types::I64, unit);
3078 let flag_addr = builder.ins().iadd(clean_ptr, unit_wide);
3079 let flag = builder
3080 .ins()
3081 .load(types::I8, ir::MemFlags::trusted(), flag_addr, 0);
3082 builder.ins().brif(flag, skip, &[], dispatch_unit, &[]);
3083 builder.switch_to_block(skip);
3085 builder.seal_block(skip);
3086 let next = builder.ins().iadd_imm(i, 1);
3087 builder.ins().stack_store(next, at, 0);
3088 builder.ins().jump(head, &[]);
3089 builder.switch_to_block(dispatch_unit);
3091 builder.seal_block(dispatch_unit);
3092 let default = builder.func.dfg.block_call(exit, &[]);
3093 let targets: Vec<ir::BlockCall> = unit_blocks
3094 .iter()
3095 .map(|&b| builder.func.dfg.block_call(b, &[]))
3096 .collect();
3097 let table = builder.create_jump_table(ir::JumpTableData::new(default, &targets));
3098 builder.ins().br_table(unit, table);
3099 (at, head, exit, unit_blocks, clean_ptr)
3100 });
3101
3102 for (unit_idx, members) in schedule.iter().enumerate() {
3104 if let Some((_, _, _, unit_blocks, _)) = &dispatcher {
3105 builder.switch_to_block(unit_blocks[unit_idx]);
3106 builder.seal_block(unit_blocks[unit_idx]);
3107 }
3108 for &step_idx in members {
3109 let (jit_op, input_slots, output_slots) = &steps[step_idx];
3110 let tracker_store = tracker.map(|t| {
3114 let idx = builder.ins().iconst(types::I64, step_idx as i64);
3115 let inst = store_slot(&mut builder, buffer_ptr, t, idx);
3116 (inst, builder.func.dfg.num_insts())
3117 });
3118 match jit_op {
3119 JitOp::Identity => {
3120 for (&i, &o) in input_slots.iter().zip(output_slots.iter()) {
3123 let val = load_slot(&mut builder, buffer_ptr, i);
3124 store_slot(&mut builder, buffer_ptr, o, val);
3125 }
3126 }
3127 JitOp::AddConst(c) => {
3128 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3129 let c_val = builder.ins().iconst(types::I64, *c as i64);
3130 let result = builder.ins().iadd(val, c_val);
3131 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3132 }
3133 JitOp::MulConst(c) => {
3134 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3135 let c_val = builder.ins().iconst(types::I64, *c as i64);
3136 let result = builder.ins().imul(val, c_val);
3137 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3138 }
3139 JitOp::DivConst(c) | JitOp::ModConst(c) => {
3140 let is_div = matches!(jit_op, JitOp::DivConst(_));
3143 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3144 if *c == 0 {
3145 let kind =
3146 builder.ins().iconst(types::I64, if is_div { 0 } else { 1 });
3147 let _ = builder.ins().call(div_zero_fail_ref, &[kind]);
3148 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3149 } else {
3150 let c_val = builder.ins().iconst(types::I64, *c as i64);
3151 let result = if is_div {
3152 builder.ins().udiv(val, c_val)
3153 } else {
3154 builder.ins().urem(val, c_val)
3155 };
3156 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3157 }
3158 }
3159 JitOp::U64DivWire | JitOp::U64ModWire => {
3160 let is_div = matches!(jit_op, JitOp::U64DivWire);
3164 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3165 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3166 let zero = builder.ins().iconst(types::I64, 0);
3167 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, b, zero);
3168 let fail_block = builder.create_block();
3169 let ok_block = builder.create_block();
3170 builder.ins().brif(is_zero, fail_block, &[], ok_block, &[]);
3171 builder.switch_to_block(fail_block);
3172 builder.seal_block(fail_block);
3173 let kind = builder.ins().iconst(types::I64, if is_div { 0 } else { 1 });
3174 let _ = builder.ins().call(div_zero_fail_ref, &[kind]);
3175 builder.ins().jump(ok_block, &[]);
3176 builder.switch_to_block(ok_block);
3177 builder.seal_block(ok_block);
3178 let result = if is_div {
3179 builder.ins().udiv(a, b)
3180 } else {
3181 builder.ins().urem(a, b)
3182 };
3183 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3184 }
3185 JitOp::ClampConst(min, max) => {
3186 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3187 let min_val = builder.ins().iconst(types::I64, *min as i64);
3188 let max_val = builder.ins().iconst(types::I64, *max as i64);
3189 let clamped_lo = builder.ins().umax(val, min_val);
3190 let clamped = builder.ins().umin(clamped_lo, max_val);
3191 store_slot(&mut builder, buffer_ptr, output_slots[0], clamped);
3192 }
3193 JitOp::Interleave => {
3194 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3195 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3196 let call = builder.ins().call(interleave_func_ref, &[a, b]);
3197 let result = builder.inst_results(call)[0];
3198 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3199 }
3200 JitOp::MixedRadixConst(radixes) => {
3201 let mut remainder = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3203 for (i, &radix) in radixes.iter().enumerate() {
3204 if radix == 0 {
3205 store_slot(
3207 &mut builder,
3208 buffer_ptr,
3209 output_slots[i],
3210 remainder,
3211 );
3212 } else {
3213 let r = builder.ins().iconst(types::I64, radix as i64);
3214 let digit = builder.ins().urem(remainder, r);
3215 store_slot(&mut builder, buffer_ptr, output_slots[i], digit);
3216 remainder = builder.ins().udiv(remainder, r);
3217 }
3218 }
3219 }
3220 JitOp::Hash => {
3221 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3222 let call = builder.ins().call(hash_func_ref, &[val]);
3223 let result = builder.inst_results(call)[0];
3224 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3225 }
3226 JitOp::SplitMix64 => {
3227 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3228 let c_gamma = builder
3229 .ins()
3230 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3231 let x1 = builder.ins().iadd(x0, c_gamma);
3232 let s30 = builder.ins().ushr_imm(x1, 30);
3233 let x2 = builder.ins().bxor(x1, s30);
3234 let c_m1 = builder
3235 .ins()
3236 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3237 let x3 = builder.ins().imul(x2, c_m1);
3238 let s27 = builder.ins().ushr_imm(x3, 27);
3239 let x4 = builder.ins().bxor(x3, s27);
3240 let c_m2 = builder
3241 .ins()
3242 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3243 let x5 = builder.ins().imul(x4, c_m2);
3244 let s31 = builder.ins().ushr_imm(x5, 31);
3245 let result = builder.ins().bxor(x5, s31);
3246 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3247 }
3248 JitOp::FairCoin => {
3249 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3250 let c_gamma = builder
3251 .ins()
3252 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3253 let x1 = builder.ins().iadd(x0, c_gamma);
3254 let s30 = builder.ins().ushr_imm(x1, 30);
3255 let x2 = builder.ins().bxor(x1, s30);
3256 let c_m1 = builder
3257 .ins()
3258 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3259 let x3 = builder.ins().imul(x2, c_m1);
3260 let s27 = builder.ins().ushr_imm(x3, 27);
3261 let x4 = builder.ins().bxor(x3, s27);
3262 let c_m2 = builder
3263 .ins()
3264 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3265 let x5 = builder.ins().imul(x4, c_m2);
3266 let s31 = builder.ins().ushr_imm(x5, 31);
3267 let h = builder.ins().bxor(x5, s31);
3268 let one = builder.ins().iconst(types::I64, 1);
3269 let result = builder.ins().band(h, one);
3270 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3271 }
3272 JitOp::CoinFlipConst(threshold) => {
3273 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3274 let thr = builder.ins().iconst(types::I64, *threshold as i64);
3275 let cmp =
3276 builder
3277 .ins()
3278 .icmp(ir::condcodes::IntCC::UnsignedLessThan, x, thr);
3279 let zero = builder.ins().iconst(types::I64, 0);
3280 let one = builder.ins().iconst(types::I64, 1);
3281 let result = builder.ins().select(cmp, one, zero);
3282 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3283 }
3284 JitOp::UnfairCoinConst(p_bits) => {
3285 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3286 let c_gamma = builder
3287 .ins()
3288 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3289 let x1 = builder.ins().iadd(x0, c_gamma);
3290 let s30 = builder.ins().ushr_imm(x1, 30);
3291 let x2 = builder.ins().bxor(x1, s30);
3292 let c_m1 = builder
3293 .ins()
3294 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3295 let x3 = builder.ins().imul(x2, c_m1);
3296 let s27 = builder.ins().ushr_imm(x3, 27);
3297 let x4 = builder.ins().bxor(x3, s27);
3298 let c_m2 = builder
3299 .ins()
3300 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3301 let x5 = builder.ins().imul(x4, c_m2);
3302 let s31 = builder.ins().ushr_imm(x5, 31);
3303 let h = builder.ins().bxor(x5, s31);
3304
3305 let fval = builder.ins().fcvt_from_uint(types::F64, h);
3306 let max_f = builder.ins().f64const(u64::MAX as f64);
3307 let unit = builder.ins().fdiv(fval, max_f);
3308 let p_f = builder.ins().f64const(f64::from_bits(*p_bits));
3309 let cmp =
3310 builder
3311 .ins()
3312 .fcmp(ir::condcodes::FloatCC::LessThan, unit, p_f);
3313 let zero = builder.ins().iconst(types::I64, 0);
3314 let one = builder.ins().iconst(types::I64, 1);
3315 let result = builder.ins().select(cmp, one, zero);
3316 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3317 }
3318 JitOp::ChanceConst(p_bits) => {
3319 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3320 let c_gamma = builder
3321 .ins()
3322 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3323 let x1 = builder.ins().iadd(x0, c_gamma);
3324 let s30 = builder.ins().ushr_imm(x1, 30);
3325 let x2 = builder.ins().bxor(x1, s30);
3326 let c_m1 = builder
3327 .ins()
3328 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3329 let x3 = builder.ins().imul(x2, c_m1);
3330 let s27 = builder.ins().ushr_imm(x3, 27);
3331 let x4 = builder.ins().bxor(x3, s27);
3332 let c_m2 = builder
3333 .ins()
3334 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3335 let x5 = builder.ins().imul(x4, c_m2);
3336 let s31 = builder.ins().ushr_imm(x5, 31);
3337 let h = builder.ins().bxor(x5, s31);
3338
3339 let fval = builder.ins().fcvt_from_uint(types::F64, h);
3340 let max_f = builder.ins().f64const(u64::MAX as f64);
3341 let unit = builder.ins().fdiv(fval, max_f);
3342 let p_f = builder.ins().f64const(f64::from_bits(*p_bits));
3343 let cmp =
3344 builder
3345 .ins()
3346 .fcmp(ir::condcodes::FloatCC::LessThan, unit, p_f);
3347 let zero_bits =
3348 builder.ins().iconst(types::I64, 0.0_f64.to_bits() as i64);
3349 let one_bits =
3350 builder.ins().iconst(types::I64, 1.0_f64.to_bits() as i64);
3351 let result = builder.ins().select(cmp, one_bits, zero_bits);
3352 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3353 }
3354 JitOp::ShuffleConst(feedback, size, min) => {
3355 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3356 let fb = builder.ins().iconst(types::I64, *feedback as i64);
3357 let sz = builder.ins().iconst(types::I64, *size as i64);
3358 let mn = builder.ins().iconst(types::I64, *min as i64);
3359 let call = builder.ins().call(shuffle_func_ref, &[val, fb, sz, mn]);
3360 let result = builder.inst_results(call)[0];
3361 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3362 }
3363
3364 JitOp::UnitInterval => {
3366 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3368 let fval = builder.ins().fcvt_from_uint(types::F64, val);
3369 let max_f = builder.ins().f64const(u64::MAX as f64);
3370 let result = builder.ins().fdiv(fval, max_f);
3371 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3372 }
3373 JitOp::F64ToU64 => {
3374 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3375 let result = builder.ins().fcvt_to_uint_sat(types::I64, fval);
3376 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3377 }
3378 JitOp::RoundToU64 => {
3379 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3380 let rounded = round_half_away(&mut builder, fval);
3381 let result = builder.ins().fcvt_to_uint_sat(types::I64, rounded);
3382 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3383 }
3384 JitOp::FloorToU64 => {
3385 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3386 let floored = builder.ins().floor(fval);
3387 let result = builder.ins().fcvt_to_uint_sat(types::I64, floored);
3388 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3389 }
3390 JitOp::CeilToU64 => {
3391 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3392 let ceiled = builder.ins().ceil(fval);
3393 let result = builder.ins().fcvt_to_uint_sat(types::I64, ceiled);
3394 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3395 }
3396 JitOp::ClampF64Const(min_bits, max_bits) => {
3397 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3398 let fmin = builder.ins().f64const(f64::from_bits(*min_bits));
3399 let fmax = builder.ins().f64const(f64::from_bits(*max_bits));
3400 let clamped = clamp_ir(&mut builder, fval, fmin, fmax);
3401 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], clamped);
3402 }
3403 JitOp::LerpConst(a_bits, b_bits) => {
3404 let t = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3406 let a = builder.ins().f64const(f64::from_bits(*a_bits));
3407 let b = builder.ins().f64const(f64::from_bits(*b_bits));
3408 let diff = builder.ins().fsub(b, a);
3409 let scaled = builder.ins().fmul(t, diff);
3410 let result = builder.ins().fadd(a, scaled);
3411 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3412 }
3413 JitOp::ScaleRangeConst(min_bits, range_bits) => {
3414 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3416 let fval = builder.ins().fcvt_from_uint(types::F64, val);
3417 let max_f = builder.ins().f64const(u64::MAX as f64);
3418 let t = builder.ins().fdiv(fval, max_f);
3419 let fmin = builder.ins().f64const(f64::from_bits(*min_bits));
3420 let frange = builder.ins().f64const(f64::from_bits(*range_bits));
3421 let scaled = builder.ins().fmul(t, frange);
3422 let result = builder.ins().fadd(fmin, scaled);
3423 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3424 }
3425 JitOp::QuantizeConst(step_bits) => {
3426 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3428 let step = builder.ins().f64const(f64::from_bits(*step_bits));
3429 let divided = builder.ins().fdiv(fval, step);
3430 let rounded = round_half_away(&mut builder, divided);
3431 let result = builder.ins().fmul(rounded, step);
3432 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3433 }
3434
3435 JitOp::LutSampleConst(lut_ptr, lut_len) => {
3436 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3438 let ptr_val = builder.ins().iconst(types::I64, *lut_ptr as i64);
3439 let len_val = builder.ins().iconst(types::I64, *lut_len as i64);
3440 let call = builder
3441 .ins()
3442 .call(lut_sample_func_ref, &[input, ptr_val, len_val]);
3443 let result = builder.inst_results(call)[0];
3444 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3445 }
3446 JitOp::DiscretizeConst(range_bits, buckets) => {
3447 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3452 let range = f64::from_bits(*range_bits);
3453 let fzero = builder.ins().f64const(0.0);
3454 let frange = builder.ins().f64const(range);
3455 let fbuckets = builder.ins().f64const(*buckets as f64);
3456 let clamped = clamp_ir(&mut builder, fval, fzero, frange);
3457 let divided = builder.ins().fdiv(clamped, frange);
3458 let scaled = builder.ins().fmul(divided, fbuckets);
3459 let as_u64 = builder.ins().fcvt_to_uint_sat(types::I64, scaled);
3460 let max_bucket =
3461 builder.ins().iconst(types::I64, (*buckets - 1) as i64);
3462 let result = builder.ins().umin(as_u64, max_bucket);
3463 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3464 }
3465
3466 JitOp::WeightedPickConst(
3467 values_ptr,
3468 biases_ptr,
3469 primaries_ptr,
3470 aliases_ptr,
3471 n,
3472 ) => {
3473 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3475 let v_ptr = builder.ins().iconst(types::I64, *values_ptr as i64);
3476 let b_ptr = builder.ins().iconst(types::I64, *biases_ptr as i64);
3477 let p_ptr = builder.ins().iconst(types::I64, *primaries_ptr as i64);
3478 let a_ptr = builder.ins().iconst(types::I64, *aliases_ptr as i64);
3479 let n_val = builder.ins().iconst(types::I64, *n as i64);
3480 let call = builder.ins().call(
3481 weighted_pick_func_ref,
3482 &[input, v_ptr, b_ptr, p_ptr, a_ptr, n_val],
3483 );
3484 let result = builder.inst_results(call)[0];
3485 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3486 }
3487
3488 JitOp::MathUnary(idx) => {
3489 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3490 let func_ref = math_unary_refs[*idx as usize];
3491 let call = builder.ins().call(func_ref, &[input]);
3492 let result = builder.inst_results(call)[0];
3493 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3494 }
3495
3496 JitOp::MathBinary(idx) => {
3497 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3498 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3499 let func_ref = math_binary_refs[*idx as usize];
3500 let call = builder.ins().call(func_ref, &[a, b]);
3501 let result = builder.inst_results(call)[0];
3502 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3503 }
3504
3505 JitOp::ToF64 => {
3506 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3507 let fval = builder.ins().fcvt_from_uint(types::F64, val);
3508 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], fval);
3509 }
3510
3511 JitOp::RegBinOp(lane, arith) => {
3513 let vt = reg_lane_type(*lane);
3514 let a = load_reg128(&mut builder, buffer_ptr, input_slots[0], vt);
3515 let b = load_reg128(&mut builder, buffer_ptr, input_slots[2], vt);
3516 let is_float = matches!(*lane, 4 | 5);
3517 let r = match (arith, is_float) {
3518 (0, false) => builder.ins().iadd(a, b),
3519 (1, false) => builder.ins().isub(a, b),
3520 (2, false) => builder.ins().imul(a, b),
3521 (0, true) => builder.ins().fadd(a, b),
3522 (1, true) => builder.ins().fsub(a, b),
3523 (2, true) => builder.ins().fmul(a, b),
3524 _ => unreachable!("RegBinOp arith index out of range"),
3525 };
3526 store_reg128(&mut builder, buffer_ptr, output_slots[0], r);
3527 }
3528 JitOp::RegCopy => {
3529 let v =
3530 load_reg128(&mut builder, buffer_ptr, input_slots[0], types::I64X2);
3531 store_reg128(&mut builder, buffer_ptr, output_slots[0], v);
3532 }
3533 JitOp::RegSplat(lane) => {
3534 let vt = reg_lane_type(*lane);
3535 let scalar = match *lane {
3536 0 => {
3538 let v = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3539 builder.ins().ireduce(types::I8, v)
3540 }
3541 1 => {
3542 let v = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3543 builder.ins().ireduce(types::I16, v)
3544 }
3545 2 => {
3546 let v = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3547 builder.ins().ireduce(types::I32, v)
3548 }
3549 3 => load_slot(&mut builder, buffer_ptr, input_slots[0]),
3550 4 => {
3552 let f = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3553 builder.ins().fdemote(types::F32, f)
3554 }
3555 5 => load_slot_f64(&mut builder, buffer_ptr, input_slots[0]),
3556 _ => unreachable!("RegSplat lane index out of range"),
3557 };
3558 let v = builder.ins().splat(vt, scalar);
3559 store_reg128(&mut builder, buffer_ptr, output_slots[0], v);
3560 }
3561
3562 JitOp::U64Add2 => {
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().iadd(a, b);
3567 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3568 }
3569 JitOp::U64Sub2 => {
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 result = builder.ins().isub(a, b);
3573 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3574 }
3575 JitOp::U64Mul2 => {
3576 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3577 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3578 let result = builder.ins().imul(a, b);
3579 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3580 }
3581 JitOp::U64Div2 => {
3582 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3583 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3584 let zero = builder.ins().iconst(types::I64, 0);
3587 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, b, zero);
3588 let div_block = builder.create_block();
3589 let merge_block = builder.create_block();
3590 builder.append_block_param(merge_block, types::I64);
3591 builder
3592 .ins()
3593 .brif(is_zero, merge_block, &[zero], div_block, &[]);
3594 builder.switch_to_block(div_block);
3595 builder.seal_block(div_block);
3596 let div_result = builder.ins().udiv(a, b);
3597 builder.ins().jump(merge_block, &[div_result]);
3598 builder.switch_to_block(merge_block);
3599 builder.seal_block(merge_block);
3600 let result = builder.block_params(merge_block)[0];
3601 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3602 }
3603 JitOp::U64Mod2 => {
3604 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3605 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3606 let zero = builder.ins().iconst(types::I64, 0);
3609 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, b, zero);
3610 let rem_block = builder.create_block();
3611 let merge_block = builder.create_block();
3612 builder.append_block_param(merge_block, types::I64);
3613 builder
3614 .ins()
3615 .brif(is_zero, merge_block, &[zero], rem_block, &[]);
3616 builder.switch_to_block(rem_block);
3617 builder.seal_block(rem_block);
3618 let rem_result = builder.ins().urem(a, b);
3619 builder.ins().jump(merge_block, &[rem_result]);
3620 builder.switch_to_block(merge_block);
3621 builder.seal_block(merge_block);
3622 let result = builder.block_params(merge_block)[0];
3623 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3624 }
3625 JitOp::U64And => {
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().band(a, b);
3629 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3630 }
3631 JitOp::U64Or => {
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().bor(a, b);
3635 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3636 }
3637 JitOp::U64Xor => {
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().bxor(a, b);
3641 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3642 }
3643 JitOp::U64Shl => {
3644 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3645 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3646 let result = builder.ins().ishl(a, b);
3647 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3648 }
3649 JitOp::U64Shr => {
3650 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3651 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3652 let result = builder.ins().ushr(a, b);
3653 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3654 }
3655 JitOp::U64Not => {
3656 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3657 let result = builder.ins().bnot(a);
3658 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3659 }
3660
3661 JitOp::F64Add => {
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().fadd(a, b);
3666 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3667 }
3668 JitOp::F64Sub => {
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 result = builder.ins().fsub(a, b);
3672 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3673 }
3674 JitOp::F64Mul => {
3675 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3676 let b = load_slot_f64(&mut builder, buffer_ptr, input_slots[1]);
3677 let result = builder.ins().fmul(a, b);
3678 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3679 }
3680 JitOp::F64Div => {
3681 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3682 let b = load_slot_f64(&mut builder, buffer_ptr, input_slots[1]);
3683 let zero = builder.ins().f64const(0.0);
3685 let is_zero =
3686 builder.ins().fcmp(ir::condcodes::FloatCC::Equal, b, zero);
3687 let div_result = builder.ins().fdiv(a, b);
3688 let result = builder.ins().select(is_zero, zero, div_result);
3689 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3690 }
3691 JitOp::F64Mod => {
3692 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3694 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3695 let call = builder
3696 .ins()
3697 .call(math_binary_refs[F64_MOD_HELPER], &[a, b]);
3698 let result = builder.inst_results(call)[0];
3699 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3700 }
3701
3702 JitOp::IsPositiveCheck { name_ptr, name_len } => {
3703 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3709 let zero = builder.ins().iconst(types::I64, 0);
3710 let is_zero =
3711 builder.ins().icmp(ir::condcodes::IntCC::Equal, val, zero);
3712 let fail_block = builder.create_block();
3713 let ok_block = builder.create_block();
3714 builder.ins().brif(is_zero, fail_block, &[], ok_block, &[]);
3715
3716 builder.switch_to_block(fail_block);
3717 builder.seal_block(fail_block);
3718 let np = builder.ins().iconst(types::I64, *name_ptr as i64);
3719 let nl = builder.ins().iconst(types::I64, *name_len as i64);
3720 let _ = builder.ins().call(is_positive_fail_ref, &[val, np, nl]);
3721 builder.ins().jump(ok_block, &[]);
3725
3726 builder.switch_to_block(ok_block);
3727 builder.seal_block(ok_block);
3728 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3729 }
3730
3731 JitOp::InRangeCheck(lo, hi) => {
3732 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3736 let lo_v = builder.ins().iconst(types::I64, *lo as i64);
3737 let hi_v = builder.ins().iconst(types::I64, *hi as i64);
3738 let below = builder.ins().icmp(
3739 ir::condcodes::IntCC::UnsignedLessThan,
3740 val,
3741 lo_v,
3742 );
3743 let above = builder.ins().icmp(
3744 ir::condcodes::IntCC::UnsignedGreaterThan,
3745 val,
3746 hi_v,
3747 );
3748 let out_of_range = builder.ins().bor(below, above);
3749
3750 let fail_block = builder.create_block();
3751 let ok_block = builder.create_block();
3752 builder
3753 .ins()
3754 .brif(out_of_range, fail_block, &[], ok_block, &[]);
3755
3756 builder.switch_to_block(fail_block);
3757 builder.seal_block(fail_block);
3758 let _ = builder.ins().call(in_range_fail_ref, &[val, lo_v, hi_v]);
3759 builder.ins().jump(ok_block, &[]);
3760
3761 builder.switch_to_block(ok_block);
3762 builder.seal_block(ok_block);
3763 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3764 }
3765
3766 JitOp::IsOneOfCheck {
3767 allowed,
3768 set_ptr,
3769 set_len,
3770 } => {
3771 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3777 let mut any_match = builder.ins().iconst(types::I8, 0);
3778 for allow in allowed.iter() {
3779 let c = builder.ins().iconst(types::I64, *allow as i64);
3780 let eq = builder.ins().icmp(ir::condcodes::IntCC::Equal, val, c);
3781 any_match = builder.ins().bor(any_match, eq);
3782 }
3783 let fail_block = builder.create_block();
3784 let ok_block = builder.create_block();
3785 builder
3789 .ins()
3790 .brif(any_match, ok_block, &[], fail_block, &[]);
3791
3792 builder.switch_to_block(fail_block);
3793 builder.seal_block(fail_block);
3794 let sp = builder.ins().iconst(types::I64, *set_ptr as i64);
3795 let sl = builder.ins().iconst(types::I64, *set_len as i64);
3796 let _ = builder.ins().call(is_one_of_fail_ref, &[val, sp, sl]);
3797 builder.ins().jump(ok_block, &[]);
3798
3799 builder.switch_to_block(ok_block);
3800 builder.seal_block(ok_block);
3801 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3802 }
3803
3804 JitOp::U64Cmp(cc) => {
3805 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3806 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3807 let cmp = builder.ins().icmp(*cc, a, b);
3808 let zero = builder.ins().iconst(types::I64, 0);
3809 let one = builder.ins().iconst(types::I64, 1);
3810 let result = builder.ins().select(cmp, one, zero);
3811 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3812 }
3813 JitOp::F64Cmp(cc) => {
3814 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3815 let b = load_slot_f64(
3816 &mut builder,
3817 buffer_ptr,
3818 if input_slots.len() > 1 {
3819 input_slots[1]
3820 } else {
3821 input_slots[0]
3822 },
3823 );
3824 let cmp = builder.ins().fcmp(*cc, a, b);
3825 let zero = builder.ins().iconst(types::I64, 0);
3826 let one = builder.ins().iconst(types::I64, 1);
3827 let result = builder.ins().select(cmp, one, zero);
3828 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3829 }
3830 JitOp::SelectU64 => {
3831 let cond = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3832 let a = load_slot(
3833 &mut builder,
3834 buffer_ptr,
3835 if input_slots.len() > 1 {
3836 input_slots[1]
3837 } else {
3838 input_slots[0]
3839 },
3840 );
3841 let b = load_slot(
3842 &mut builder,
3843 buffer_ptr,
3844 if input_slots.len() > 2 {
3845 input_slots[2]
3846 } else {
3847 input_slots[0]
3848 },
3849 );
3850 let zero = builder.ins().iconst(types::I64, 0);
3851 let is_nonzero =
3852 builder
3853 .ins()
3854 .icmp(ir::condcodes::IntCC::NotEqual, cond, zero);
3855 let result = builder.ins().select(is_nonzero, a, b);
3856 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3857 }
3858 JitOp::SelectF64 => {
3859 let cond = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3860 let a = load_slot_f64(
3861 &mut builder,
3862 buffer_ptr,
3863 if input_slots.len() > 1 {
3864 input_slots[1]
3865 } else {
3866 input_slots[0]
3867 },
3868 );
3869 let b = load_slot_f64(
3870 &mut builder,
3871 buffer_ptr,
3872 if input_slots.len() > 2 {
3873 input_slots[2]
3874 } else {
3875 input_slots[0]
3876 },
3877 );
3878 let zero = builder.ins().iconst(types::I64, 0);
3879 let is_nonzero =
3880 builder
3881 .ins()
3882 .icmp(ir::condcodes::IntCC::NotEqual, cond, zero);
3883 let result = builder.ins().select(is_nonzero, a, b);
3884 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3885 }
3886
3887 JitOp::I64ToF64 => {
3888 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3889 let fval = builder.ins().fcvt_from_sint(types::F64, val);
3890 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], fval);
3891 }
3892 JitOp::ToBool => {
3893 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3894 let zero = builder.ins().iconst(types::I64, 0);
3895 let one = builder.ins().iconst(types::I64, 1);
3896 let cmp = builder
3897 .ins()
3898 .icmp(ir::condcodes::IntCC::NotEqual, val, zero);
3899 let result = builder.ins().select(cmp, one, zero);
3900 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3901 }
3902 JitOp::ConstU64(v) | JitOp::ConstF64(v) => {
3903 let result = builder.ins().iconst(types::I64, *v as i64);
3904 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3905 }
3906 JitOp::HashRangeConst(max) => {
3907 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3908 let c_gamma = builder
3909 .ins()
3910 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3911 let x1 = builder.ins().iadd(input, c_gamma);
3912 let s30 = builder.ins().ushr_imm(x1, 30);
3913 let x2 = builder.ins().bxor(x1, s30);
3914 let c_m1 = builder
3915 .ins()
3916 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3917 let x3 = builder.ins().imul(x2, c_m1);
3918 let s27 = builder.ins().ushr_imm(x3, 27);
3919 let x4 = builder.ins().bxor(x3, s27);
3920 let c_m2 = builder
3921 .ins()
3922 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3923 let x5 = builder.ins().imul(x4, c_m2);
3924 let s31 = builder.ins().ushr_imm(x5, 31);
3925 let h = builder.ins().bxor(x5, s31);
3926 if *max == 0 {
3927 let zero = builder.ins().iconst(types::I64, 0);
3928 store_slot(&mut builder, buffer_ptr, output_slots[0], zero);
3929 } else {
3930 let m = builder.ins().iconst(types::I64, *max as i64);
3931 let rem = builder.ins().urem(h, m);
3932 store_slot(&mut builder, buffer_ptr, output_slots[0], rem);
3933 }
3934 }
3935 JitOp::HashIntervalConst(min_bits, max_bits) => {
3936 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3937 let c_gamma = builder
3938 .ins()
3939 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3940 let x1 = builder.ins().iadd(input, c_gamma);
3941 let s30 = builder.ins().ushr_imm(x1, 30);
3942 let x2 = builder.ins().bxor(x1, s30);
3943 let c_m1 = builder
3944 .ins()
3945 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3946 let x3 = builder.ins().imul(x2, c_m1);
3947 let s27 = builder.ins().ushr_imm(x3, 27);
3948 let x4 = builder.ins().bxor(x3, s27);
3949 let c_m2 = builder
3950 .ins()
3951 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3952 let x5 = builder.ins().imul(x4, c_m2);
3953 let s31 = builder.ins().ushr_imm(x5, 31);
3954 let h = builder.ins().bxor(x5, s31);
3955
3956 let h_f = builder.ins().fcvt_from_uint(types::F64, h);
3957 let denom = builder.ins().f64const(u64::MAX as f64);
3958 let unit = builder.ins().fdiv(h_f, denom);
3959 let min_f = f64::from_bits(*min_bits);
3960 let max_f = f64::from_bits(*max_bits);
3961 let span = builder.ins().f64const(max_f - min_f);
3962 let min_val = builder.ins().f64const(min_f);
3963 let scaled = builder.ins().fmul(unit, span);
3964 let res_f = builder.ins().fadd(min_val, scaled);
3965 let res = builder
3966 .ins()
3967 .bitcast(types::I64, ir::MemFlags::new(), res_f);
3968 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3969 }
3970 JitOp::InvLerpConst(a_bits, b_bits) => {
3971 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3972 let in_f =
3973 builder
3974 .ins()
3975 .bitcast(types::F64, ir::MemFlags::new(), input);
3976 let a_f = f64::from_bits(*a_bits);
3977 let b_f = f64::from_bits(*b_bits);
3978 let a_val = builder.ins().f64const(a_f);
3979 let inv_span = builder.ins().f64const(1.0 / (b_f - a_f));
3983 let diff = builder.ins().fsub(in_f, a_val);
3984 let t = builder.ins().fmul(diff, inv_span);
3985 let zero = builder.ins().f64const(0.0);
3986 let one = builder.ins().f64const(1.0);
3987 let res_f = clamp_ir(&mut builder, t, zero, one);
3988 let res = builder
3989 .ins()
3990 .bitcast(types::I64, ir::MemFlags::new(), res_f);
3991 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3992 }
3993 JitOp::RemapConst(in_min_bits, in_max_bits, out_min_bits, out_max_bits) => {
3994 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3995 let in_f =
3996 builder
3997 .ins()
3998 .bitcast(types::F64, ir::MemFlags::new(), input);
3999 let in_min = f64::from_bits(*in_min_bits);
4000 let in_max = f64::from_bits(*in_max_bits);
4001 let out_min = f64::from_bits(*out_min_bits);
4002 let out_max = f64::from_bits(*out_max_bits);
4003 let in_span_val = builder.ins().f64const(in_max - in_min);
4007 let in_min_val = builder.ins().f64const(in_min);
4008 let out_min_val = builder.ins().f64const(out_min);
4009 let out_span_val = builder.ins().f64const(out_max - out_min);
4010 let diff = builder.ins().fsub(in_f, in_min_val);
4011 let t = builder.ins().fdiv(diff, in_span_val);
4012 let scaled = builder.ins().fmul(t, out_span_val);
4013 let res_f = builder.ins().fadd(out_min_val, scaled);
4014 let res = builder
4015 .ins()
4016 .bitcast(types::I64, ir::MemFlags::new(), res_f);
4017 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4018 }
4019 JitOp::EpochOffsetConst(base) => {
4020 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4021 let b = builder.ins().iconst(types::I64, *base as i64);
4022 let res = builder.ins().iadd(val, b);
4023 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4024 }
4025 JitOp::EpochScaleConst(factor) => {
4026 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4027 let f = builder.ins().iconst(types::I64, *factor as i64);
4028 let res = builder.ins().imul(val, f);
4029 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4030 }
4031 JitOp::ThreadId => {
4032 let call = builder.ins().call(thread_id_func_ref, &[]);
4033 let res = builder.inst_results(call)[0];
4034 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4035 }
4036 JitOp::CurrentEpochMillis => {
4037 let call = builder.ins().call(current_epoch_millis_func_ref, &[]);
4038 let res = builder.inst_results(call)[0];
4039 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4040 }
4041 JitOp::Perlin1dConst(perm_ptr, freq_bits) => {
4042 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4043 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
4044 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
4045 let call = builder.ins().call(perlin_1d_func_ref, &[input, p, fb]);
4046 let res = builder.inst_results(call)[0];
4047 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4048 }
4049 JitOp::Perlin2dConst(perm_ptr, freq_bits) => {
4050 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4051 let y = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4052 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
4053 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
4054 let call = builder.ins().call(perlin_2d_func_ref, &[x, y, p, fb]);
4055 let res = builder.inst_results(call)[0];
4056 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4057 }
4058 JitOp::Simplex2dConst(perm_ptr, freq_bits) => {
4059 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4060 let y = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4061 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
4062 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
4063 let call = builder.ins().call(simplex_2d_func_ref, &[x, y, p, fb]);
4064 let res = builder.inst_results(call)[0];
4065 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4066 }
4067 JitOp::FractalNoise1dConst(perm_ptr, freq_bits, octaves) => {
4068 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
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_1d_func_ref, &[input, 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::FractalNoise2dConst(perm_ptr, freq_bits, octaves) => {
4079 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4080 let y = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4081 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
4082 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
4083 let oct = builder.ins().iconst(types::I64, *octaves as i64);
4084 let call = builder
4085 .ins()
4086 .call(fractal_noise_2d_func_ref, &[x, y, p, fb, oct]);
4087 let res = builder.inst_results(call)[0];
4088 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4089 }
4090 JitOp::CycleWalkConst(range, seed, inc) => {
4091 let pos = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4092 let r = builder.ins().iconst(types::I64, *range as i64);
4093 let s = builder.ins().iconst(types::I64, *seed as i64);
4094 let i = builder.ins().iconst(types::I64, *inc as i64);
4095 let call = builder.ins().call(cycle_walk_func_ref, &[pos, r, s, i]);
4096 let res = builder.inst_results(call)[0];
4097 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4098 }
4099
4100 JitOp::VariadicSum => {
4101 if input_slots.is_empty() {
4102 let zero = builder.ins().iconst(types::I64, 0);
4103 store_slot(&mut builder, buffer_ptr, output_slots[0], zero);
4104 } else {
4105 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4106 for &slot in &input_slots[1..] {
4107 let v = load_slot(&mut builder, buffer_ptr, slot);
4108 acc = builder.ins().iadd(acc, v);
4109 }
4110 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4111 }
4112 }
4113 JitOp::VariadicProduct => {
4114 if input_slots.is_empty() {
4115 let one = builder.ins().iconst(types::I64, 1);
4116 store_slot(&mut builder, buffer_ptr, output_slots[0], one);
4117 } else {
4118 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4119 for &slot in &input_slots[1..] {
4120 let v = load_slot(&mut builder, buffer_ptr, slot);
4121 acc = builder.ins().imul(acc, v);
4122 }
4123 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4124 }
4125 }
4126 JitOp::VariadicMin => {
4127 if input_slots.is_empty() {
4128 let ident = builder.ins().iconst(types::I64, u64::MAX as i64);
4137 store_slot(&mut builder, buffer_ptr, output_slots[0], ident);
4138 } else {
4139 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4140 for &slot in &input_slots[1..] {
4141 let v = load_slot(&mut builder, buffer_ptr, slot);
4142 let cmp = builder.ins().icmp(
4143 ir::condcodes::IntCC::UnsignedLessThan,
4144 v,
4145 acc,
4146 );
4147 acc = builder.ins().select(cmp, v, acc);
4148 }
4149 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4150 }
4151 }
4152 JitOp::VariadicMax => {
4153 if input_slots.is_empty() {
4154 let zero = builder.ins().iconst(types::I64, 0);
4155 store_slot(&mut builder, buffer_ptr, output_slots[0], zero);
4156 } else {
4157 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4158 for &slot in &input_slots[1..] {
4159 let v = load_slot(&mut builder, buffer_ptr, slot);
4160 let cmp = builder.ins().icmp(
4161 ir::condcodes::IntCC::UnsignedGreaterThan,
4162 v,
4163 acc,
4164 );
4165 acc = builder.ins().select(cmp, v, acc);
4166 }
4167 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4168 }
4169 }
4170
4171 JitOp::CeilToMultiple => {
4172 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4173 let m = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4174 let zero = builder.ins().iconst(types::I64, 0);
4175 let one = builder.ins().iconst(types::I64, 1);
4176 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, m, zero);
4177 let calc_block = builder.create_block();
4178 let merge_block = builder.create_block();
4179 builder.append_block_param(merge_block, types::I64);
4180 builder
4181 .ins()
4182 .brif(is_zero, merge_block, &[val], calc_block, &[]);
4183 builder.switch_to_block(calc_block);
4184 builder.seal_block(calc_block);
4185 let div = div_ceil(&mut builder, val, m, one);
4188 let high = builder.ins().umulhi(div, m);
4189 let low = builder.ins().imul(div, m);
4190 let zero_hi = builder.ins().iconst(types::I64, 0);
4191 let overflows =
4192 builder
4193 .ins()
4194 .icmp(ir::condcodes::IntCC::NotEqual, high, zero_hi);
4195 let max = builder.ins().iconst(types::I64, -1);
4196 let mul = builder.ins().select(overflows, max, low);
4197 builder.ins().jump(merge_block, &[mul]);
4198 builder.switch_to_block(merge_block);
4199 builder.seal_block(merge_block);
4200 let result = builder.block_params(merge_block)[0];
4201 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4202 }
4203 JitOp::CheckedAdd => {
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 sum = builder.ins().iadd(a, b);
4207 let is_overflow =
4208 builder
4209 .ins()
4210 .icmp(ir::condcodes::IntCC::UnsignedLessThan, sum, a);
4211 let zero = builder.ins().iconst(types::I64, 0);
4212 let result = builder.ins().select(is_overflow, zero, sum);
4213 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4214 }
4215 JitOp::CheckedSub => {
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 is_lt =
4219 builder
4220 .ins()
4221 .icmp(ir::condcodes::IntCC::UnsignedLessThan, a, b);
4222 let diff = builder.ins().isub(a, b);
4223 let zero = builder.ins().iconst(types::I64, 0);
4224 let result = builder.ins().select(is_lt, zero, diff);
4225 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4226 }
4227 JitOp::CheckedMul => {
4228 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4229 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4230 let prod = builder.ins().imul(a, b);
4231 let zero = builder.ins().iconst(types::I64, 0);
4232 let a_is_zero =
4233 builder.ins().icmp(ir::condcodes::IntCC::Equal, a, zero);
4234 let div_block = builder.create_block();
4235 let merge_block = builder.create_block();
4236 builder.append_block_param(merge_block, types::I64);
4237 builder
4238 .ins()
4239 .brif(a_is_zero, merge_block, &[zero], div_block, &[]);
4240 builder.switch_to_block(div_block);
4241 builder.seal_block(div_block);
4242 let div = builder.ins().udiv(prod, a);
4243 let ok = builder.ins().icmp(ir::condcodes::IntCC::Equal, div, b);
4244 let mul_res = builder.ins().select(ok, prod, zero);
4245 builder.ins().jump(merge_block, &[mul_res]);
4246 builder.switch_to_block(merge_block);
4247 builder.seal_block(merge_block);
4248 let result = builder.block_params(merge_block)[0];
4249 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4250 }
4251 JitOp::MultiplesAtLeast => {
4252 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4253 let m = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4254 let zero = builder.ins().iconst(types::I64, 0);
4255 let one = builder.ins().iconst(types::I64, 1);
4256 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, m, zero);
4257 let calc_block = builder.create_block();
4258 let merge_block = builder.create_block();
4259 builder.append_block_param(merge_block, types::I64);
4260 builder
4261 .ins()
4262 .brif(is_zero, merge_block, &[zero], calc_block, &[]);
4263 builder.switch_to_block(calc_block);
4264 builder.seal_block(calc_block);
4265 let div = div_ceil(&mut builder, val, m, one);
4266 builder.ins().jump(merge_block, &[div]);
4267 builder.switch_to_block(merge_block);
4268 builder.seal_block(merge_block);
4269 let result = builder.block_params(merge_block)[0];
4270 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4271 }
4272
4273 JitOp::BlendConst(mix_bits) => {
4274 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4279 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4280 let fa = builder.ins().bitcast(types::F64, ir::MemFlags::new(), a);
4281 let fb = builder.ins().bitcast(types::F64, ir::MemFlags::new(), b);
4282 let mix_f64 = f64::from_bits(*mix_bits);
4283 let mix_val = builder.ins().f64const(mix_f64);
4284 let one = builder.ins().f64const(1.0);
4285 let one_minus_mix = builder.ins().fsub(one, mix_val);
4286 let a_part = builder.ins().fmul(fa, one_minus_mix);
4287 let b_part = builder.ins().fmul(fb, mix_val);
4288 let sum = builder.ins().fadd(a_part, b_part);
4289 let result =
4290 builder.ins().bitcast(types::I64, ir::MemFlags::new(), sum);
4291 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4292 }
4293 JitOp::LfsrStepConst(feedback) => {
4294 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4295 let feedback = builder.ins().iconst(types::I64, *feedback as i64);
4296 let one = builder.ins().iconst(types::I64, 1);
4297 let zero = builder.ins().iconst(types::I64, 0);
4298 let shifted = builder.ins().ushr(val, one);
4299 let lsb = builder.ins().band(val, one);
4300 let is_odd =
4301 builder
4302 .ins()
4303 .icmp(ir::condcodes::IntCC::NotEqual, lsb, zero);
4304 let fb_mask = builder.ins().select(is_odd, feedback, zero);
4305 let result = builder.ins().bxor(shifted, fb_mask);
4306 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4307 }
4308 JitOp::PcgConst(seed, stream) => {
4309 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4310 let s = builder.ins().iconst(types::I64, *seed as i64);
4311 let st = builder.ins().iconst(types::I64, *stream as i64);
4312 let call = builder.ins().call(pcg_func_ref, &[input, s, st]);
4313 let result = builder.inst_results(call)[0];
4314 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4315 }
4316 JitOp::PcgStreamConst(seed) => {
4317 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4318 let st = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4319 let s = builder.ins().iconst(types::I64, *seed as i64);
4320 let call = builder.ins().call(pcg_stream_func_ref, &[input, st, s]);
4321 let result = builder.inst_results(call)[0];
4322 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4323 }
4324 JitOp::NOfConst(n, m) => {
4325 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4326 let n_val = builder.ins().iconst(types::I64, *n as i64);
4327 let m_val = builder.ins().iconst(types::I64, *m as i64);
4328 let call = builder.ins().call(n_of_func_ref, &[input, n_val, m_val]);
4329 let result = builder.inst_results(call)[0];
4330 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4331 }
4332
4333 JitOp::SlotCall { kit, scratch_base } => {
4334 emit_slot_call(
4335 &mut builder,
4336 buffer_ptr,
4337 scratch_ptr,
4338 slot_call_ref,
4339 kit,
4340 *scratch_base,
4341 input_slots,
4342 output_slots,
4343 );
4344 }
4345
4346 JitOp::Convert {
4347 from,
4348 to,
4349 kit,
4350 scratch_base,
4351 } => {
4352 emit_conversion(
4353 &mut builder,
4354 buffer_ptr,
4355 input_slots[0],
4356 output_slots[0],
4357 *from,
4358 *to,
4359 |builder| {
4360 emit_slot_call(
4361 builder,
4362 buffer_ptr,
4363 scratch_ptr,
4364 slot_call_ref,
4365 kit,
4366 *scratch_base,
4367 input_slots,
4368 output_slots,
4369 )
4370 },
4371 );
4372 }
4373
4374 JitOp::U64ToStr { scratch_base }
4375 | JitOp::I64ToStr { scratch_base }
4376 | JitOp::F64ToStr { scratch_base } => {
4377 let func = match jit_op {
4380 JitOp::U64ToStr { .. } => u64_to_str_ref,
4381 JitOp::I64ToStr { .. } => i64_to_str_ref,
4382 _ => f64_to_str_ref,
4383 };
4384 let value = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4385 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4386 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4387 builder
4388 .ins()
4389 .call(func, &[scratch_ptr, base_v, buffer_ptr, out_v, value]);
4390 }
4391 JitOp::JsonToStr { scratch_base } => {
4392 let ptr = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4393 let len = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4394 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4395 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4396 builder.ins().call(
4397 json_to_str_ref,
4398 &[scratch_ptr, base_v, buffer_ptr, out_v, ptr, len],
4399 );
4400 }
4401 JitOp::StrConcat { scratch_base } => {
4402 let n_words = input_slots.len();
4405 let frame = builder.create_sized_stack_slot(ir::StackSlotData::new(
4406 ir::StackSlotKind::ExplicitSlot,
4407 (n_words.max(1) * 8) as u32,
4408 3,
4409 ));
4410 for (k, &s) in input_slots.iter().enumerate() {
4411 let v = load_slot(&mut builder, buffer_ptr, s);
4412 builder.ins().stack_store(v, frame, (k * 8) as i32);
4413 }
4414 let pairs_ptr = builder.ins().stack_addr(types::I64, frame, 0);
4415 let n_v = builder.ins().iconst(types::I64, (n_words / 2) as i64);
4416 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4417 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4418 builder.ins().call(
4419 str_concat_ref,
4420 &[scratch_ptr, base_v, buffer_ptr, out_v, pairs_ptr, n_v],
4421 );
4422 }
4423
4424 JitOp::VecProduce { kind, scratch_base } => {
4425 let func = func_of(&vec_producer_refs, *kind);
4429 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4430 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4431 let words = load_words(&mut builder, buffer_ptr, input_slots, 4);
4432 let mut args = vec![scratch_ptr, base_v, buffer_ptr, out_v];
4433 args.extend(words);
4434 builder.ins().call(func, &args);
4435 }
4436 JitOp::VecReduce(kind) => {
4437 let func = func_of(&vec_reducer_refs, *kind);
4438 let words = load_words(&mut builder, buffer_ptr, input_slots, 4);
4439 let call = builder.ins().call(func, &words);
4440 let bits = builder.inst_results(call)[0];
4441 store_slot(&mut builder, buffer_ptr, output_slots[0], bits);
4442 }
4443 JitOp::RegLane(kind) => {
4444 let func = func_of(®_lane_refs, *kind);
4445 let words = load_words(&mut builder, buffer_ptr, input_slots, 3);
4446 let call = builder.ins().call(func, &words);
4447 let word = builder.inst_results(call)[0];
4448 store_slot(&mut builder, buffer_ptr, output_slots[0], word);
4449 }
4450 JitOp::RegProduce(kind) => {
4451 let func = func_of(®_producer_refs, *kind);
4452 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4453 let words = load_words(&mut builder, buffer_ptr, input_slots, 4);
4454 let mut args = vec![buffer_ptr, out_v];
4455 args.extend(words);
4456 builder.ins().call(func, &args);
4457 }
4458 JitOp::RegDotF32 => {
4459 let a =
4463 load_reg128(&mut builder, buffer_ptr, input_slots[0], types::F32X4);
4464 let b =
4465 load_reg128(&mut builder, buffer_ptr, input_slots[2], types::F32X4);
4466 let p = builder.ins().fmul(a, b);
4467 let p0 = builder.ins().extractlane(p, 0);
4468 let p1 = builder.ins().extractlane(p, 1);
4469 let p2 = builder.ins().extractlane(p, 2);
4470 let p3 = builder.ins().extractlane(p, 3);
4471 let s01 = builder.ins().fadd(p0, p1);
4472 let s23 = builder.ins().fadd(p2, p3);
4473 let s = builder.ins().fadd(s01, s23);
4474 let wide = builder.ins().fpromote(types::F64, s);
4475 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], wide);
4476 }
4477 JitOp::RegShuffleConst(mask) => {
4478 let x =
4482 load_reg128(&mut builder, buffer_ptr, input_slots[0], types::I8X16);
4483 let imm = builder
4484 .func
4485 .dfg
4486 .immediates
4487 .push(ir::ConstantData::from(&mask[..]));
4488 let r = builder.ins().shuffle(x, x, imm);
4489 store_reg128(&mut builder, buffer_ptr, output_slots[0], r);
4490 }
4491
4492 JitOp::Fallback => {
4493 }
4496 }
4497 if let Some((inst, mark)) = tracker_store {
4498 let calls = (mark..builder.func.dfg.num_insts()).any(|i| {
4499 builder.func.dfg.insts[ir::Inst::from_u32(i as u32)]
4500 .opcode()
4501 .is_call()
4502 });
4503 if !calls {
4504 builder.func.layout.remove_inst(inst);
4505 }
4506 }
4507 }
4508 if let Some((at, head, _, _, clean_ptr)) = &dispatcher {
4511 let one = builder.ins().iconst(types::I8, 1);
4512 builder
4513 .ins()
4514 .store(ir::MemFlags::trusted(), one, *clean_ptr, unit_idx as i32);
4515 let i = builder.ins().stack_load(types::I64, *at, 0);
4516 let next = builder.ins().iadd_imm(i, 1);
4517 builder.ins().stack_store(next, *at, 0);
4518 builder.ins().jump(*head, &[]);
4519 }
4520 }
4521
4522 if let Some((_, head, exit, _, _)) = dispatcher {
4523 builder.seal_block(head);
4526 builder.switch_to_block(exit);
4527 builder.seal_block(exit);
4528 }
4529 builder.ins().return_(&[]);
4530 builder.finalize();
4531 }
4532 let fallible = ctx.func.layout.blocks().any(|block| {
4535 ctx.func
4536 .layout
4537 .block_insts(block)
4538 .any(|inst| ctx.func.dfg.insts[inst].opcode().is_call())
4539 });
4540
4541 module
4542 .define_function(func_id, &mut ctx)
4543 .map_err(|e| format!("define function: {e}"))?;
4544 module.clear_context(&mut ctx);
4545 defined.push((func_id, fallible));
4546 }
4547 module
4548 .finalize_definitions()
4549 .map_err(|e| format!("finalize: {e}"))?;
4550
4551 let entries: Vec<JitEntry> = defined
4552 .iter()
4553 .map(|&(func_id, fallible)| {
4554 let code_ptr = module.get_finalized_function(func_id);
4555 let straight_fn: NativeFn = unsafe { mem::transmute(code_ptr) };
4556 let dispatch_fn: NativeDispatchFn = unsafe { mem::transmute(code_ptr) };
4557 (straight_fn, dispatch_fn, fallible)
4558 })
4559 .collect();
4560 let kits: Vec<SlotKitRef> = functions
4562 .iter()
4563 .flat_map(|(steps, _)| steps.iter())
4564 .filter_map(|(op, _, _)| op.slot_kit().cloned())
4565 .collect();
4566 let any_fallible = defined.iter().any(|&(_, f)| f);
4567 let code = super::kernels::JitCode::new(module, kits, any_fallible);
4568 Ok((entries, code))
4569}
4570
4571#[allow(clippy::too_many_arguments)]
4580fn emit_slot_call(
4581 builder: &mut FunctionBuilder,
4582 buffer_ptr: ir::Value,
4583 scratch_ptr: ir::Value,
4584 slot_call_ref: ir::FuncRef,
4585 kit: &SlotKitRef,
4586 scratch_base: usize,
4587 input_slots: &[usize],
4588 output_slots: &[usize],
4589) {
4590 let n_in = input_slots.len();
4591 let n_out = output_slots.len();
4592 let frame = |builder: &mut FunctionBuilder, n: usize| {
4593 builder.create_sized_stack_slot(ir::StackSlotData::new(
4594 ir::StackSlotKind::ExplicitSlot,
4595 (n.max(1) * 8) as u32,
4596 3,
4597 ))
4598 };
4599 let in_frame = frame(builder, n_in);
4600 let out_frame = frame(builder, n_out);
4601 for (k, &s) in input_slots.iter().enumerate() {
4602 let v = load_slot(builder, buffer_ptr, s);
4603 builder.ins().stack_store(v, in_frame, (k * 8) as i32);
4604 }
4605 let kit_ptr = builder
4606 .ins()
4607 .iconst(types::I64, std::sync::Arc::as_ptr(&kit.0) as usize as i64);
4608 let in_ptr = builder.ins().stack_addr(types::I64, in_frame, 0);
4609 let n_in_v = builder.ins().iconst(types::I64, n_in as i64);
4610 let out_ptr = builder.ins().stack_addr(types::I64, out_frame, 0);
4611 let n_out_v = builder.ins().iconst(types::I64, n_out as i64);
4612 let base_v = builder.ins().iconst(types::I64, scratch_base as i64);
4613 let n_sc_v = builder.ins().iconst(types::I64, kit.0.scratch.len() as i64);
4614 builder.ins().call(
4615 slot_call_ref,
4616 &[
4617 kit_ptr,
4618 in_ptr,
4619 n_in_v,
4620 out_ptr,
4621 n_out_v,
4622 scratch_ptr,
4623 base_v,
4624 n_sc_v,
4625 ],
4626 );
4627 for (k, &s) in output_slots.iter().enumerate() {
4628 let v = builder
4629 .ins()
4630 .stack_load(types::I64, out_frame, (k * 8) as i32);
4631 store_slot(builder, buffer_ptr, s, v);
4632 }
4633}
4634
4635fn emit_conversion(
4648 builder: &mut FunctionBuilder,
4649 buffer_ptr: ir::Value,
4650 in_slot: usize,
4651 out_slot: usize,
4652 from: crate::ast::PortType,
4653 to: crate::ast::PortType,
4654 slow: impl FnOnce(&mut FunctionBuilder),
4655) {
4656 use Scalar::{Bool, F32, F64, Signed, Unsigned};
4657 use ir::condcodes::{FloatCC, IntCC};
4658 let (Some(src), Some(dst)) = (Scalar::of(from), Scalar::of(to)) else {
4659 slow(builder);
4660 return;
4661 };
4662 let raw = load_slot(builder, buffer_ptr, in_slot);
4663
4664 let float_of = |builder: &mut FunctionBuilder| -> ir::Value {
4666 match src {
4667 F32 => {
4668 let bits = builder.ins().ireduce(types::I32, raw);
4669 let x = builder.ins().bitcast(types::F32, ir::MemFlags::new(), bits);
4670 builder.ins().fpromote(types::F64, x)
4671 }
4672 _ => builder.ins().bitcast(types::F64, ir::MemFlags::new(), raw),
4673 }
4674 };
4675 let store_float = |builder: &mut FunctionBuilder, x: ir::Value| {
4678 let word = if dst == F32 {
4679 let bits = builder.ins().bitcast(types::I32, ir::MemFlags::new(), x);
4680 builder.ins().uextend(types::I64, bits)
4681 } else {
4682 builder.ins().bitcast(types::I64, ir::MemFlags::new(), x)
4683 };
4684 store_slot(builder, buffer_ptr, out_slot, word);
4685 };
4686
4687 match (src, dst) {
4688 (_, Bool) => {
4691 let truth = match src {
4692 F32 | F64 => {
4693 let x = float_of(builder);
4698 let zero = builder.ins().f64const(0.0);
4699 let nonzero = builder.ins().fcmp(FloatCC::NotEqual, x, zero);
4700 let ordered = builder.ins().fcmp(FloatCC::Ordered, x, x);
4701 builder.ins().band(nonzero, ordered)
4702 }
4703 _ => builder.ins().icmp_imm(IntCC::NotEqual, raw, 0),
4704 };
4705 let word = builder.ins().uextend(types::I64, truth);
4706 store_slot(builder, buffer_ptr, out_slot, word);
4707 }
4708 (s, d) if s.int_range().is_some() && d.int_range().is_some() => {
4711 let (smin, smax) = s.int_range().expect("an integer");
4712 let (dmin, dmax) = d.int_range().expect("an integer");
4713 let mut fits = Vec::new();
4714 if dmin > smin {
4715 fits.push(builder.ins().icmp_imm(
4717 IntCC::SignedGreaterThanOrEqual,
4718 raw,
4719 dmin as i64,
4720 ));
4721 }
4722 if dmax < smax {
4723 let cc = if matches!(s, Signed(_)) {
4724 IntCC::SignedLessThanOrEqual
4725 } else {
4726 IntCC::UnsignedLessThanOrEqual
4727 };
4728 fits.push(builder.ins().icmp_imm(cc, raw, dmax as u64 as i64));
4729 }
4730 branch_on(
4731 builder,
4732 fits,
4733 |b| {
4734 store_slot(b, buffer_ptr, out_slot, raw);
4735 },
4736 slow,
4737 );
4738 }
4739 (s, F32 | F64) if s.int_range().is_some() => {
4741 let ty = if dst == F32 { types::F32 } else { types::F64 };
4742 let x = if matches!(s, Signed(_)) {
4743 builder.ins().fcvt_from_sint(ty, raw)
4744 } else {
4745 builder.ins().fcvt_from_uint(ty, raw)
4746 };
4747 store_float(builder, x);
4748 }
4749 (F32, F64) => {
4750 let x = float_of(builder);
4751 store_float(builder, x);
4752 }
4753 (F64, F32) => {
4754 let x = float_of(builder);
4755 let narrow = builder.ins().fdemote(types::F32, x);
4756 store_float(builder, narrow);
4757 }
4758 (F32 | F64, d) => {
4762 let (lo, hi) = match d {
4763 Unsigned(b) => (0.0, 2f64.powi(b as i32)),
4764 Signed(b) => (-(2f64.powi(b as i32 - 1)), 2f64.powi(b as i32 - 1)),
4765 _ => {
4766 slow(builder);
4767 return;
4768 }
4769 };
4770 let x = float_of(builder);
4771 let lo_v = builder.ins().f64const(lo);
4772 let hi_v = builder.ins().f64const(hi);
4773 let above = builder.ins().fcmp(FloatCC::GreaterThanOrEqual, x, lo_v);
4774 let below = builder.ins().fcmp(FloatCC::LessThan, x, hi_v);
4775 branch_on(
4776 builder,
4777 vec![above, below],
4778 |b| {
4779 let word = if matches!(d, Signed(_)) {
4780 b.ins().fcvt_to_sint_sat(types::I64, x)
4781 } else {
4782 b.ins().fcvt_to_uint_sat(types::I64, x)
4783 };
4784 store_slot(b, buffer_ptr, out_slot, word);
4785 },
4786 slow,
4787 );
4788 }
4789 _ => slow(builder),
4790 }
4791}
4792
4793fn branch_on(
4796 builder: &mut FunctionBuilder,
4797 conds: Vec<ir::Value>,
4798 fast: impl FnOnce(&mut FunctionBuilder),
4799 slow: impl FnOnce(&mut FunctionBuilder),
4800) {
4801 let mut conds = conds.into_iter();
4802 let Some(first) = conds.next() else {
4803 fast(builder);
4804 return;
4805 };
4806 let mut ok = first;
4807 for c in conds {
4808 ok = builder.ins().band(ok, c);
4809 }
4810 let fast_block = builder.create_block();
4811 let slow_block = builder.create_block();
4812 let done = builder.create_block();
4813 builder.ins().brif(ok, fast_block, &[], slow_block, &[]);
4814 builder.switch_to_block(fast_block);
4815 builder.seal_block(fast_block);
4816 fast(builder);
4817 builder.ins().jump(done, &[]);
4818 builder.switch_to_block(slow_block);
4819 builder.seal_block(slow_block);
4820 slow(builder);
4821 builder.ins().jump(done, &[]);
4822 builder.switch_to_block(done);
4823 builder.seal_block(done);
4824}
4825
4826fn load_slot(builder: &mut FunctionBuilder, buffer_ptr: ir::Value, slot: usize) -> ir::Value {
4827 let offset = (slot * 8) as i32;
4828 builder
4829 .ins()
4830 .load(types::I64, ir::MemFlags::trusted(), buffer_ptr, offset)
4831}
4832
4833fn store_slot(
4835 builder: &mut FunctionBuilder,
4836 buffer_ptr: ir::Value,
4837 slot: usize,
4838 value: ir::Value,
4839) -> ir::Inst {
4840 let offset = (slot * 8) as i32;
4841 builder
4842 .ins()
4843 .store(ir::MemFlags::trusted(), value, buffer_ptr, offset)
4844}
4845
4846fn reg_lane_type(lane: u8) -> ir::Type {
4849 match lane {
4850 0 => types::I8X16,
4851 1 => types::I16X8,
4852 2 => types::I32X4,
4853 3 => types::I64X2,
4854 4 => types::F32X4,
4855 5 => types::F64X2,
4856 _ => unreachable!("register lane index out of range"),
4857 }
4858}
4859
4860fn load_reg128(
4865 builder: &mut FunctionBuilder,
4866 buffer_ptr: ir::Value,
4867 first_slot: usize,
4868 vt: ir::Type,
4869) -> ir::Value {
4870 let offset = (first_slot * 8) as i32;
4871 builder
4872 .ins()
4873 .load(vt, ir::MemFlags::new(), buffer_ptr, offset)
4874}
4875
4876fn store_reg128(
4878 builder: &mut FunctionBuilder,
4879 buffer_ptr: ir::Value,
4880 first_slot: usize,
4881 value: ir::Value,
4882) {
4883 let offset = (first_slot * 8) as i32;
4884 builder
4885 .ins()
4886 .store(ir::MemFlags::new(), value, buffer_ptr, offset);
4887}
4888
4889fn round_half_away(builder: &mut FunctionBuilder, x: ir::Value) -> ir::Value {
4894 let t = builder.ins().trunc(x);
4895 let frac = builder.ins().fsub(x, t);
4896 let mag = builder.ins().fabs(frac);
4897 let half = builder.ins().f64const(0.5);
4898 let reaches = builder
4899 .ins()
4900 .fcmp(ir::condcodes::FloatCC::GreaterThanOrEqual, mag, half);
4901 let one = builder.ins().f64const(1.0);
4902 let step = builder.ins().fcopysign(one, x);
4903 let up = builder.ins().fadd(t, step);
4904 builder.ins().select(reaches, up, t)
4905}
4906
4907fn clamp_ir(
4912 builder: &mut FunctionBuilder,
4913 x: ir::Value,
4914 lo: ir::Value,
4915 hi: ir::Value,
4916) -> ir::Value {
4917 let below = builder.ins().fcmp(ir::condcodes::FloatCC::LessThan, x, lo);
4918 let above = builder
4919 .ins()
4920 .fcmp(ir::condcodes::FloatCC::GreaterThan, x, hi);
4921 let capped = builder.ins().select(above, hi, x);
4922 builder.ins().select(below, lo, capped)
4923}
4924
4925fn div_ceil(
4928 builder: &mut FunctionBuilder,
4929 val: ir::Value,
4930 m: ir::Value,
4931 one: ir::Value,
4932) -> ir::Value {
4933 let q = builder.ins().udiv(val, m);
4934 let r = builder.ins().urem(val, m);
4935 let zero = builder.ins().iconst(types::I64, 0);
4936 let inexact = builder.ins().icmp(ir::condcodes::IntCC::NotEqual, r, zero);
4937 let q1 = builder.ins().iadd(q, one);
4938 builder.ins().select(inexact, q1, q)
4939}
4940
4941fn func_of<K: PartialEq + Copy>(refs: &[(K, ir::FuncRef)], key: K) -> ir::FuncRef {
4943 refs.iter()
4944 .find(|(k, _)| *k == key)
4945 .map(|(_, r)| *r)
4946 .expect("every helper of the group is declared")
4947}
4948
4949fn load_words(
4952 builder: &mut FunctionBuilder,
4953 buffer_ptr: ir::Value,
4954 input_slots: &[usize],
4955 n: usize,
4956) -> Vec<ir::Value> {
4957 (0..n)
4958 .map(|k| match input_slots.get(k) {
4959 Some(&s) => load_slot(builder, buffer_ptr, s),
4960 None => builder.ins().iconst(types::I64, 0),
4961 })
4962 .collect()
4963}
4964
4965fn load_slot_f64(builder: &mut FunctionBuilder, buffer_ptr: ir::Value, slot: usize) -> ir::Value {
4967 let i64_val = load_slot(builder, buffer_ptr, slot);
4968 builder
4969 .ins()
4970 .bitcast(types::F64, ir::MemFlags::new(), i64_val)
4971}
4972
4973fn store_slot_f64(
4975 builder: &mut FunctionBuilder,
4976 buffer_ptr: ir::Value,
4977 slot: usize,
4978 value: ir::Value,
4979) {
4980 let i64_val = builder
4981 .ins()
4982 .bitcast(types::I64, ir::MemFlags::new(), value);
4983 store_slot(builder, buffer_ptr, slot, i64_val);
4984}
4985
4986#[cfg(test)]
4989mod tests {
4990 use super::*;
4991
4992 #[test]
4997 fn every_scalar_conversion_lowers_by_its_types() {
4998 let mut missing = Vec::new();
4999 for &from in crate::ast::PortType::ALL {
5000 for &to in crate::ast::PortType::ALL {
5001 if from == to || Scalar::of(from).is_none() || Scalar::of(to).is_none() {
5002 continue;
5003 }
5004 let Some(node) = crate::compile::assembly::boundary_adapter(from, to) else {
5005 continue;
5006 };
5007 match classify_node(node.as_ref()) {
5008 JitOp::Convert { from: f, to: t, .. } if f == from && t == to => {}
5009 other => missing.push(format!(
5010 "{} ({from:?} -> {to:?}) classified as {other:?}",
5011 node.meta().name
5012 )),
5013 }
5014 }
5015 }
5016 assert!(missing.is_empty(), "{}", missing.join("\n"));
5017 }
5018
5019 #[test]
5020 fn jit_identity() {
5021 let steps = vec![(JitOp::Identity, vec![0], vec![1])];
5022 let mut output_map = HashMap::new();
5023 output_map.insert("out".into(), 1);
5024 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5025 kernel.eval(&[42]);
5026 assert_eq!(kernel.get("out"), 42);
5027 }
5028
5029 #[test]
5030 fn jit_add_const() {
5031 let steps = vec![(JitOp::AddConst(100), vec![0], vec![1])];
5032 let mut output_map = HashMap::new();
5033 output_map.insert("out".into(), 1);
5034 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5035 kernel.eval(&[5]);
5036 assert_eq!(kernel.get("out"), 105);
5037 }
5038
5039 #[test]
5040 fn jit_mul_const() {
5041 let steps = vec![(JitOp::MulConst(7), vec![0], vec![1])];
5042 let mut output_map = HashMap::new();
5043 output_map.insert("out".into(), 1);
5044 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5045 kernel.eval(&[6]);
5046 assert_eq!(kernel.get("out"), 42);
5047 }
5048
5049 #[test]
5050 fn jit_mod_const() {
5051 let steps = vec![(JitOp::ModConst(100), vec![0], vec![1])];
5052 let mut output_map = HashMap::new();
5053 output_map.insert("out".into(), 1);
5054 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5055 kernel.eval(&[542]);
5056 assert_eq!(kernel.get("out"), 42);
5057 }
5058
5059 #[test]
5060 fn jit_hash() {
5061 let steps = vec![(JitOp::Hash, vec![0], vec![1])];
5062 let mut output_map = HashMap::new();
5063 output_map.insert("out".into(), 1);
5064 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5065
5066 kernel.eval(&[42]);
5067 let v1 = kernel.get("out");
5068
5069 let expected = xxhash_rust::xxh3::xxh3_64(&42u64.to_le_bytes());
5071 assert_eq!(v1, expected);
5072 }
5073
5074 #[test]
5075 fn jit_hash_deterministic() {
5076 let steps = vec![(JitOp::Hash, vec![0], vec![1])];
5077 let mut output_map = HashMap::new();
5078 output_map.insert("out".into(), 1);
5079 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5080
5081 kernel.eval(&[42]);
5082 let v1 = kernel.get("out");
5083 kernel.eval(&[42]);
5084 let v2 = kernel.get("out");
5085 assert_eq!(v1, v2);
5086 }
5087
5088 #[test]
5089 fn jit_chain_hash_mod() {
5090 let steps = vec![
5092 (JitOp::Hash, vec![0], vec![1]), (JitOp::ModConst(1_000_000), vec![1], vec![2]), ];
5095 let mut output_map = HashMap::new();
5096 output_map.insert("user_id".into(), 2);
5097 let mut kernel = compile_jit_raw(1, 3, steps, output_map, Vec::new()).unwrap();
5098
5099 kernel.eval(&[42]);
5100 let uid = kernel.get("user_id");
5101 assert!(uid < 1_000_000, "got {uid}");
5102 }
5103
5104 #[test]
5105 fn jit_clamp_const() {
5106 let steps = vec![(JitOp::ClampConst(10, 50), vec![0], vec![1])];
5107 let mut output_map = HashMap::new();
5108 output_map.insert("out".into(), 1);
5109 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5110
5111 kernel.eval(&[5]);
5112 assert_eq!(kernel.get("out"), 10); kernel.eval(&[30]);
5115 assert_eq!(kernel.get("out"), 30); kernel.eval(&[100]);
5118 assert_eq!(kernel.get("out"), 50); }
5120
5121 #[test]
5122 fn jit_interleave() {
5123 let steps = vec![(JitOp::Interleave, vec![0, 1], vec![2])];
5124 let mut output_map = HashMap::new();
5125 output_map.insert("out".into(), 2);
5126 let mut kernel = compile_jit_raw(2, 3, steps, output_map, Vec::new()).unwrap();
5127
5128 kernel.eval(&[0b101, 0b010]);
5129 assert_eq!(kernel.get("out"), 0b01_10_01);
5131 }
5132
5133 #[test]
5134 fn jit_mixed_radix() {
5135 let steps = vec![(
5137 JitOp::MixedRadixConst(vec![100, 1000, 0]),
5138 vec![0],
5139 vec![1, 2, 3],
5140 )];
5141 let mut output_map = HashMap::new();
5142 output_map.insert("d0".into(), 1);
5143 output_map.insert("d1".into(), 2);
5144 output_map.insert("d2".into(), 3);
5145 let mut kernel = compile_jit_raw(1, 4, steps, output_map, Vec::new()).unwrap();
5146
5147 kernel.eval(&[4_201_337]);
5149 assert_eq!(kernel.get("d0"), 37);
5150 assert_eq!(kernel.get("d1"), 13);
5151 assert_eq!(kernel.get("d2"), 42);
5152 }
5153
5154 #[test]
5155 fn jit_unit_interval() {
5156 let steps = vec![(JitOp::UnitInterval, vec![0], vec![1])];
5157 let mut output_map = HashMap::new();
5158 output_map.insert("out".into(), 1);
5159 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5160
5161 kernel.eval(&[0]);
5162 let v = f64::from_bits(kernel.get("out"));
5163 assert!((v - 0.0).abs() < 1e-10);
5164
5165 kernel.eval(&[u64::MAX]);
5166 let v = f64::from_bits(kernel.get("out"));
5167 assert!((v - 1.0).abs() < 1e-10);
5168 }
5169
5170 #[test]
5171 fn jit_f64_to_u64() {
5172 let steps = vec![(JitOp::F64ToU64, vec![0], vec![1])];
5174 let mut output_map = HashMap::new();
5175 output_map.insert("out".into(), 1);
5176 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5177
5178 kernel.eval(&[3.7f64.to_bits()]);
5179 assert_eq!(kernel.get("out"), 3); }
5181
5182 #[test]
5183 fn jit_round_to_u64() {
5184 let steps = vec![(JitOp::RoundToU64, vec![0], vec![1])];
5185 let mut output_map = HashMap::new();
5186 output_map.insert("out".into(), 1);
5187 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5188
5189 kernel.eval(&[3.7f64.to_bits()]);
5190 assert_eq!(kernel.get("out"), 4);
5191
5192 kernel.eval(&[3.2f64.to_bits()]);
5193 assert_eq!(kernel.get("out"), 3);
5194 }
5195
5196 #[test]
5197 fn jit_clamp_f64() {
5198 let steps = vec![(
5199 JitOp::ClampF64Const(0.0f64.to_bits(), 1.0f64.to_bits()),
5200 vec![0],
5201 vec![1],
5202 )];
5203 let mut output_map = HashMap::new();
5204 output_map.insert("out".into(), 1);
5205 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5206
5207 kernel.eval(&[(-0.5f64).to_bits()]);
5208 assert_eq!(f64::from_bits(kernel.get("out")), 0.0);
5209
5210 kernel.eval(&[0.5f64.to_bits()]);
5211 assert_eq!(f64::from_bits(kernel.get("out")), 0.5);
5212
5213 kernel.eval(&[1.5f64.to_bits()]);
5214 assert_eq!(f64::from_bits(kernel.get("out")), 1.0);
5215 }
5216
5217 #[test]
5218 fn jit_lerp() {
5219 let steps = vec![(
5220 JitOp::LerpConst(10.0f64.to_bits(), 20.0f64.to_bits()),
5221 vec![0],
5222 vec![1],
5223 )];
5224 let mut output_map = HashMap::new();
5225 output_map.insert("out".into(), 1);
5226 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5227
5228 kernel.eval(&[0.0f64.to_bits()]);
5229 assert_eq!(f64::from_bits(kernel.get("out")), 10.0);
5230
5231 kernel.eval(&[1.0f64.to_bits()]);
5232 assert_eq!(f64::from_bits(kernel.get("out")), 20.0);
5233
5234 kernel.eval(&[0.5f64.to_bits()]);
5235 assert_eq!(f64::from_bits(kernel.get("out")), 15.0);
5236 }
5237
5238 #[test]
5239 fn jit_scale_range() {
5240 let steps = vec![(
5241 JitOp::ScaleRangeConst(10.0f64.to_bits(), 10.0f64.to_bits()),
5242 vec![0],
5243 vec![1],
5244 )];
5245 let mut output_map = HashMap::new();
5246 output_map.insert("out".into(), 1);
5247 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5248
5249 kernel.eval(&[0]);
5250 let v = f64::from_bits(kernel.get("out"));
5251 assert!((v - 10.0).abs() < 0.001);
5252
5253 kernel.eval(&[u64::MAX]);
5254 let v = f64::from_bits(kernel.get("out"));
5255 assert!((v - 20.0).abs() < 0.001);
5256 }
5257
5258 #[test]
5259 fn jit_quantize() {
5260 let steps = vec![(JitOp::QuantizeConst(10.0f64.to_bits()), vec![0], vec![1])];
5261 let mut output_map = HashMap::new();
5262 output_map.insert("out".into(), 1);
5263 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5264
5265 kernel.eval(&[13.0f64.to_bits()]);
5266 assert_eq!(f64::from_bits(kernel.get("out")), 10.0);
5267
5268 kernel.eval(&[17.0f64.to_bits()]);
5269 assert_eq!(f64::from_bits(kernel.get("out")), 20.0);
5270 }
5271
5272 #[test]
5273 fn jit_discretize() {
5274 let steps = vec![(
5275 JitOp::DiscretizeConst(100.0f64.to_bits(), 10),
5276 vec![0],
5277 vec![1],
5278 )];
5279 let mut output_map = HashMap::new();
5280 output_map.insert("out".into(), 1);
5281 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5282
5283 kernel.eval(&[0.0f64.to_bits()]);
5284 assert_eq!(kernel.get("out"), 0);
5285
5286 kernel.eval(&[55.0f64.to_bits()]);
5287 assert_eq!(kernel.get("out"), 5);
5288
5289 kernel.eval(&[99.0f64.to_bits()]);
5290 assert_eq!(kernel.get("out"), 9);
5291
5292 kernel.eval(&[200.0f64.to_bits()]);
5294 assert_eq!(kernel.get("out"), 9);
5295 }
5296
5297 #[test]
5298 fn jit_chain_unit_interval_lerp() {
5299 let steps = vec![
5301 (JitOp::UnitInterval, vec![0], vec![1]),
5302 (
5303 JitOp::LerpConst(100.0f64.to_bits(), 200.0f64.to_bits()),
5304 vec![1],
5305 vec![2],
5306 ),
5307 ];
5308 let mut output_map = HashMap::new();
5309 output_map.insert("out".into(), 2);
5310 let mut kernel = compile_jit_raw(1, 3, steps, output_map, Vec::new()).unwrap();
5311
5312 kernel.eval(&[0]);
5313 let v = f64::from_bits(kernel.get("out"));
5314 assert!((v - 100.0).abs() < 0.001);
5315
5316 kernel.eval(&[u64::MAX]);
5317 let v = f64::from_bits(kernel.get("out"));
5318 assert!((v - 200.0).abs() < 0.001);
5319 }
5320
5321 #[test]
5322 fn jit_multi_step_chain() {
5323 let steps = vec![
5325 (JitOp::AddConst(10), vec![0], vec![1]),
5326 (JitOp::MulConst(3), vec![1], vec![2]),
5327 (JitOp::ModConst(100), vec![2], vec![3]),
5328 ];
5329 let mut output_map = HashMap::new();
5330 output_map.insert("out".into(), 3);
5331 let mut kernel = compile_jit_raw(1, 4, steps, output_map, Vec::new()).unwrap();
5332
5333 kernel.eval(&[5]);
5334 assert_eq!(kernel.get("out"), 45);
5336 }
5337
5338 #[test]
5341 fn jit_is_positive_check_passes_positive() {
5342 let steps = vec![(
5343 JitOp::IsPositiveCheck {
5344 name_ptr: 0,
5345 name_len: 0,
5346 },
5347 vec![0],
5348 vec![1],
5349 )];
5350 let mut output_map = HashMap::new();
5351 output_map.insert("out".into(), 1);
5352 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5353 kernel.eval(&[42]);
5354 assert_eq!(kernel.get("out"), 42);
5355 kernel.eval(&[u64::MAX]);
5358 assert_eq!(kernel.get("out"), u64::MAX);
5359 }
5360
5361 #[test]
5362 fn jit_in_range_check_passes_interior() {
5363 let steps = vec![(JitOp::InRangeCheck(10, 100), vec![0], vec![1])];
5364 let mut output_map = HashMap::new();
5365 output_map.insert("out".into(), 1);
5366 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5367 kernel.eval(&[50]);
5368 assert_eq!(kernel.get("out"), 50);
5369 kernel.eval(&[10]);
5371 assert_eq!(kernel.get("out"), 10);
5372 kernel.eval(&[100]);
5373 assert_eq!(kernel.get("out"), 100);
5374 }
5375
5376 #[test]
5381 fn jit_is_one_of_check_passes_allowed_values() {
5382 let steps = vec![(
5383 JitOp::IsOneOfCheck {
5384 allowed: vec![1, 2, 3, 5, 8],
5385 set_ptr: 0,
5386 set_len: 0,
5387 },
5388 vec![0],
5389 vec![1],
5390 )];
5391 let mut output_map = HashMap::new();
5392 output_map.insert("out".into(), 1);
5393 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5394 for v in [1u64, 2, 3, 5, 8] {
5396 kernel.eval(&[v]);
5397 assert_eq!(kernel.get("out"), v);
5398 }
5399 }
5400
5401 #[test]
5402 fn jit_is_one_of_check_accepts_single_element_allow_list() {
5403 let steps = vec![(
5406 JitOp::IsOneOfCheck {
5407 allowed: vec![42],
5408 set_ptr: 0,
5409 set_len: 0,
5410 },
5411 vec![0],
5412 vec![1],
5413 )];
5414 let mut output_map = HashMap::new();
5415 output_map.insert("out".into(), 1);
5416 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5417 kernel.eval(&[42]);
5418 assert_eq!(kernel.get("out"), 42);
5419 }
5420
5421 fn extract_panic_msg(payload: Box<dyn std::any::Any + Send + 'static>) -> String {
5431 payload
5432 .downcast_ref::<String>()
5433 .cloned()
5434 .or_else(|| payload.downcast_ref::<&str>().map(|s| s.to_string()))
5435 .unwrap_or_else(|| "(non-string panic)".into())
5436 }
5437
5438 #[test]
5439 fn jit_is_positive_violation_is_catchable() {
5440 let steps = vec![(
5441 JitOp::IsPositiveCheck {
5442 name_ptr: 0,
5443 name_len: 0,
5444 },
5445 vec![0],
5446 vec![1],
5447 )];
5448 let mut output_map = HashMap::new();
5449 output_map.insert("out".into(), 1);
5450 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5451 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[0])))
5452 .expect_err("JIT violation should panic");
5453 assert!(extract_panic_msg(err).contains("must be > 0"));
5454 }
5455
5456 #[test]
5457 fn jit_in_range_violation_is_catchable() {
5458 let steps = vec![(JitOp::InRangeCheck(10, 100), vec![0], vec![1])];
5459 let mut output_map = HashMap::new();
5460 output_map.insert("out".into(), 1);
5461 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5462
5463 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[5])))
5464 .expect_err("below-range should panic");
5465 assert!(extract_panic_msg(err).contains("outside [10, 100]"));
5466
5467 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[500])))
5468 .expect_err("above-range should panic");
5469 assert!(extract_panic_msg(err).contains("outside [10, 100]"));
5470 }
5471
5472 #[test]
5473 fn jit_is_one_of_violation_is_catchable() {
5474 let steps = vec![(
5475 JitOp::IsOneOfCheck {
5476 allowed: vec![1, 3, 5],
5477 set_ptr: 0,
5478 set_len: 0,
5479 },
5480 vec![0],
5481 vec![1],
5482 )];
5483 let mut output_map = HashMap::new();
5484 output_map.insert("out".into(), 1);
5485 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5486 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[2])))
5487 .expect_err("disallowed value should panic");
5488 assert!(extract_panic_msg(err).contains("not in allowed set"));
5489 }
5490
5491 #[test]
5492 fn invoke_with_catch_restores_slot_after_foreign_panic() {
5493 let caught = std::panic::catch_unwind(|| {
5501 invoke_with_catch(|| panic!("foreign panic"));
5502 });
5503 assert!(caught.is_err(), "foreign panic should propagate out");
5504
5505 let steps = vec![(
5507 JitOp::IsPositiveCheck {
5508 name_ptr: 0,
5509 name_len: 0,
5510 },
5511 vec![0],
5512 vec![1],
5513 )];
5514 let mut output_map = HashMap::new();
5515 output_map.insert("out".into(), 1);
5516 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5517 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[0])))
5518 .expect_err("JIT violation should panic cleanly after foreign panic");
5519 assert!(extract_panic_msg(err).contains("must be > 0"));
5520
5521 kernel.eval(&[42]);
5523 assert_eq!(kernel.get("out"), 42);
5524 }
5525
5526 #[test]
5527 fn jit_kernel_survives_multiple_violations() {
5528 let steps = vec![(
5532 JitOp::IsPositiveCheck {
5533 name_ptr: 0,
5534 name_len: 0,
5535 },
5536 vec![0],
5537 vec![1],
5538 )];
5539 let mut output_map = HashMap::new();
5540 output_map.insert("out".into(), 1);
5541 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5542
5543 for _ in 0..3 {
5544 let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[0])))
5545 .expect_err("violation should still panic");
5546 }
5547 kernel.eval(&[42]);
5549 assert_eq!(kernel.get("out"), 42);
5550 }
5551}