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))]
452struct JitJmpBuf([u8; 512]);
453
454#[cfg(not(windows))]
455unsafe extern "C" {
456 fn _setjmp(env: *mut JitJmpBuf) -> i32;
457 fn _longjmp(env: *mut JitJmpBuf, val: i32) -> !;
458}
459
460#[cfg(windows)]
473unsafe extern "C" {
474 fn _setjmp(env: *mut JitJmpBuf, frame: *mut std::ffi::c_void) -> i32;
475 #[link_name = "longjmp"]
476 fn _longjmp(env: *mut JitJmpBuf, val: i32) -> !;
477}
478
479use std::cell::{Cell, RefCell};
480thread_local! {
481 static JIT_JMP_BUF: Cell<Option<*mut JitJmpBuf>> = const { Cell::new(None) };
488 static JIT_VIOLATION_MSG: RefCell<Option<String>> = const { RefCell::new(None) };
492}
493
494fn jit_violation_longjmp(msg: String) -> ! {
501 JIT_VIOLATION_MSG.with(|m| *m.borrow_mut() = Some(msg.clone()));
502 let buf_ptr: Option<*mut JitJmpBuf> = JIT_JMP_BUF.with(|b| b.get());
503 match buf_ptr {
504 Some(ptr) => unsafe { _longjmp(ptr, 1) },
505 None => {
506 let mut err = std::io::stderr().lock();
507 use std::io::Write;
508 let _ = writeln!(err, "{msg}");
509 let _ = err.flush();
510 std::process::abort();
511 }
512 }
513}
514
515struct JmpBufGuard {
523 prev: Option<*mut JitJmpBuf>,
524}
525
526impl Drop for JmpBufGuard {
527 fn drop(&mut self) {
528 JIT_JMP_BUF.with(|b| b.set(self.prev));
529 }
530}
531
532pub(crate) fn invoke_with_catch<F: FnOnce()>(f: F) {
543 use std::mem::MaybeUninit;
544 let mut buf: MaybeUninit<JitJmpBuf> = MaybeUninit::uninit();
545 let buf_ptr = buf.as_mut_ptr();
546 let prev: Option<*mut JitJmpBuf> = JIT_JMP_BUF.with(|b| b.replace(Some(buf_ptr)));
551 let _guard = JmpBufGuard { prev };
552 #[cfg(not(windows))]
553 let jmpval = unsafe { _setjmp(buf_ptr) };
554 #[cfg(windows)]
556 let jmpval = unsafe { _setjmp(buf_ptr, std::ptr::null_mut()) };
557 if jmpval == 0 {
558 f();
559 } else {
560 let msg = JIT_VIOLATION_MSG
565 .with(|m| m.borrow_mut().take())
566 .unwrap_or_else(|| "JIT predicate violation (no message)".into());
567 std::panic::resume_unwind(Box::new(msg));
571 }
572}
573
574extern "C" fn jit_is_positive_fail(value: u64, name_ptr: u64, name_len: u64) -> u64 {
578 let name = if name_ptr != 0 {
583 unsafe {
584 std::str::from_utf8_unchecked(std::slice::from_raw_parts(
585 name_ptr as *const u8,
586 name_len as usize,
587 ))
588 }
589 } else {
590 "value"
591 };
592 jit_violation_longjmp(format!(
593 "is_positive({name}): value must be > 0, got {value}"
594 ));
595}
596
597extern "C" fn jit_in_range_fail(value: u64, lo: u64, hi: u64) -> u64 {
600 jit_violation_longjmp(format!("in_range: value {value} outside [{lo}, {hi}]"));
601}
602
603extern "C" fn jit_div_zero_fail(kind: u64) -> u64 {
607 jit_violation_longjmp(
608 if kind == 0 {
609 "attempt to divide by zero"
610 } else {
611 "attempt to calculate the remainder with a divisor of zero"
612 }
613 .to_string(),
614 );
615}
616
617extern "C" fn jit_f64_mod(a_bits: u64, b_bits: u64) -> u64 {
621 let (a, b) = (f64::from_bits(a_bits), f64::from_bits(b_bits));
622 (if b != 0.0 { a % b } else { 0.0 }).to_bits()
623}
624
625extern "C" fn jit_is_one_of_fail(value: u64, set_ptr: u64, set_len: u64) -> u64 {
633 let msg = if set_ptr != 0 {
634 let set = unsafe { std::slice::from_raw_parts(set_ptr as *const u64, set_len as usize) };
635 format!("is_one_of: value {value} not in allowed set {set:?}")
636 } else {
637 format!("is_one_of: value {value} not in allowed set [..]")
638 };
639 jit_violation_longjmp(msg);
640}
641
642extern "C" fn jit_weighted_pick(
647 input: u64,
648 values_ptr: u64,
649 biases_ptr: u64,
650 primaries_ptr: u64,
651 aliases_ptr: u64,
652 n: u64,
653) -> u64 {
654 guarded(|| {
655 let n = n as usize;
656 let slot = (input as usize) % n;
657 let bias_test = ((input >> 32) as f64) / (u32::MAX as f64);
658 unsafe {
659 let biases = std::slice::from_raw_parts(biases_ptr as *const f64, n);
660 let primaries = std::slice::from_raw_parts(primaries_ptr as *const u64, n);
661 let aliases = std::slice::from_raw_parts(aliases_ptr as *const u64, n);
662 let values = std::slice::from_raw_parts(values_ptr as *const u64, n);
663 let index = if bias_test < biases[slot] {
664 primaries[slot]
665 } else {
666 aliases[slot]
667 };
668 values[index as usize]
669 }
670 })
671}
672
673fn guarded<T>(body: impl FnOnce() -> T) -> T {
678 match std::panic::catch_unwind(std::panic::AssertUnwindSafe(body)) {
679 Ok(v) => v,
680 Err(payload) => {
681 let msg = payload
682 .downcast_ref::<String>()
683 .cloned()
684 .or_else(|| payload.downcast_ref::<&str>().map(|s| s.to_string()))
685 .unwrap_or_else(|| "panic in a compiled helper".to_string());
686 jit_violation_longjmp(msg)
687 }
688 }
689}
690
691#[derive(Clone)]
694pub struct SlotKitRef(pub std::sync::Arc<crate::ast::CompiledSlotKit>);
695
696impl SlotKitRef {
697 fn new(kit: crate::ast::CompiledSlotKit) -> Self {
698 SlotKitRef(std::sync::Arc::new(kit))
699 }
700}
701
702impl std::fmt::Debug for SlotKitRef {
703 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
704 write!(
705 f,
706 "SlotKitRef({:p}, {} scratch)",
707 std::sync::Arc::as_ptr(&self.0),
708 self.0.scratch.len()
709 )
710 }
711}
712
713impl PartialEq for SlotKitRef {
714 fn eq(&self, other: &Self) -> bool {
715 std::sync::Arc::ptr_eq(&self.0, &other.0)
716 }
717}
718
719extern "C" fn jit_slot_call(
731 kit: *const crate::ast::CompiledSlotKit,
732 inputs: *const u64,
733 n_in: u64,
734 outputs: *mut u64,
735 n_out: u64,
736 scratch: *mut crate::ast::ScratchBuf,
737 base: u64,
738 n_scratch: u64,
739) {
740 guarded(|| unsafe {
741 let kit = &*kit;
742 let ins = std::slice::from_raw_parts(inputs, n_in as usize);
743 let outs = std::slice::from_raw_parts_mut(outputs, n_out as usize);
744 let sc = std::slice::from_raw_parts_mut(scratch.add(base as usize), n_scratch as usize);
745 (kit.op)(ins, outs, sc)
746 })
747}
748
749impl JitOp {
750 pub(crate) fn slot_kit(&self) -> Option<&SlotKitRef> {
752 match self {
753 JitOp::SlotCall { kit, .. } | JitOp::Convert { kit, .. } => Some(kit),
754 _ => None,
755 }
756 }
757
758 pub(crate) fn scratch_elems(&self) -> &[crate::ast::ScratchElem] {
760 const STR_ENTRY: [crate::ast::ScratchElem; 1] = [crate::ast::ScratchElem::Str];
761 const F32_ENTRY: [crate::ast::ScratchElem; 1] = [crate::ast::ScratchElem::F32];
762 match self {
763 JitOp::SlotCall { kit, .. } | JitOp::Convert { kit, .. } => &kit.0.scratch,
764 JitOp::U64ToStr { .. }
765 | JitOp::I64ToStr { .. }
766 | JitOp::F64ToStr { .. }
767 | JitOp::StrConcat { .. }
768 | JitOp::JsonToStr { .. } => &STR_ENTRY,
769 JitOp::VecProduce { .. } => &F32_ENTRY,
770 _ => &[],
771 }
772 }
773
774 pub(crate) fn place_scratch(&mut self, base: usize) {
777 match self {
778 JitOp::SlotCall { scratch_base, .. }
779 | JitOp::Convert { scratch_base, .. }
780 | JitOp::U64ToStr { scratch_base }
781 | JitOp::I64ToStr { scratch_base }
782 | JitOp::F64ToStr { scratch_base }
783 | JitOp::StrConcat { scratch_base }
784 | JitOp::JsonToStr { scratch_base }
785 | JitOp::VecProduce { scratch_base, .. } => *scratch_base = base,
786 _ => {}
787 }
788 }
789}
790
791#[derive(Debug, Clone, Copy, PartialEq, Eq)]
795pub enum VecProducer {
796 Add,
798 Scale,
800 Norm,
802 HashVec,
804 XxHash3Vec,
806 RegToVec,
808}
809
810#[derive(Debug, Clone, Copy, PartialEq, Eq)]
813pub enum VecReducer {
814 Dot,
816 L2,
818 Cosine,
820 LidMle,
822}
823
824#[derive(Debug, Clone, Copy, PartialEq, Eq)]
827pub enum RegLaneRead {
828 F32,
830 I16,
832 I64,
834}
835
836#[derive(Debug, Clone, Copy, PartialEq, Eq)]
840pub enum RegProducer {
841 WithLaneF32,
843 GatherF32,
845 VecToRegF32,
847 MulI8,
849}
850
851unsafe fn vec_f32_of<'a>(ptr: u64, len: u64) -> &'a [f32] {
857 if len == 0 {
858 &[]
859 } else {
860 unsafe { std::slice::from_raw_parts(ptr as usize as *const f32, len as usize) }
861 }
862}
863
864unsafe fn write_f32_entry(
871 scratch: *mut crate::ast::ScratchBuf,
872 base: u64,
873 buffer: *mut u64,
874 out_slot: u64,
875 f: impl FnOnce(&mut Vec<f32>),
876) {
877 unsafe {
878 let entry = &mut *scratch.add(base as usize);
879 let crate::ast::ScratchBuf::F32(v) = entry else {
880 panic!("a vector lowering's scratch entry is not an f32 vector");
881 };
882 f(v);
883 *buffer.add(out_slot as usize) = v.as_ptr() as usize as u64;
884 *buffer.add(out_slot as usize + 1) = v.len() as u64;
885 }
886}
887
888macro_rules! vec_producer {
892 ($name:ident, |$out:ident, $w0:ident, $w1:ident, $w2:ident, $w3:ident| $body:expr) => {
893 extern "C" fn $name(
894 scratch: *mut crate::ast::ScratchBuf,
895 base: u64,
896 buffer: *mut u64,
897 out_slot: u64,
898 $w0: u64,
899 $w1: u64,
900 $w2: u64,
901 $w3: u64,
902 ) {
903 guarded(|| unsafe {
904 let _ = ($w2, $w3);
905 write_f32_entry(scratch, base, buffer, out_slot, |$out| $body)
906 })
907 }
908 };
909}
910
911vec_producer!(jit_vec_add, |out, a_ptr, a_len, b_ptr, b_len| {
912 let (a, b) = (vec_f32_of(a_ptr, a_len), vec_f32_of(b_ptr, b_len));
913 crate::numeric::vector::check_lens("vec_add", a.len(), b.len());
914 crate::numeric::vector::add_f32_into(a, b, out)
915});
916vec_producer!(jit_vec_scale, |out, a_ptr, a_len, k_bits, _z| {
917 let a = vec_f32_of(a_ptr, a_len);
918 crate::numeric::vector::scale_f32_into(a, f64::from_bits(k_bits) as f32, out)
919});
920vec_producer!(jit_vec_norm, |out, a_ptr, a_len, _y, _z| {
921 crate::numeric::vector::norm_f32_into(vec_f32_of(a_ptr, a_len), out)
922});
923vec_producer!(jit_hash_vec, |out, seed, dim, _y, _z| {
924 crate::numeric::vector::hash_vec_into(seed, dim, out)
925});
926vec_producer!(jit_xxhash3_vec, |out, seed, dim, _y, _z| {
927 crate::numeric::vector::xxhash3_vec_into(seed, dim, out)
928});
929vec_producer!(jit_reg_to_vec_f32, |out, lo, hi, _y, _z| {
930 out.clear();
931 out.extend_from_slice(&crate::ast::Bits128([lo, hi]).lanes_f32())
932});
933
934macro_rules! vec_reducer {
937 ($name:ident, |$w0:ident, $w1:ident, $w2:ident, $w3:ident| $body:expr) => {
938 extern "C" fn $name($w0: u64, $w1: u64, $w2: u64, $w3: u64) -> u64 {
939 guarded(|| unsafe {
940 let _ = ($w2, $w3);
941 let r: f64 = $body;
942 r.to_bits()
943 })
944 }
945 };
946}
947
948vec_reducer!(jit_vec_dot, |a_ptr, a_len, b_ptr, b_len| {
949 let (a, b) = (vec_f32_of(a_ptr, a_len), vec_f32_of(b_ptr, b_len));
950 crate::numeric::vector::check_lens("vec_dot", a.len(), b.len());
951 crate::numeric::vector::dot_f32(a, b) as f64
952});
953vec_reducer!(jit_vec_l2, |a_ptr, a_len, b_ptr, b_len| {
954 let (a, b) = (vec_f32_of(a_ptr, a_len), vec_f32_of(b_ptr, b_len));
955 crate::numeric::vector::check_lens("vec_l2", a.len(), b.len());
956 (crate::numeric::vector::l2sq_f32(a, b) as f64).sqrt()
957});
958vec_reducer!(jit_vec_cosine, |a_ptr, a_len, b_ptr, b_len| {
959 let (a, b) = (vec_f32_of(a_ptr, a_len), vec_f32_of(b_ptr, b_len));
960 crate::numeric::vector::check_lens("vec_cosine", a.len(), b.len());
961 crate::numeric::vector::cosine_f32(a, b)
962});
963vec_reducer!(jit_lid_mle, |d_ptr, d_len, k_bits, _z| {
964 crate::numeric::vector::lid_mle_of(vec_f32_of(d_ptr, d_len), f64::from_bits(k_bits))
965});
966
967extern "C" fn jit_reg_lane_f32(lo: u64, hi: u64, i: u64) -> u64 {
969 guarded(|| crate::numeric::register::lane_f32(crate::ast::Bits128([lo, hi]), i).to_bits())
970}
971
972extern "C" fn jit_reg_lane_i16(lo: u64, hi: u64, i: u64) -> u64 {
974 guarded(|| crate::numeric::register::lane_i16(crate::ast::Bits128([lo, hi]), i) as i64 as u64)
975}
976
977extern "C" fn jit_reg_lane_i64(lo: u64, hi: u64, i: u64) -> u64 {
979 guarded(|| crate::numeric::register::lane_i64(crate::ast::Bits128([lo, hi]), i) as u64)
980}
981
982macro_rules! reg_producer {
985 ($name:ident, |$w0:ident, $w1:ident, $w2:ident, $w3:ident| $body:expr) => {
986 extern "C" fn $name(
987 buffer: *mut u64,
988 out_slot: u64,
989 $w0: u64,
990 $w1: u64,
991 $w2: u64,
992 $w3: u64,
993 ) {
994 guarded(|| unsafe {
995 let _ = ($w2, $w3);
996 let r: crate::ast::Bits128 = $body;
997 *buffer.add(out_slot as usize) = r.0[0];
998 *buffer.add(out_slot as usize + 1) = r.0[1];
999 })
1000 }
1001 };
1002}
1003
1004reg_producer!(jit_reg_with_lane_f32, |lo, hi, i, v_bits| {
1005 crate::numeric::register::with_lane_f32(
1006 crate::ast::Bits128([lo, hi]),
1007 i,
1008 f64::from_bits(v_bits),
1009 )
1010});
1011reg_producer!(jit_reg_gather_f32, |v_ptr, v_len, offset, _z| {
1012 crate::numeric::register::gather_f32(vec_f32_of(v_ptr, v_len), offset)
1013});
1014reg_producer!(jit_vec_to_reg_f32, |v_ptr, v_len, _y, _z| {
1015 crate::numeric::register::to_reg_f32(vec_f32_of(v_ptr, v_len))
1016});
1017reg_producer!(jit_reg_mul_i8, |a_lo, a_hi, b_lo, b_hi| {
1018 crate::numeric::register::mul_i8(
1019 crate::ast::Bits128([a_lo, a_hi]),
1020 crate::ast::Bits128([b_lo, b_hi]),
1021 )
1022});
1023
1024unsafe fn write_str_entry(
1033 scratch: *mut crate::ast::ScratchBuf,
1034 base: u64,
1035 buffer: *mut u64,
1036 out_slot: u64,
1037 f: impl FnOnce(&mut Vec<u8>),
1038) {
1039 unsafe {
1040 let entry = &mut *scratch.add(base as usize);
1041 let crate::ast::ScratchBuf::Str(v) = entry else {
1042 panic!("a string lowering's scratch entry is not a string");
1043 };
1044 v.clear();
1045 f(v);
1046 *buffer.add(out_slot as usize) = v.as_ptr() as usize as u64;
1047 *buffer.add(out_slot as usize + 1) = v.len() as u64;
1048 }
1049}
1050
1051extern "C" fn jit_u64_to_str(
1053 scratch: *mut crate::ast::ScratchBuf,
1054 base: u64,
1055 buffer: *mut u64,
1056 out_slot: u64,
1057 value: u64,
1058) {
1059 use std::io::Write;
1060 guarded(|| unsafe {
1061 write_str_entry(scratch, base, buffer, out_slot, |v| {
1062 write!(v, "{value}").expect("a vector accepts every write")
1063 })
1064 })
1065}
1066
1067extern "C" fn jit_i64_to_str(
1069 scratch: *mut crate::ast::ScratchBuf,
1070 base: u64,
1071 buffer: *mut u64,
1072 out_slot: u64,
1073 value: u64,
1074) {
1075 use std::io::Write;
1076 guarded(|| unsafe {
1077 write_str_entry(scratch, base, buffer, out_slot, |v| {
1078 write!(v, "{}", value as i64).expect("a vector accepts every write")
1079 })
1080 })
1081}
1082
1083extern "C" fn jit_f64_to_str(
1086 scratch: *mut crate::ast::ScratchBuf,
1087 base: u64,
1088 buffer: *mut u64,
1089 out_slot: u64,
1090 bits: u64,
1091) {
1092 use std::io::Write;
1093 guarded(|| unsafe {
1094 write_str_entry(scratch, base, buffer, out_slot, |v| {
1095 write!(v, "{}", f64::from_bits(bits)).expect("a vector accepts every write")
1096 })
1097 })
1098}
1099
1100extern "C" fn jit_str_concat(
1104 scratch: *mut crate::ast::ScratchBuf,
1105 base: u64,
1106 buffer: *mut u64,
1107 out_slot: u64,
1108 pairs: *const u64,
1109 n: u64,
1110) {
1111 guarded(|| unsafe {
1112 let words = std::slice::from_raw_parts(pairs, 2 * n as usize);
1113 write_str_entry(scratch, base, buffer, out_slot, |v| {
1114 for pair in words.as_chunks::<2>().0 {
1115 let bytes =
1119 std::slice::from_raw_parts(pair[0] as usize as *const u8, pair[1] as usize);
1120 v.extend_from_slice(bytes);
1121 }
1122 })
1123 })
1124}
1125
1126extern "C" fn jit_json_to_str(
1129 scratch: *mut crate::ast::ScratchBuf,
1130 base: u64,
1131 buffer: *mut u64,
1132 out_slot: u64,
1133 ptr: u64,
1134 len: u64,
1135) {
1136 guarded(|| unsafe {
1137 let pair = [ptr, len];
1138 let value = crate::derive_support::ref_value(&pair);
1139 let json = match value {
1140 crate::ast::Value::Json(j) => j.as_ref(),
1141 other => panic!("expected Json wire, got {other:?}"),
1142 };
1143 write_str_entry(scratch, base, buffer, out_slot, |v| {
1144 serde_json::to_writer(v, json).expect("a vector accepts every write")
1145 })
1146 })
1147}
1148
1149#[derive(Debug, Clone, Copy, PartialEq)]
1156enum Scalar {
1157 Unsigned(u32),
1158 Signed(u32),
1159 Bool,
1160 F32,
1161 F64,
1162}
1163
1164impl Scalar {
1165 fn of(t: crate::ast::PortType) -> Option<Self> {
1166 use crate::ast::PortType as P;
1167 Some(match t {
1168 P::U8 => Self::Unsigned(8),
1169 P::U16 => Self::Unsigned(16),
1170 P::U32 => Self::Unsigned(32),
1171 P::U64 => Self::Unsigned(64),
1172 P::I8 => Self::Signed(8),
1173 P::I16 => Self::Signed(16),
1174 P::I32 => Self::Signed(32),
1175 P::I64 => Self::Signed(64),
1176 P::Bool => Self::Bool,
1177 P::F32 => Self::F32,
1178 P::F64 => Self::F64,
1179 _ => return None,
1180 })
1181 }
1182
1183 fn int_range(self) -> Option<(i128, i128)> {
1185 match self {
1186 Self::Unsigned(b) => Some((0, (1i128 << b) - 1)),
1187 Self::Signed(b) => Some((-(1i128 << (b - 1)), (1i128 << (b - 1)) - 1)),
1188 Self::Bool => Some((0, 1)),
1189 Self::F32 | Self::F64 => None,
1190 }
1191 }
1192}
1193
1194fn conversion_op(node: &dyn PolydatNode) -> Option<JitOp> {
1205 use crate::ast::Slot;
1206 let meta = node.meta();
1207 let [Slot::Wire(input)] = meta.ins.as_slice() else {
1208 return None;
1209 };
1210 let [output] = meta.outs.as_slice() else {
1211 return None;
1212 };
1213 let (from, to) = (input.typ, output.typ);
1214 Scalar::of(from)?;
1215 Scalar::of(to)?;
1216 let canonical = crate::compile::assembly::boundary_adapter(from, to)?;
1217 if canonical.meta().name != meta.name {
1218 return None;
1219 }
1220 let op = node.compiled_u64()?;
1221 Some(JitOp::Convert {
1222 from,
1223 to,
1224 kit: SlotKitRef::new(crate::ast::CompiledSlotKit {
1225 scratch: Vec::new(),
1226 op: Box::new(move |inputs, outputs, _| op(inputs, outputs)),
1227 }),
1228 scratch_base: 0,
1229 })
1230}
1231
1232pub fn classify_node_typed(node: &dyn PolydatNode, wire_types: &[crate::ast::PortType]) -> JitOp {
1241 use crate::ast::PortType as PT;
1242 let is_ref = |t: &crate::ast::PortType| t.slot_color() == crate::ast::SlotColor::Ref2;
1243 let vec_produce = |kind: VecProducer| JitOp::VecProduce {
1244 kind,
1245 scratch_base: 0,
1246 };
1247 let ref_copy = |ty: crate::ast::PortType| {
1248 crate::compile::assembly::ref_copy_kit(ty)
1249 .map(|kit| JitOp::SlotCall {
1250 kit: SlotKitRef::new(kit),
1251 scratch_base: 0,
1252 })
1253 .unwrap_or(JitOp::Fallback)
1254 };
1255 let meta = node.meta();
1256 let named = match meta.name.as_str() {
1257 "__u64_to_string" => JitOp::U64ToStr { scratch_base: 0 },
1260 "__i64_to_string" => JitOp::I64ToStr { scratch_base: 0 },
1261 "__f64_to_string" => JitOp::F64ToStr { scratch_base: 0 },
1262 "json_to_str" if wire_types == [crate::ast::PortType::Json] => {
1263 JitOp::JsonToStr { scratch_base: 0 }
1264 }
1265 "str_concat"
1266 if !wire_types.is_empty()
1267 && wire_types.iter().all(|t| *t == crate::ast::PortType::Str) =>
1268 {
1269 JitOp::StrConcat { scratch_base: 0 }
1270 }
1271 "vec_add" if wire_types == [PT::VecF32, PT::VecF32] => vec_produce(VecProducer::Add),
1276 "vec_scale" if wire_types == [PT::VecF32, PT::F64] => vec_produce(VecProducer::Scale),
1277 "vec_norm" if wire_types == [PT::VecF32] => vec_produce(VecProducer::Norm),
1278 "hash_vec" if wire_types == [PT::U64, PT::U64] => vec_produce(VecProducer::HashVec),
1279 "xxhash3_vec" if wire_types == [PT::U64, PT::U64] => vec_produce(VecProducer::XxHash3Vec),
1280 "reg_to_vec_f32" if wire_types == [PT::RegF32x4] => vec_produce(VecProducer::RegToVec),
1281 "vec_dot" if wire_types == [PT::VecF32, PT::VecF32] => JitOp::VecReduce(VecReducer::Dot),
1282 "vec_l2" if wire_types == [PT::VecF32, PT::VecF32] => JitOp::VecReduce(VecReducer::L2),
1283 "vec_cosine" if wire_types == [PT::VecF32, PT::VecF32] => {
1284 JitOp::VecReduce(VecReducer::Cosine)
1285 }
1286 "lid_mle" if wire_types == [PT::VecF32, PT::F64] => JitOp::VecReduce(VecReducer::LidMle),
1287 "reg_lane_f32" if wire_types == [PT::RegF32x4, PT::U64] => JitOp::RegLane(RegLaneRead::F32),
1291 "reg_lane_i16" if wire_types == [PT::RegI16x8, PT::U64] => JitOp::RegLane(RegLaneRead::I16),
1292 "reg_lane_i64" if wire_types == [PT::RegI64x2, PT::U64] => JitOp::RegLane(RegLaneRead::I64),
1293 "reg_with_lane_f32" if wire_types == [PT::RegF32x4, PT::U64, PT::F64] => {
1294 JitOp::RegProduce(RegProducer::WithLaneF32)
1295 }
1296 "reg_gather_f32" if wire_types == [PT::VecF32, PT::U64] => {
1297 JitOp::RegProduce(RegProducer::GatherF32)
1298 }
1299 "vec_to_reg_f32" if wire_types == [PT::VecF32] => {
1300 JitOp::RegProduce(RegProducer::VecToRegF32)
1301 }
1302 "reg_mul_i8" if wire_types == [PT::RegI8x16, PT::RegI8x16] => {
1303 JitOp::RegProduce(RegProducer::MulI8)
1304 }
1305 "reg_dot_f32" if wire_types == [PT::RegF32x4, PT::RegF32x4] => JitOp::RegDotF32,
1306 n if n.starts_with("__port_") || n == "default_or" => match meta.outs.first() {
1312 Some(o) if is_ref(&o.typ) => return ref_copy(o.typ),
1313 _ => JitOp::Identity,
1314 },
1315 "select" | "select_u64" if wire_types.iter().skip(1).any(|t| t.slot_width() != 1) => {
1318 JitOp::Fallback
1319 }
1320 _ if wire_types.iter().any(is_ref) || meta.outs.iter().any(|o| is_ref(&o.typ)) => {
1321 JitOp::Fallback
1322 }
1323 _ => classify_node(node),
1324 };
1325 if !matches!(named, JitOp::Fallback) {
1326 return named;
1327 }
1328 if let Some(kit) = node.compiled_slot(
1329 wire_types,
1330 crate::compile::select::Engine::Native(crate::compile::select::Provenance::Auto),
1331 ) {
1332 return JitOp::SlotCall {
1333 kit: SlotKitRef::new(kit),
1334 scratch_base: 0,
1335 };
1336 }
1337 if let Some(op) = node.compiled_u64() {
1338 return JitOp::SlotCall {
1339 kit: SlotKitRef::new(crate::ast::CompiledSlotKit {
1340 scratch: Vec::new(),
1341 op: Box::new(move |inputs, outputs, _| op(inputs, outputs)),
1342 }),
1343 scratch_base: 0,
1344 };
1345 }
1346 JitOp::Fallback
1347}
1348
1349#[derive(Debug, Clone, PartialEq)]
1356pub enum JitOp {
1357 Identity,
1360 AddConst(u64),
1362 MulConst(u64),
1364 DivConst(u64),
1366 ModConst(u64),
1368 ClampConst(u64, u64),
1370 Interleave,
1372 MixedRadixConst(Vec<u64>),
1374 Hash,
1376 SplitMix64,
1378 ShuffleConst(u64, u64, u64),
1380
1381 UnitInterval,
1384 F64ToU64,
1386 RoundToU64,
1389 FloorToU64,
1391 CeilToU64,
1393 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),
1415 MathBinary(u8),
1418
1419 U64Add2,
1422 U64Sub2,
1424 U64Mul2,
1426 U64Div2,
1428 U64Mod2,
1430 U64And,
1432 U64Or,
1434 U64Xor,
1436 U64Shl,
1438 U64Shr,
1440 U64Not,
1442
1443 ToF64,
1446
1447 F64Add,
1449 F64Sub,
1451 F64Mul,
1453 F64Div,
1455 F64Mod,
1457 U64DivWire,
1460 U64ModWire,
1463
1464 SlotCall {
1472 kit: SlotKitRef,
1475 scratch_base: usize,
1478 },
1479
1480 Convert {
1488 from: crate::ast::PortType,
1490 to: crate::ast::PortType,
1492 kit: SlotKitRef,
1494 scratch_base: usize,
1496 },
1497
1498 U64ToStr {
1504 scratch_base: usize,
1506 },
1507 I64ToStr {
1509 scratch_base: usize,
1511 },
1512 F64ToStr {
1514 scratch_base: usize,
1516 },
1517 StrConcat {
1520 scratch_base: usize,
1522 },
1523 JsonToStr {
1525 scratch_base: usize,
1527 },
1528
1529 VecProduce {
1533 kind: VecProducer,
1535 scratch_base: usize,
1537 },
1538 VecReduce(VecReducer),
1540 RegLane(RegLaneRead),
1543 RegProduce(RegProducer),
1545 RegDotF32,
1549 RegShuffleConst([u8; 16]),
1552
1553 IsPositiveCheck {
1560 name_ptr: u64,
1562 name_len: u64,
1564 },
1565 InRangeCheck(u64, u64),
1569 IsOneOfCheck {
1577 allowed: Vec<u64>,
1579 set_ptr: u64,
1581 set_len: u64,
1583 },
1584
1585 RegBinOp(u8, u8),
1593 RegCopy,
1597 RegSplat(u8),
1601
1602 U64Cmp(ir::condcodes::IntCC),
1605 F64Cmp(ir::condcodes::FloatCC),
1607 SelectU64,
1609 SelectF64,
1611
1612 I64ToF64,
1615 ToBool,
1617 ConstU64(u64),
1619 ConstF64(u64),
1621
1622 HashRangeConst(u64),
1625 HashIntervalConst(u64, u64),
1627 InvLerpConst(u64, u64),
1629 RemapConst(u64, u64, u64, u64),
1631
1632 EpochOffsetConst(u64),
1635 EpochScaleConst(u64),
1637 ThreadId,
1639 CurrentEpochMillis,
1641
1642 Perlin1dConst(u64, u64),
1645 Perlin2dConst(u64, u64),
1647 Simplex2dConst(u64, u64),
1649 FractalNoise1dConst(u64, u64, u64),
1651 FractalNoise2dConst(u64, u64, u64),
1653
1654 VariadicSum,
1657 VariadicProduct,
1659 VariadicMin,
1661 VariadicMax,
1663 CheckedAdd,
1665 CheckedSub,
1667 CheckedMul,
1669 CeilToMultiple,
1671 MultiplesAtLeast,
1673
1674 FairCoin,
1677 BlendConst(u64),
1679 LfsrStepConst(u64),
1682 PcgConst(u64, u64),
1684 PcgStreamConst(u64),
1686 CycleWalkConst(u64, u64, u64),
1688 UnfairCoinConst(u64),
1690 CoinFlipConst(u64),
1693 ChanceConst(u64),
1695 NOfConst(u64, u64),
1697
1698 Fallback,
1701}
1702
1703pub fn classify_node(node: &dyn PolydatNode) -> JitOp {
1710 if let Some(op) = conversion_op(node) {
1711 return op;
1712 }
1713 let name = node.meta().name.as_str();
1714 let consts = node.jit_constants();
1715
1716 match name {
1717 "identity" => JitOp::Identity,
1718 "hash" | "splitmix64" | "scatter" => JitOp::SplitMix64,
1719 "fair_coin" => JitOp::FairCoin,
1720 "unfair_coin" => {
1721 if let Some(&p) = consts.first() {
1722 JitOp::UnfairCoinConst(p)
1723 } else {
1724 JitOp::Fallback
1725 }
1726 }
1727 "chance" => {
1728 if let Some(&p) = consts.first() {
1729 JitOp::ChanceConst(p)
1730 } else {
1731 JitOp::Fallback
1732 }
1733 }
1734 "xxhash3" | "xxh3" => JitOp::Hash,
1735 "hash_range" => {
1736 if let Some(&c) = consts.first() {
1737 JitOp::HashRangeConst(c)
1738 } else {
1739 JitOp::Fallback
1740 }
1741 }
1742 "hash_interval" => {
1743 if consts.len() >= 2 {
1744 JitOp::HashIntervalConst(consts[0], consts[1])
1745 } else {
1746 JitOp::Fallback
1747 }
1748 }
1749 "add" => {
1750 if let Some(&c) = consts.first() {
1751 JitOp::AddConst(c)
1752 } else {
1753 JitOp::Fallback
1754 }
1755 }
1756 "mul" => {
1757 if let Some(&c) = consts.first() {
1758 JitOp::MulConst(c)
1759 } else {
1760 JitOp::Fallback
1761 }
1762 }
1763 "div" => {
1764 if let Some(&c) = consts.first() {
1765 JitOp::DivConst(c)
1766 } else {
1767 JitOp::Fallback
1768 }
1769 }
1770 "mod" => {
1771 if let Some(&c) = consts.first() {
1772 JitOp::ModConst(c)
1773 } else {
1774 JitOp::Fallback
1775 }
1776 }
1777 "clamp" => {
1778 if consts.len() >= 2 && consts[0] <= consts[1] {
1784 JitOp::ClampConst(consts[0], consts[1])
1785 } else {
1786 JitOp::Fallback
1787 }
1788 }
1789 "interleave" => JitOp::Interleave,
1790 "mixed_radix" => {
1791 if consts.is_empty() {
1792 JitOp::Fallback
1793 } else {
1794 JitOp::MixedRadixConst(consts)
1795 }
1796 }
1797 "shuffle" => {
1798 if consts.len() >= 3 {
1799 JitOp::ShuffleConst(consts[0], consts[1], consts[2])
1800 } else {
1801 JitOp::Fallback
1802 }
1803 }
1804 "unit_interval" => JitOp::UnitInterval,
1806 "f64_to_u64" => JitOp::F64ToU64,
1807 "round_to_u64" => JitOp::RoundToU64,
1808 "floor_to_u64" => JitOp::FloorToU64,
1809 "ceil_to_u64" => JitOp::CeilToU64,
1810 "clamp_f64" => {
1811 if consts.len() >= 2 {
1812 JitOp::ClampF64Const(consts[0], consts[1])
1813 } else {
1814 JitOp::Fallback
1815 }
1816 }
1817 "lerp" => {
1818 if consts.len() >= 2 {
1819 JitOp::LerpConst(consts[0], consts[1])
1820 } else {
1821 JitOp::Fallback
1822 }
1823 }
1824 "scale_range" => {
1825 if consts.len() >= 2 {
1826 JitOp::ScaleRangeConst(consts[0], consts[1])
1827 } else {
1828 JitOp::Fallback
1829 }
1830 }
1831 "quantize" => {
1832 if let Some(&c) = consts.first() {
1833 JitOp::QuantizeConst(c)
1834 } else {
1835 JitOp::Fallback
1836 }
1837 }
1838 "discretize" => {
1839 if consts.len() >= 2 {
1840 JitOp::DiscretizeConst(consts[0], consts[1])
1841 } else {
1842 JitOp::Fallback
1843 }
1844 }
1845 "lut_sample" | "dist_normal" | "icd_normal" | "dist_exponential" | "icd_exponential"
1846 | "dist_uniform" | "dist_pareto" | "dist_zipf" | "dist_empirical" => {
1847 if consts.len() >= 2 {
1848 JitOp::LutSampleConst(consts[0], consts[1])
1849 } else {
1850 JitOp::Fallback
1851 }
1852 }
1853 "sin" => JitOp::MathUnary(0),
1855 "cos" => JitOp::MathUnary(1),
1856 "tan" => JitOp::MathUnary(2),
1857 "asin" => JitOp::MathUnary(3),
1858 "acos" => JitOp::MathUnary(4),
1859 "atan" => JitOp::MathUnary(5),
1860 "sqrt" => JitOp::MathUnary(6),
1861 "abs_f64" => JitOp::MathUnary(7),
1862 "ln" => JitOp::MathUnary(8),
1863 "exp" => JitOp::MathUnary(9),
1864 "floor_base10" => JitOp::MathUnary(10),
1865 "ceiling_base10" => JitOp::MathUnary(11),
1866 "closest_base10" => JitOp::MathUnary(12),
1867 "floor_decade" => JitOp::MathUnary(13),
1868 "ceiling_decade" => JitOp::MathUnary(14),
1869 "closest_decade" => JitOp::MathUnary(15),
1870 "floor_binomial" => JitOp::MathUnary(16),
1871 "ceiling_binomial" => JitOp::MathUnary(17),
1872 "closest_binomial" => JitOp::MathUnary(18),
1873 "floor_fibonacci" => JitOp::MathUnary(19),
1874 "ceiling_fibonacci" => JitOp::MathUnary(20),
1875 "closest_fibonacci" => JitOp::MathUnary(21),
1876 "atan2" => JitOp::MathBinary(0),
1877 "pow" => JitOp::MathBinary(1),
1878 "round_nearest" => JitOp::MathBinary(2),
1879 "round_floor" => JitOp::MathBinary(3),
1880 "round_ceiling" => JitOp::MathBinary(4),
1881 "to_f64" => JitOp::ToF64,
1882 "u64_add" => JitOp::U64Add2,
1884 "u64_sub" => JitOp::U64Sub2,
1885 "u64_mul" => JitOp::U64Mul2,
1886 "u64_div" => JitOp::U64Div2,
1887 "u64_mod" => JitOp::U64Mod2,
1888 "u64_and" => JitOp::U64And,
1889 "u64_or" => JitOp::U64Or,
1890 "u64_xor" => JitOp::U64Xor,
1891 "u64_shl" => JitOp::U64Shl,
1892 "u64_shr" => JitOp::U64Shr,
1893 "u64_not" => JitOp::U64Not,
1894
1895 "reg_add_i8" => JitOp::RegBinOp(0, 0),
1897 "reg_sub_i8" => JitOp::RegBinOp(0, 1),
1898 "reg_shuffle_bytes" => {
1903 let mut mask = [0u8; 16];
1904 if consts.len() == 16 && consts.iter().all(|&m| m < 16) {
1905 for (m, &c) in mask.iter_mut().zip(consts.iter()) {
1906 *m = c as u8;
1907 }
1908 JitOp::RegShuffleConst(mask)
1909 } else {
1910 JitOp::Fallback
1911 }
1912 }
1913 "reg_add_i16" => JitOp::RegBinOp(1, 0),
1914 "reg_sub_i16" => JitOp::RegBinOp(1, 1),
1915 "reg_mul_i16" => JitOp::RegBinOp(1, 2),
1916 "reg_add_i32" => JitOp::RegBinOp(2, 0),
1917 "reg_sub_i32" => JitOp::RegBinOp(2, 1),
1918 "reg_mul_i32" => JitOp::RegBinOp(2, 2),
1919 "reg_add_i64" => JitOp::RegBinOp(3, 0),
1920 "reg_sub_i64" => JitOp::RegBinOp(3, 1),
1921 "reg_mul_i64" => JitOp::RegBinOp(3, 2),
1922 "reg_add_f32" => JitOp::RegBinOp(4, 0),
1923 "reg_sub_f32" => JitOp::RegBinOp(4, 1),
1924 "reg_mul_f32" => JitOp::RegBinOp(4, 2),
1925 "reg_add_f64" => JitOp::RegBinOp(5, 0),
1926 "reg_sub_f64" => JitOp::RegBinOp(5, 1),
1927 "reg_mul_f64" => JitOp::RegBinOp(5, 2),
1928 "__reg_view_raw" | "__reg_view_i8x16" | "__reg_view_i16x8" | "__reg_view_i32x4"
1929 | "__reg_view_i64x2" | "__reg_view_f16x8" | "__reg_view_f32x4" | "__reg_view_f64x2" => {
1930 JitOp::RegCopy
1931 }
1932 "reg_splat_i8" => JitOp::RegSplat(0),
1933 "reg_splat_i16" => JitOp::RegSplat(1),
1934 "reg_splat_i32" => JitOp::RegSplat(2),
1935 "reg_splat_i64" => JitOp::RegSplat(3),
1936 "reg_splat_f32" => JitOp::RegSplat(4),
1937 "reg_splat_f64" => JitOp::RegSplat(5),
1938
1939 "f64_add" => JitOp::F64Add,
1940 "f64_sub" => JitOp::F64Sub,
1941 "f64_mul" => JitOp::F64Mul,
1942 "f64_div" => JitOp::F64Div,
1943 "f64_mod" => JitOp::F64Mod,
1944
1945 "u64_eq" => JitOp::U64Cmp(ir::condcodes::IntCC::Equal),
1947 "u64_ne" => JitOp::U64Cmp(ir::condcodes::IntCC::NotEqual),
1948 "u64_lt" => JitOp::U64Cmp(ir::condcodes::IntCC::UnsignedLessThan),
1949 "u64_le" => JitOp::U64Cmp(ir::condcodes::IntCC::UnsignedLessThanOrEqual),
1950 "u64_gt" => JitOp::U64Cmp(ir::condcodes::IntCC::UnsignedGreaterThan),
1951 "u64_ge" => JitOp::U64Cmp(ir::condcodes::IntCC::UnsignedGreaterThanOrEqual),
1952 "f64_eq" => JitOp::F64Cmp(ir::condcodes::FloatCC::Equal),
1953 "f64_ne" => JitOp::F64Cmp(ir::condcodes::FloatCC::NotEqual),
1954 "f64_lt" => JitOp::F64Cmp(ir::condcodes::FloatCC::LessThan),
1955 "f64_le" => JitOp::F64Cmp(ir::condcodes::FloatCC::LessThanOrEqual),
1956 "f64_gt" => JitOp::F64Cmp(ir::condcodes::FloatCC::GreaterThan),
1957 "f64_ge" => JitOp::F64Cmp(ir::condcodes::FloatCC::GreaterThanOrEqual),
1958 "select_u64" | "select" => JitOp::SelectU64,
1959 "select_f64" => JitOp::SelectF64,
1960
1961 "div_wire" => JitOp::U64DivWire,
1963 "mod_wire" => JitOp::U64ModWire,
1964 "ceil_to_multiple" => JitOp::CeilToMultiple,
1965 "multiples_at_least" => JitOp::MultiplesAtLeast,
1966 "checked_add" => JitOp::CheckedAdd,
1967 "checked_sub" => JitOp::CheckedSub,
1968 "checked_mul" => JitOp::CheckedMul,
1969
1970 "sum" => JitOp::VariadicSum,
1972 "product" => JitOp::VariadicProduct,
1973 "min" => JitOp::VariadicMin,
1974 "max" => JitOp::VariadicMax,
1975
1976 "blend" => {
1978 if let Some(&c) = consts.first() {
1982 JitOp::BlendConst(c)
1983 } else {
1984 JitOp::Fallback
1985 }
1986 }
1987 "lfsr_step" => {
1988 if let Some(&fb) = consts.first() {
1989 JitOp::LfsrStepConst(fb)
1990 } else {
1991 JitOp::Fallback
1992 }
1993 }
1994 "pcg" => {
1995 if consts.len() >= 2 {
1996 JitOp::PcgConst(consts[0], consts[1])
1997 } else {
1998 JitOp::Fallback
1999 }
2000 }
2001 "pcg_stream" => {
2002 if let Some(&seed) = consts.first() {
2003 JitOp::PcgStreamConst(seed)
2004 } else {
2005 JitOp::Fallback
2006 }
2007 }
2008 "n_of" => {
2009 if consts.len() >= 2 {
2010 JitOp::NOfConst(consts[0], consts[1])
2011 } else {
2012 JitOp::Fallback
2013 }
2014 }
2015
2016 "cycle_walk" => {
2017 if consts.len() >= 3 {
2018 JitOp::CycleWalkConst(consts[0], consts[1], consts[2])
2019 } else {
2020 JitOp::Fallback
2021 }
2022 }
2023 "coin_flip" => {
2024 if let Some(&threshold) = consts.first() {
2029 JitOp::CoinFlipConst(threshold)
2030 } else {
2031 JitOp::Fallback
2032 }
2033 }
2034 "default_or" => JitOp::Identity,
2037 "const_u64" | "const_bool" => {
2038 if let Some(&c) = consts.first() {
2039 JitOp::ConstU64(c)
2040 } else {
2041 JitOp::Fallback
2042 }
2043 }
2044 "const_f64" => {
2045 if let Some(&c) = consts.first() {
2046 JitOp::ConstF64(c)
2047 } else {
2048 JitOp::Fallback
2049 }
2050 }
2051 "inv_lerp" => {
2052 if consts.len() >= 2 {
2053 JitOp::InvLerpConst(consts[0], consts[1])
2054 } else {
2055 JitOp::Fallback
2056 }
2057 }
2058 "remap" => {
2059 if consts.len() >= 4 {
2060 JitOp::RemapConst(consts[0], consts[1], consts[2], consts[3])
2061 } else {
2062 JitOp::Fallback
2063 }
2064 }
2065 "epoch_offset" => {
2066 if let Some(&c) = consts.first() {
2067 JitOp::EpochOffsetConst(c)
2068 } else {
2069 JitOp::Fallback
2070 }
2071 }
2072 "epoch_scale" => {
2073 if let Some(&c) = consts.first() {
2074 JitOp::EpochScaleConst(c)
2075 } else {
2076 JitOp::Fallback
2077 }
2078 }
2079 "thread_id" => JitOp::ThreadId,
2080 "current_epoch_millis" => JitOp::CurrentEpochMillis,
2081 "perlin_1d" => {
2082 if consts.len() >= 2 {
2083 JitOp::Perlin1dConst(consts[0], consts[1])
2084 } else {
2085 JitOp::Fallback
2086 }
2087 }
2088 "perlin_2d" => {
2089 if consts.len() >= 2 {
2090 JitOp::Perlin2dConst(consts[0], consts[1])
2091 } else {
2092 JitOp::Fallback
2093 }
2094 }
2095 "simplex_2d" => {
2096 if consts.len() >= 2 {
2097 JitOp::Simplex2dConst(consts[0], consts[1])
2098 } else {
2099 JitOp::Fallback
2100 }
2101 }
2102 "fractal_noise_1d" => {
2103 if consts.len() >= 3 {
2104 JitOp::FractalNoise1dConst(consts[0], consts[1], consts[2])
2105 } else {
2106 JitOp::Fallback
2107 }
2108 }
2109 "fractal_noise_2d" => {
2110 if consts.len() >= 3 {
2111 JitOp::FractalNoise2dConst(consts[0], consts[1], consts[2])
2112 } else {
2113 JitOp::Fallback
2114 }
2115 }
2116
2117 "trunc_u64" => JitOp::F64ToU64,
2121 "round_u64" => JitOp::RoundToU64,
2124
2125 "weighted_pick" => {
2126 if consts.len() >= 5 {
2127 JitOp::WeightedPickConst(consts[0], consts[1], consts[2], consts[3], consts[4])
2128 } else {
2129 JitOp::Fallback
2130 }
2131 }
2132
2133 "is_positive" => {
2139 let name = node.meta().ins.iter().find_map(|slot| match slot {
2140 crate::ast::Slot::Const {
2141 name,
2142 value: crate::ast::ConstValue::Str(v),
2143 } if name == "name" => Some(v),
2144 _ => None,
2145 });
2146 match name {
2147 Some(v) => JitOp::IsPositiveCheck {
2148 name_ptr: v.as_ptr() as u64,
2149 name_len: v.len() as u64,
2150 },
2151 None => JitOp::IsPositiveCheck {
2152 name_ptr: 0,
2153 name_len: 0,
2154 },
2155 }
2156 }
2157 "in_range" => {
2158 if consts.len() >= 2 {
2159 JitOp::InRangeCheck(consts[0], consts[1])
2160 } else {
2161 JitOp::Fallback
2162 }
2163 }
2164 "is_one_of" => {
2165 if consts.is_empty() {
2166 JitOp::Fallback
2167 } else {
2168 let set = node.meta().ins.iter().find_map(|slot| match slot {
2169 crate::ast::Slot::Const {
2170 name,
2171 value: crate::ast::ConstValue::VecU64(v),
2172 } if name == "allowed" => Some(v),
2173 _ => None,
2174 });
2175 let (set_ptr, set_len) = match set {
2176 Some(v) => (v.as_ptr() as u64, v.len() as u64),
2177 None => (0, 0),
2178 };
2179 JitOp::IsOneOfCheck {
2180 allowed: consts,
2181 set_ptr,
2182 set_len,
2183 }
2184 }
2185 }
2186 _ => JitOp::Fallback,
2191 }
2192}
2193
2194#[doc(hidden)]
2202pub fn compile_jit_raw(
2203 coord_count: usize,
2204 total_slots: usize,
2205 steps: Vec<(JitOp, Vec<usize>, Vec<usize>)>,
2206 output_map: HashMap<String, usize>,
2207 nodes: Vec<Box<dyn PolydatNode>>,
2208) -> Result<JitKernelRaw, String> {
2209 let alone = vec![false; steps.len()];
2210 compile_jit_raw_with(
2211 coord_count,
2212 total_slots,
2213 steps,
2214 output_map,
2215 nodes,
2216 crate::compile::externs::Externs::coordinates_only(coord_count),
2217 super::kernels::ScratchPlan::default(),
2218 Vec::new(),
2219 alone,
2220 )
2221}
2222
2223fn pure_units(
2232 steps: &[(JitOp, Vec<usize>, Vec<usize>)],
2233 total_slots: usize,
2234 alone: &[bool],
2235 volatile: &[usize],
2236 unset_read: &[usize],
2237) -> crate::compile::fusion_units::UnitPlan {
2238 let mut producer = vec![usize::MAX; total_slots + 1];
2239 for (i, (_, _, outs)) in steps.iter().enumerate() {
2240 for &s in outs {
2241 if s < producer.len() {
2242 producer[s] = i;
2243 }
2244 }
2245 }
2246 let preds: Vec<Vec<usize>> = steps
2247 .iter()
2248 .map(|(_, ins, _)| {
2249 let mut p: Vec<usize> = ins
2250 .iter()
2251 .filter_map(|&s| producer.get(s).copied().filter(|&p| p != usize::MAX))
2252 .collect();
2253 p.sort_unstable();
2254 p.dedup();
2255 p
2256 })
2257 .collect();
2258 let inputs_read: Vec<Vec<usize>> = steps
2260 .iter()
2261 .map(|(_, ins, _)| {
2262 ins.iter()
2263 .copied()
2264 .filter(|&s| producer.get(s).is_some_and(|&p| p == usize::MAX))
2265 .collect()
2266 })
2267 .collect();
2268 let fusible: Vec<bool> = (0..steps.len())
2269 .map(|i| !alone.get(i).copied().unwrap_or(false))
2270 .collect();
2271 let mut class = vec![0u64; steps.len()];
2272 for &v in volatile {
2273 if v < class.len() {
2274 class[v] = 1;
2275 }
2276 }
2277 let reads: Vec<Vec<usize>> = inputs_read
2278 .iter()
2279 .map(|ins| {
2280 ins.iter()
2281 .copied()
2282 .filter(|s| unset_read.binary_search(s).is_ok())
2283 .collect()
2284 })
2285 .collect();
2286 let class = crate::compile::fusion_units::refine_by_externs(&preds, &reads, &class);
2287 let rank: Vec<usize> = (0..steps.len()).collect();
2288 crate::compile::fusion_units::plan_units(&preds, &inputs_read, &fusible, &class, &rank, &|_| {
2289 false
2290 })
2291}
2292
2293fn unit_dependents(
2296 input_dependents: Vec<Vec<usize>>,
2297 plan: &crate::compile::fusion_units::UnitPlan,
2298) -> Vec<Vec<usize>> {
2299 input_dependents
2300 .into_iter()
2301 .map(|steps| {
2302 let mut units: Vec<usize> = steps.iter().map(|&s| plan.unit_of[s]).collect();
2303 units.sort_unstable();
2304 units.dedup();
2305 units
2306 })
2307 .collect()
2308}
2309
2310#[allow(clippy::too_many_arguments)]
2313pub(crate) fn compile_jit_raw_with(
2314 coord_count: usize,
2315 total_slots: usize,
2316 steps: Vec<(JitOp, Vec<usize>, Vec<usize>)>,
2317 output_map: HashMap<String, usize>,
2318 nodes: Vec<Box<dyn PolydatNode>>,
2319 externs: crate::compile::externs::Externs,
2320 scratch: super::kernels::ScratchPlan,
2321 volatile: Vec<usize>,
2322 alone: Vec<bool>,
2323) -> Result<JitKernelRaw, String> {
2324 let plan = pure_units(
2327 &steps,
2328 total_slots,
2329 &alone,
2330 &volatile,
2331 &externs.unset_read_slots(),
2332 );
2333 let (_, entry, code) = compile_jit_impl(&steps, Some(&plan.units), Some(total_slots))?;
2334 let cones =
2335 super::kernels::ConePlan::new(&steps, total_slots, &plan, output_map.values().copied());
2336 let mut core = JitCore::new(
2337 total_slots,
2338 coord_count,
2339 output_map,
2340 code,
2341 nodes,
2342 scratch,
2343 volatile,
2344 entry,
2345 cones,
2346 );
2347 core.set_externs(externs);
2348 core.engine =
2349 crate::compile::select::Engine::PureNative(crate::compile::select::Provenance::Raw);
2350 Ok(JitKernelRaw { core })
2351}
2352
2353pub(crate) type JitSegmentCode = (NativeFn, super::kernels::JitCode);
2357
2358pub(crate) fn compile_jit_entry(
2362 steps: &[(JitOp, Vec<usize>, Vec<usize>)],
2363 tracker: Option<usize>,
2364) -> Result<JitSegmentCode, String> {
2365 let (raw_fn, _, code) = compile_jit_impl(steps, None, tracker)?;
2366 Ok((raw_fn, code))
2367}
2368
2369#[allow(clippy::too_many_arguments)]
2371pub(crate) fn compile_jit_push_pull(
2372 coord_count: usize,
2373 total_slots: usize,
2374 steps: Vec<(JitOp, Vec<usize>, Vec<usize>)>,
2375 output_map: HashMap<String, usize>,
2376 nodes: Vec<Box<dyn PolydatNode>>,
2377 input_dependents: Vec<Vec<usize>>,
2378 externs: crate::compile::externs::Externs,
2379 scratch: super::kernels::ScratchPlan,
2380 volatile: Vec<usize>,
2381 alone: Vec<bool>,
2382) -> Result<JitKernelPushPull, String> {
2383 let buffer_len = total_slots;
2384 let plan = pure_units(
2385 &steps,
2386 total_slots,
2387 &alone,
2388 &volatile,
2389 &externs.unset_read_slots(),
2390 );
2391 let (_, entry, code) = compile_jit_impl(&steps, Some(&plan.units), Some(total_slots))?;
2392 let step_outs: Vec<&[usize]> = steps.iter().map(|(_, _, o)| o.as_slice()).collect();
2393 let slot_provenance =
2394 crate::compile::slot_provenance(coord_count, buffer_len, &step_outs, &input_dependents);
2395 let cones =
2396 super::kernels::ConePlan::new(&steps, total_slots, &plan, output_map.values().copied());
2397 let input_dependents = unit_dependents(input_dependents, &plan);
2398 let mut core = JitCore::new(
2399 total_slots,
2400 coord_count,
2401 output_map,
2402 code,
2403 nodes,
2404 scratch,
2405 volatile,
2406 entry,
2407 cones,
2408 );
2409 core.set_externs(externs);
2410 Ok(JitKernelPushPull {
2411 core,
2412 input_dependents,
2413 slot_provenance,
2414 changed_mask: crate::kernel::ProvMask::all_below(coord_count),
2415 force_run: false,
2416 })
2417}
2418
2419pub type NativeFn = unsafe fn(*const u64, *mut u64, *mut crate::ast::ScratchBuf);
2423pub type NativeDispatchFn =
2427 unsafe fn(*const u64, *mut u64, *mut crate::ast::ScratchBuf, *const u32, u64, *mut u8);
2428
2429type JitCompiled = (NativeFn, NativeDispatchFn, super::kernels::JitCode);
2435
2436type JitEntry = (NativeFn, NativeDispatchFn, bool);
2439
2440type JitFunctionSpec<'a> = (
2443 &'a [(JitOp, Vec<usize>, Vec<usize>)],
2444 Option<&'a [Vec<usize>]>,
2445);
2446
2447fn compile_jit_impl(
2451 steps: &[(JitOp, Vec<usize>, Vec<usize>)],
2452 dispatch: Option<&[Vec<usize>]>,
2453 tracker: Option<usize>,
2454) -> Result<JitCompiled, String> {
2455 let (entries, code) = compile_jit_module(&[(steps, dispatch)], tracker)?;
2456 let (straight_fn, dispatch_fn, _) = entries[0];
2457 Ok((straight_fn, dispatch_fn, code))
2458}
2459
2460pub(crate) type JitStep = (JitOp, Vec<usize>, Vec<usize>);
2462
2463pub(crate) fn compile_jit_entries(
2467 batches: &[&[JitStep]],
2468 tracker: Option<usize>,
2469) -> Result<(Vec<(NativeFn, bool)>, super::kernels::JitCode), String> {
2470 let specs: Vec<JitFunctionSpec> = batches.iter().map(|&b| (b, None)).collect();
2471 let (entries, code) = compile_jit_module(&specs, tracker)?;
2472 Ok((
2473 entries
2474 .into_iter()
2475 .map(|(f, _, fallible)| (f, fallible))
2476 .collect(),
2477 code,
2478 ))
2479}
2480
2481fn compile_jit_module(
2483 functions: &[JitFunctionSpec],
2484 tracker: Option<usize>,
2485) -> Result<(Vec<JitEntry>, super::kernels::JitCode), String> {
2486 let mut flag_builder = settings::builder();
2487 flag_builder.set("opt_level", "speed").unwrap();
2488 flag_builder.set("unwind_info", "true").unwrap();
2493 flag_builder.set("preserve_frame_pointers", "true").unwrap();
2494 let isa = super::host_isa::build_host_isa(flag_builder)?;
2495
2496 let mut jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
2497
2498 jit_builder.symbol("jit_xxh3_hash", jit_xxh3_hash as *const u8);
2500 jit_builder.symbol("jit_interleave", jit_interleave as *const u8);
2501 jit_builder.symbol("jit_shuffle", jit_shuffle as *const u8);
2502 jit_builder.symbol("jit_lut_sample", jit_lut_sample as *const u8);
2503 jit_builder.symbol("jit_weighted_pick", jit_weighted_pick as *const u8);
2504 jit_builder.symbol("jit_pcg", jit_pcg as *const u8);
2505 jit_builder.symbol("jit_pcg_stream", jit_pcg_stream as *const u8);
2506 jit_builder.symbol("jit_n_of", jit_n_of as *const u8);
2507 jit_builder.symbol("jit_cycle_walk", jit_cycle_walk as *const u8);
2508 jit_builder.symbol("jit_perlin_1d", jit_perlin_1d as *const u8);
2509 jit_builder.symbol("jit_perlin_2d", jit_perlin_2d as *const u8);
2510 jit_builder.symbol("jit_simplex_2d", jit_simplex_2d as *const u8);
2511 jit_builder.symbol("jit_fractal_noise_1d", jit_fractal_noise_1d as *const u8);
2512 jit_builder.symbol("jit_fractal_noise_2d", jit_fractal_noise_2d as *const u8);
2513 jit_builder.symbol("jit_thread_id", jit_thread_id as *const u8);
2514 jit_builder.symbol(
2515 "jit_current_epoch_millis",
2516 jit_current_epoch_millis as *const u8,
2517 );
2518 jit_builder.symbol("jit_is_positive_fail", jit_is_positive_fail as *const u8);
2522 jit_builder.symbol("jit_in_range_fail", jit_in_range_fail as *const u8);
2523 jit_builder.symbol("jit_is_one_of_fail", jit_is_one_of_fail as *const u8);
2524 jit_builder.symbol("jit_slot_call", jit_slot_call as *const u8);
2527 jit_builder.symbol("jit_u64_to_str", jit_u64_to_str as *const u8);
2528 jit_builder.symbol("jit_i64_to_str", jit_i64_to_str as *const u8);
2529 jit_builder.symbol("jit_f64_to_str", jit_f64_to_str as *const u8);
2530 jit_builder.symbol("jit_str_concat", jit_str_concat as *const u8);
2531 jit_builder.symbol("jit_json_to_str", jit_json_to_str as *const u8);
2532 jit_builder.symbol("jit_vec_add", jit_vec_add as *const u8);
2533 jit_builder.symbol("jit_vec_scale", jit_vec_scale as *const u8);
2534 jit_builder.symbol("jit_vec_norm", jit_vec_norm as *const u8);
2535 jit_builder.symbol("jit_hash_vec", jit_hash_vec as *const u8);
2536 jit_builder.symbol("jit_xxhash3_vec", jit_xxhash3_vec as *const u8);
2537 jit_builder.symbol("jit_reg_to_vec_f32", jit_reg_to_vec_f32 as *const u8);
2538 jit_builder.symbol("jit_vec_dot", jit_vec_dot as *const u8);
2539 jit_builder.symbol("jit_vec_l2", jit_vec_l2 as *const u8);
2540 jit_builder.symbol("jit_vec_cosine", jit_vec_cosine as *const u8);
2541 jit_builder.symbol("jit_lid_mle", jit_lid_mle as *const u8);
2542 jit_builder.symbol("jit_reg_lane_f32", jit_reg_lane_f32 as *const u8);
2543 jit_builder.symbol("jit_reg_lane_i16", jit_reg_lane_i16 as *const u8);
2544 jit_builder.symbol("jit_reg_lane_i64", jit_reg_lane_i64 as *const u8);
2545 jit_builder.symbol("jit_reg_with_lane_f32", jit_reg_with_lane_f32 as *const u8);
2546 jit_builder.symbol("jit_reg_gather_f32", jit_reg_gather_f32 as *const u8);
2547 jit_builder.symbol("jit_vec_to_reg_f32", jit_vec_to_reg_f32 as *const u8);
2548 jit_builder.symbol("jit_reg_mul_i8", jit_reg_mul_i8 as *const u8);
2549 jit_builder.symbol("jit_sin", jit_sin as *const u8);
2551 jit_builder.symbol("jit_cos", jit_cos as *const u8);
2552 jit_builder.symbol("jit_tan", jit_tan as *const u8);
2553 jit_builder.symbol("jit_asin", jit_asin as *const u8);
2554 jit_builder.symbol("jit_acos", jit_acos as *const u8);
2555 jit_builder.symbol("jit_atan", jit_atan as *const u8);
2556 jit_builder.symbol("jit_sqrt", jit_sqrt as *const u8);
2557 jit_builder.symbol("jit_abs_f64", jit_abs_f64 as *const u8);
2558 jit_builder.symbol("jit_ln", jit_ln as *const u8);
2559 jit_builder.symbol("jit_exp", jit_exp as *const u8);
2560 jit_builder.symbol("jit_floor_base10", jit_floor_base10 as *const u8);
2561 jit_builder.symbol("jit_ceiling_base10", jit_ceiling_base10 as *const u8);
2562 jit_builder.symbol("jit_closest_base10", jit_closest_base10 as *const u8);
2563 jit_builder.symbol("jit_floor_decade", jit_floor_decade as *const u8);
2564 jit_builder.symbol("jit_ceiling_decade", jit_ceiling_decade as *const u8);
2565 jit_builder.symbol("jit_closest_decade", jit_closest_decade as *const u8);
2566 jit_builder.symbol("jit_floor_binomial", jit_floor_binomial as *const u8);
2567 jit_builder.symbol("jit_ceiling_binomial", jit_ceiling_binomial as *const u8);
2568 jit_builder.symbol("jit_closest_binomial", jit_closest_binomial as *const u8);
2569 jit_builder.symbol("jit_floor_fibonacci", jit_floor_fibonacci as *const u8);
2570 jit_builder.symbol("jit_ceiling_fibonacci", jit_ceiling_fibonacci as *const u8);
2571 jit_builder.symbol("jit_closest_fibonacci", jit_closest_fibonacci as *const u8);
2572 jit_builder.symbol("jit_atan2", jit_atan2 as *const u8);
2573 jit_builder.symbol("jit_pow", jit_pow as *const u8);
2574 jit_builder.symbol("jit_round_nearest", jit_round_nearest as *const u8);
2575 jit_builder.symbol("jit_round_floor", jit_round_floor as *const u8);
2576 jit_builder.symbol("jit_round_ceiling", jit_round_ceiling as *const u8);
2577 jit_builder.symbol("jit_f64_mod", jit_f64_mod as *const u8);
2578 jit_builder.symbol("jit_div_zero_fail", jit_div_zero_fail as *const u8);
2579
2580 let mut module = JITModule::new(jit_builder);
2581
2582 let hash_func_id = {
2584 let mut sig = module.make_signature();
2585 sig.params.push(AbiParam::new(types::I64));
2586 sig.returns.push(AbiParam::new(types::I64));
2587 module
2588 .declare_function("jit_xxh3_hash", Linkage::Import, &sig)
2589 .map_err(|e| format!("declare hash: {e}"))?
2590 };
2591
2592 let interleave_func_id = {
2594 let mut sig = module.make_signature();
2595 sig.params.push(AbiParam::new(types::I64));
2596 sig.params.push(AbiParam::new(types::I64));
2597 sig.returns.push(AbiParam::new(types::I64));
2598 module
2599 .declare_function("jit_interleave", Linkage::Import, &sig)
2600 .map_err(|e| format!("declare interleave: {e}"))?
2601 };
2602
2603 let shuffle_func_id = {
2605 let mut sig = module.make_signature();
2606 for _ in 0..4 {
2607 sig.params.push(AbiParam::new(types::I64));
2608 }
2609 sig.returns.push(AbiParam::new(types::I64));
2610 module
2611 .declare_function("jit_shuffle", Linkage::Import, &sig)
2612 .map_err(|e| format!("declare shuffle: {e}"))?
2613 };
2614
2615 let lut_sample_func_id = {
2617 let mut sig = module.make_signature();
2618 for _ in 0..3 {
2619 sig.params.push(AbiParam::new(types::I64));
2620 }
2621 sig.returns.push(AbiParam::new(types::I64));
2622 module
2623 .declare_function("jit_lut_sample", Linkage::Import, &sig)
2624 .map_err(|e| format!("declare lut_sample: {e}"))?
2625 };
2626
2627 let weighted_pick_func_id = {
2629 let mut sig = module.make_signature();
2630 for _ in 0..6 {
2631 sig.params.push(AbiParam::new(types::I64));
2632 }
2633 sig.returns.push(AbiParam::new(types::I64));
2634 module
2635 .declare_function("jit_weighted_pick", Linkage::Import, &sig)
2636 .map_err(|e| format!("declare weighted_pick: {e}"))?
2637 };
2638
2639 let pcg_func_id = {
2640 let mut sig = module.make_signature();
2641 for _ in 0..3 {
2642 sig.params.push(AbiParam::new(types::I64));
2643 }
2644 sig.returns.push(AbiParam::new(types::I64));
2645 module
2646 .declare_function("jit_pcg", Linkage::Import, &sig)
2647 .map_err(|e| format!("declare pcg: {e}"))?
2648 };
2649 let pcg_stream_func_id = {
2650 let mut sig = module.make_signature();
2651 for _ in 0..3 {
2652 sig.params.push(AbiParam::new(types::I64));
2653 }
2654 sig.returns.push(AbiParam::new(types::I64));
2655 module
2656 .declare_function("jit_pcg_stream", Linkage::Import, &sig)
2657 .map_err(|e| format!("declare pcg_stream: {e}"))?
2658 };
2659 let n_of_func_id = {
2660 let mut sig = module.make_signature();
2661 for _ in 0..3 {
2662 sig.params.push(AbiParam::new(types::I64));
2663 }
2664 sig.returns.push(AbiParam::new(types::I64));
2665 module
2666 .declare_function("jit_n_of", Linkage::Import, &sig)
2667 .map_err(|e| format!("declare n_of: {e}"))?
2668 };
2669 let cycle_walk_func_id = {
2670 let mut sig = module.make_signature();
2671 for _ in 0..4 {
2672 sig.params.push(AbiParam::new(types::I64));
2673 }
2674 sig.returns.push(AbiParam::new(types::I64));
2675 module
2676 .declare_function("jit_cycle_walk", Linkage::Import, &sig)
2677 .map_err(|e| format!("declare cycle_walk: {e}"))?
2678 };
2679 let perlin_1d_func_id = {
2680 let mut sig = module.make_signature();
2681 for _ in 0..3 {
2682 sig.params.push(AbiParam::new(types::I64));
2683 }
2684 sig.returns.push(AbiParam::new(types::I64));
2685 module
2686 .declare_function("jit_perlin_1d", Linkage::Import, &sig)
2687 .map_err(|e| format!("declare perlin_1d: {e}"))?
2688 };
2689 let perlin_2d_func_id = {
2690 let mut sig = module.make_signature();
2691 for _ in 0..4 {
2692 sig.params.push(AbiParam::new(types::I64));
2693 }
2694 sig.returns.push(AbiParam::new(types::I64));
2695 module
2696 .declare_function("jit_perlin_2d", Linkage::Import, &sig)
2697 .map_err(|e| format!("declare perlin_2d: {e}"))?
2698 };
2699 let simplex_2d_func_id = {
2700 let mut sig = module.make_signature();
2701 for _ in 0..4 {
2702 sig.params.push(AbiParam::new(types::I64));
2703 }
2704 sig.returns.push(AbiParam::new(types::I64));
2705 module
2706 .declare_function("jit_simplex_2d", Linkage::Import, &sig)
2707 .map_err(|e| format!("declare simplex_2d: {e}"))?
2708 };
2709 let fractal_noise_1d_func_id = {
2710 let mut sig = module.make_signature();
2711 for _ in 0..4 {
2712 sig.params.push(AbiParam::new(types::I64));
2713 }
2714 sig.returns.push(AbiParam::new(types::I64));
2715 module
2716 .declare_function("jit_fractal_noise_1d", Linkage::Import, &sig)
2717 .map_err(|e| format!("declare fractal_noise_1d: {e}"))?
2718 };
2719 let fractal_noise_2d_func_id = {
2720 let mut sig = module.make_signature();
2721 for _ in 0..5 {
2722 sig.params.push(AbiParam::new(types::I64));
2723 }
2724 sig.returns.push(AbiParam::new(types::I64));
2725 module
2726 .declare_function("jit_fractal_noise_2d", Linkage::Import, &sig)
2727 .map_err(|e| format!("declare fractal_noise_2d: {e}"))?
2728 };
2729 let thread_id_func_id = {
2730 let mut sig = module.make_signature();
2731 sig.returns.push(AbiParam::new(types::I64));
2732 module
2733 .declare_function("jit_thread_id", Linkage::Import, &sig)
2734 .map_err(|e| format!("declare thread_id: {e}"))?
2735 };
2736 let current_epoch_millis_func_id = {
2737 let mut sig = module.make_signature();
2738 sig.returns.push(AbiParam::new(types::I64));
2739 module
2740 .declare_function("jit_current_epoch_millis", Linkage::Import, &sig)
2741 .map_err(|e| format!("declare current_epoch_millis: {e}"))?
2742 };
2743
2744 let math_unary_names = [
2746 "jit_sin",
2747 "jit_cos",
2748 "jit_tan",
2749 "jit_asin",
2750 "jit_acos",
2751 "jit_atan",
2752 "jit_sqrt",
2753 "jit_abs_f64",
2754 "jit_ln",
2755 "jit_exp",
2756 "jit_floor_base10",
2757 "jit_ceiling_base10",
2758 "jit_closest_base10",
2759 "jit_floor_decade",
2760 "jit_ceiling_decade",
2761 "jit_closest_decade",
2762 "jit_floor_binomial",
2763 "jit_ceiling_binomial",
2764 "jit_closest_binomial",
2765 "jit_floor_fibonacci",
2766 "jit_ceiling_fibonacci",
2767 "jit_closest_fibonacci",
2768 ];
2769 let mut math_unary_ids = Vec::new();
2770 for name in &math_unary_names {
2771 let mut sig = module.make_signature();
2772 sig.params.push(AbiParam::new(types::I64));
2773 sig.returns.push(AbiParam::new(types::I64));
2774 math_unary_ids.push(
2775 module
2776 .declare_function(name, Linkage::Import, &sig)
2777 .map_err(|e| format!("declare {name}: {e}"))?,
2778 );
2779 }
2780
2781 let is_positive_fail_id = {
2785 let mut sig = module.make_signature();
2786 sig.params.push(AbiParam::new(types::I64));
2787 sig.params.push(AbiParam::new(types::I64));
2788 sig.params.push(AbiParam::new(types::I64));
2789 sig.returns.push(AbiParam::new(types::I64));
2790 module
2791 .declare_function("jit_is_positive_fail", Linkage::Import, &sig)
2792 .map_err(|e| format!("declare is_positive_fail: {e}"))?
2793 };
2794
2795 let in_range_fail_id = {
2797 let mut sig = module.make_signature();
2798 for _ in 0..3 {
2799 sig.params.push(AbiParam::new(types::I64));
2800 }
2801 sig.returns.push(AbiParam::new(types::I64));
2802 module
2803 .declare_function("jit_in_range_fail", Linkage::Import, &sig)
2804 .map_err(|e| format!("declare in_range_fail: {e}"))?
2805 };
2806
2807 let is_one_of_fail_id = {
2810 let mut sig = module.make_signature();
2811 sig.params.push(AbiParam::new(types::I64));
2812 sig.params.push(AbiParam::new(types::I64));
2813 sig.params.push(AbiParam::new(types::I64));
2814 sig.returns.push(AbiParam::new(types::I64));
2815 module
2816 .declare_function("jit_is_one_of_fail", Linkage::Import, &sig)
2817 .map_err(|e| format!("declare is_one_of_fail: {e}"))?
2818 };
2819
2820 let math_binary_names = [
2822 "jit_atan2",
2823 "jit_pow",
2824 "jit_round_nearest",
2825 "jit_round_floor",
2826 "jit_round_ceiling",
2827 "jit_f64_mod",
2828 ];
2829 const F64_MOD_HELPER: usize = 5;
2830
2831 let div_zero_fail_id = {
2833 let mut sig = module.make_signature();
2834 sig.params.push(AbiParam::new(types::I64));
2835 sig.returns.push(AbiParam::new(types::I64));
2836 module
2837 .declare_function("jit_div_zero_fail", Linkage::Import, &sig)
2838 .map_err(|e| format!("declare div_zero_fail: {e}"))?
2839 };
2840 let mut math_binary_ids = Vec::new();
2841 for name in &math_binary_names {
2842 let mut sig = module.make_signature();
2843 sig.params.push(AbiParam::new(types::I64));
2844 sig.params.push(AbiParam::new(types::I64));
2845 sig.returns.push(AbiParam::new(types::I64));
2846 math_binary_ids.push(
2847 module
2848 .declare_function(name, Linkage::Import, &sig)
2849 .map_err(|e| format!("declare {name}: {e}"))?,
2850 );
2851 }
2852
2853 let slot_call_id = {
2856 let mut sig = module.make_signature();
2857 for _ in 0..8 {
2858 sig.params.push(AbiParam::new(types::I64));
2859 }
2860 module
2861 .declare_function("jit_slot_call", Linkage::Import, &sig)
2862 .map_err(|e| format!("declare jit_slot_call: {e}"))?
2863 };
2864
2865 let mut declare_str = |name: &str, args: usize| -> Result<cranelift_module::FuncId, String> {
2869 let mut sig = module.make_signature();
2870 for _ in 0..args {
2871 sig.params.push(AbiParam::new(types::I64));
2872 }
2873 module
2874 .declare_function(name, Linkage::Import, &sig)
2875 .map_err(|e| format!("declare {name}: {e}"))
2876 };
2877 let u64_to_str_id = declare_str("jit_u64_to_str", 5)?;
2878 let i64_to_str_id = declare_str("jit_i64_to_str", 5)?;
2879 let f64_to_str_id = declare_str("jit_f64_to_str", 5)?;
2880 let str_concat_id = declare_str("jit_str_concat", 6)?;
2881 let json_to_str_id = declare_str("jit_json_to_str", 6)?;
2882
2883 let mut declare_words =
2888 |name: &str, args: usize, returns: bool| -> Result<cranelift_module::FuncId, String> {
2889 let mut sig = module.make_signature();
2890 for _ in 0..args {
2891 sig.params.push(AbiParam::new(types::I64));
2892 }
2893 if returns {
2894 sig.returns.push(AbiParam::new(types::I64));
2895 }
2896 module
2897 .declare_function(name, Linkage::Import, &sig)
2898 .map_err(|e| format!("declare {name}: {e}"))
2899 };
2900 let vec_producer_ids = [
2901 (VecProducer::Add, declare_words("jit_vec_add", 8, false)?),
2902 (
2903 VecProducer::Scale,
2904 declare_words("jit_vec_scale", 8, false)?,
2905 ),
2906 (VecProducer::Norm, declare_words("jit_vec_norm", 8, false)?),
2907 (
2908 VecProducer::HashVec,
2909 declare_words("jit_hash_vec", 8, false)?,
2910 ),
2911 (
2912 VecProducer::XxHash3Vec,
2913 declare_words("jit_xxhash3_vec", 8, false)?,
2914 ),
2915 (
2916 VecProducer::RegToVec,
2917 declare_words("jit_reg_to_vec_f32", 8, false)?,
2918 ),
2919 ];
2920 let vec_reducer_ids = [
2921 (VecReducer::Dot, declare_words("jit_vec_dot", 4, true)?),
2922 (VecReducer::L2, declare_words("jit_vec_l2", 4, true)?),
2923 (
2924 VecReducer::Cosine,
2925 declare_words("jit_vec_cosine", 4, true)?,
2926 ),
2927 (VecReducer::LidMle, declare_words("jit_lid_mle", 4, true)?),
2928 ];
2929 let reg_lane_ids = [
2930 (
2931 RegLaneRead::F32,
2932 declare_words("jit_reg_lane_f32", 3, true)?,
2933 ),
2934 (
2935 RegLaneRead::I16,
2936 declare_words("jit_reg_lane_i16", 3, true)?,
2937 ),
2938 (
2939 RegLaneRead::I64,
2940 declare_words("jit_reg_lane_i64", 3, true)?,
2941 ),
2942 ];
2943 let reg_producer_ids = [
2944 (
2945 RegProducer::WithLaneF32,
2946 declare_words("jit_reg_with_lane_f32", 6, false)?,
2947 ),
2948 (
2949 RegProducer::GatherF32,
2950 declare_words("jit_reg_gather_f32", 6, false)?,
2951 ),
2952 (
2953 RegProducer::VecToRegF32,
2954 declare_words("jit_vec_to_reg_f32", 6, false)?,
2955 ),
2956 (
2957 RegProducer::MulI8,
2958 declare_words("jit_reg_mul_i8", 6, false)?,
2959 ),
2960 ];
2961
2962 let mut defined: Vec<(cranelift_module::FuncId, bool)> = Vec::with_capacity(functions.len());
2967 for (function_idx, &(steps, dispatch)) in functions.iter().enumerate() {
2968 let mut sig = module.make_signature();
2972 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() {
2976 sig.params.push(AbiParam::new(types::I64)); sig.params.push(AbiParam::new(types::I64)); sig.params.push(AbiParam::new(types::I64)); }
2980 let func_id = module
2981 .declare_function(
2982 &format!("polydat_kernel_{function_idx}"),
2983 Linkage::Local,
2984 &sig,
2985 )
2986 .map_err(|e| format!("declare kernel: {e}"))?;
2987
2988 let mut ctx = module.make_context();
2989 ctx.func.signature = sig;
2990
2991 let mut fb_ctx = FunctionBuilderContext::new();
2992 {
2993 let mut builder = FunctionBuilder::new(&mut ctx.func, &mut fb_ctx);
2994 let block = builder.create_block();
2995 builder.append_block_params_for_function_params(block);
2996 builder.switch_to_block(block);
2997 builder.seal_block(block);
2998
2999 let _coords_ptr = builder.block_params(block)[0];
3000 let buffer_ptr = builder.block_params(block)[1];
3001 let scratch_ptr = builder.block_params(block)[2];
3002
3003 let hash_func_ref = module.declare_func_in_func(hash_func_id, builder.func);
3005 let interleave_func_ref = module.declare_func_in_func(interleave_func_id, builder.func);
3006 let shuffle_func_ref = module.declare_func_in_func(shuffle_func_id, builder.func);
3007 let lut_sample_func_ref = module.declare_func_in_func(lut_sample_func_id, builder.func);
3008 let weighted_pick_func_ref =
3009 module.declare_func_in_func(weighted_pick_func_id, builder.func);
3010 let is_positive_fail_ref =
3011 module.declare_func_in_func(is_positive_fail_id, builder.func);
3012 let in_range_fail_ref = module.declare_func_in_func(in_range_fail_id, builder.func);
3013 let div_zero_fail_ref = module.declare_func_in_func(div_zero_fail_id, builder.func);
3014 let is_one_of_fail_ref = module.declare_func_in_func(is_one_of_fail_id, builder.func);
3015 let slot_call_ref = module.declare_func_in_func(slot_call_id, builder.func);
3016 let u64_to_str_ref = module.declare_func_in_func(u64_to_str_id, builder.func);
3017 let i64_to_str_ref = module.declare_func_in_func(i64_to_str_id, builder.func);
3018 let f64_to_str_ref = module.declare_func_in_func(f64_to_str_id, builder.func);
3019 let str_concat_ref = module.declare_func_in_func(str_concat_id, builder.func);
3020 let json_to_str_ref = module.declare_func_in_func(json_to_str_id, builder.func);
3021 let vec_producer_refs: Vec<(VecProducer, ir::FuncRef)> = vec_producer_ids
3022 .iter()
3023 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
3024 .collect();
3025 let vec_reducer_refs: Vec<(VecReducer, ir::FuncRef)> = vec_reducer_ids
3026 .iter()
3027 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
3028 .collect();
3029 let reg_lane_refs: Vec<(RegLaneRead, ir::FuncRef)> = reg_lane_ids
3030 .iter()
3031 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
3032 .collect();
3033 let reg_producer_refs: Vec<(RegProducer, ir::FuncRef)> = reg_producer_ids
3034 .iter()
3035 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
3036 .collect();
3037 let pcg_func_ref = module.declare_func_in_func(pcg_func_id, builder.func);
3038 let pcg_stream_func_ref = module.declare_func_in_func(pcg_stream_func_id, builder.func);
3039 let n_of_func_ref = module.declare_func_in_func(n_of_func_id, builder.func);
3040 let cycle_walk_func_ref = module.declare_func_in_func(cycle_walk_func_id, builder.func);
3041 let perlin_1d_func_ref = module.declare_func_in_func(perlin_1d_func_id, builder.func);
3042 let perlin_2d_func_ref = module.declare_func_in_func(perlin_2d_func_id, builder.func);
3043 let simplex_2d_func_ref = module.declare_func_in_func(simplex_2d_func_id, builder.func);
3044 let fractal_noise_1d_func_ref =
3045 module.declare_func_in_func(fractal_noise_1d_func_id, builder.func);
3046 let fractal_noise_2d_func_ref =
3047 module.declare_func_in_func(fractal_noise_2d_func_id, builder.func);
3048 let thread_id_func_ref = module.declare_func_in_func(thread_id_func_id, builder.func);
3049 let current_epoch_millis_func_ref =
3050 module.declare_func_in_func(current_epoch_millis_func_id, builder.func);
3051 let math_unary_refs: Vec<_> = math_unary_ids
3052 .iter()
3053 .map(|id| module.declare_func_in_func(*id, builder.func))
3054 .collect();
3055 let math_binary_refs: Vec<_> = math_binary_ids
3056 .iter()
3057 .map(|id| module.declare_func_in_func(*id, builder.func))
3058 .collect();
3059 let everything: [Vec<usize>; 1] = [(0..steps.len()).collect()];
3067 let schedule: &[Vec<usize>] = dispatch.unwrap_or(&everything);
3068 let dispatcher = dispatch.map(|units| {
3069 let list_ptr = builder.block_params(block)[3];
3070 let list_len = builder.block_params(block)[4];
3071 let clean_ptr = builder.block_params(block)[5];
3072 let at = builder.create_sized_stack_slot(ir::StackSlotData::new(
3073 ir::StackSlotKind::ExplicitSlot,
3074 8,
3075 3,
3076 ));
3077 let zero = builder.ins().iconst(types::I64, 0);
3078 builder.ins().stack_store(zero, at, 0);
3079 let head = builder.create_block();
3080 let fetch = builder.create_block();
3081 let dispatch_unit = builder.create_block();
3082 let skip = builder.create_block();
3083 let exit = builder.create_block();
3084 let unit_blocks: Vec<ir::Block> =
3085 units.iter().map(|_| builder.create_block()).collect();
3086 builder.ins().jump(head, &[]);
3087 builder.switch_to_block(head);
3089 let i = builder.ins().stack_load(types::I64, at, 0);
3090 let done = builder.ins().icmp(
3091 ir::condcodes::IntCC::UnsignedGreaterThanOrEqual,
3092 i,
3093 list_len,
3094 );
3095 builder.ins().brif(done, exit, &[], fetch, &[]);
3096 builder.switch_to_block(fetch);
3099 builder.seal_block(fetch);
3100 let offset = builder.ins().ishl_imm(i, 2);
3101 let addr = builder.ins().iadd(list_ptr, offset);
3102 let unit = builder
3103 .ins()
3104 .load(types::I32, ir::MemFlags::trusted(), addr, 0);
3105 let unit_wide = builder.ins().uextend(types::I64, unit);
3106 let flag_addr = builder.ins().iadd(clean_ptr, unit_wide);
3107 let flag = builder
3108 .ins()
3109 .load(types::I8, ir::MemFlags::trusted(), flag_addr, 0);
3110 builder.ins().brif(flag, skip, &[], dispatch_unit, &[]);
3111 builder.switch_to_block(skip);
3113 builder.seal_block(skip);
3114 let next = builder.ins().iadd_imm(i, 1);
3115 builder.ins().stack_store(next, at, 0);
3116 builder.ins().jump(head, &[]);
3117 builder.switch_to_block(dispatch_unit);
3119 builder.seal_block(dispatch_unit);
3120 let default = builder.func.dfg.block_call(exit, &[]);
3121 let targets: Vec<ir::BlockCall> = unit_blocks
3122 .iter()
3123 .map(|&b| builder.func.dfg.block_call(b, &[]))
3124 .collect();
3125 let table = builder.create_jump_table(ir::JumpTableData::new(default, &targets));
3126 builder.ins().br_table(unit, table);
3127 (at, head, exit, unit_blocks, clean_ptr)
3128 });
3129
3130 for (unit_idx, members) in schedule.iter().enumerate() {
3132 if let Some((_, _, _, unit_blocks, _)) = &dispatcher {
3133 builder.switch_to_block(unit_blocks[unit_idx]);
3134 builder.seal_block(unit_blocks[unit_idx]);
3135 }
3136 for &step_idx in members {
3137 let (jit_op, input_slots, output_slots) = &steps[step_idx];
3138 let tracker_store = tracker.map(|t| {
3142 let idx = builder.ins().iconst(types::I64, step_idx as i64);
3143 let inst = store_slot(&mut builder, buffer_ptr, t, idx);
3144 (inst, builder.func.dfg.num_insts())
3145 });
3146 match jit_op {
3147 JitOp::Identity => {
3148 for (&i, &o) in input_slots.iter().zip(output_slots.iter()) {
3151 let val = load_slot(&mut builder, buffer_ptr, i);
3152 store_slot(&mut builder, buffer_ptr, o, val);
3153 }
3154 }
3155 JitOp::AddConst(c) => {
3156 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3157 let c_val = builder.ins().iconst(types::I64, *c as i64);
3158 let result = builder.ins().iadd(val, c_val);
3159 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3160 }
3161 JitOp::MulConst(c) => {
3162 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3163 let c_val = builder.ins().iconst(types::I64, *c as i64);
3164 let result = builder.ins().imul(val, c_val);
3165 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3166 }
3167 JitOp::DivConst(c) | JitOp::ModConst(c) => {
3168 let is_div = matches!(jit_op, JitOp::DivConst(_));
3171 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3172 if *c == 0 {
3173 let kind =
3174 builder.ins().iconst(types::I64, if is_div { 0 } else { 1 });
3175 let _ = builder.ins().call(div_zero_fail_ref, &[kind]);
3176 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3177 } else {
3178 let c_val = builder.ins().iconst(types::I64, *c as i64);
3179 let result = if is_div {
3180 builder.ins().udiv(val, c_val)
3181 } else {
3182 builder.ins().urem(val, c_val)
3183 };
3184 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3185 }
3186 }
3187 JitOp::U64DivWire | JitOp::U64ModWire => {
3188 let is_div = matches!(jit_op, JitOp::U64DivWire);
3192 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3193 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3194 let zero = builder.ins().iconst(types::I64, 0);
3195 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, b, zero);
3196 let fail_block = builder.create_block();
3197 let ok_block = builder.create_block();
3198 builder.ins().brif(is_zero, fail_block, &[], ok_block, &[]);
3199 builder.switch_to_block(fail_block);
3200 builder.seal_block(fail_block);
3201 let kind = builder.ins().iconst(types::I64, if is_div { 0 } else { 1 });
3202 let _ = builder.ins().call(div_zero_fail_ref, &[kind]);
3203 builder.ins().jump(ok_block, &[]);
3204 builder.switch_to_block(ok_block);
3205 builder.seal_block(ok_block);
3206 let result = if is_div {
3207 builder.ins().udiv(a, b)
3208 } else {
3209 builder.ins().urem(a, b)
3210 };
3211 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3212 }
3213 JitOp::ClampConst(min, max) => {
3214 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3215 let min_val = builder.ins().iconst(types::I64, *min as i64);
3216 let max_val = builder.ins().iconst(types::I64, *max as i64);
3217 let clamped_lo = builder.ins().umax(val, min_val);
3218 let clamped = builder.ins().umin(clamped_lo, max_val);
3219 store_slot(&mut builder, buffer_ptr, output_slots[0], clamped);
3220 }
3221 JitOp::Interleave => {
3222 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3223 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3224 let call = builder.ins().call(interleave_func_ref, &[a, b]);
3225 let result = builder.inst_results(call)[0];
3226 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3227 }
3228 JitOp::MixedRadixConst(radixes) => {
3229 let mut remainder = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3231 for (i, &radix) in radixes.iter().enumerate() {
3232 if radix == 0 {
3233 store_slot(
3235 &mut builder,
3236 buffer_ptr,
3237 output_slots[i],
3238 remainder,
3239 );
3240 } else {
3241 let r = builder.ins().iconst(types::I64, radix as i64);
3242 let digit = builder.ins().urem(remainder, r);
3243 store_slot(&mut builder, buffer_ptr, output_slots[i], digit);
3244 remainder = builder.ins().udiv(remainder, r);
3245 }
3246 }
3247 }
3248 JitOp::Hash => {
3249 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3250 let call = builder.ins().call(hash_func_ref, &[val]);
3251 let result = builder.inst_results(call)[0];
3252 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3253 }
3254 JitOp::SplitMix64 => {
3255 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3256 let c_gamma = builder
3257 .ins()
3258 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3259 let x1 = builder.ins().iadd(x0, c_gamma);
3260 let s30 = builder.ins().ushr_imm(x1, 30);
3261 let x2 = builder.ins().bxor(x1, s30);
3262 let c_m1 = builder
3263 .ins()
3264 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3265 let x3 = builder.ins().imul(x2, c_m1);
3266 let s27 = builder.ins().ushr_imm(x3, 27);
3267 let x4 = builder.ins().bxor(x3, s27);
3268 let c_m2 = builder
3269 .ins()
3270 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3271 let x5 = builder.ins().imul(x4, c_m2);
3272 let s31 = builder.ins().ushr_imm(x5, 31);
3273 let result = builder.ins().bxor(x5, s31);
3274 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3275 }
3276 JitOp::FairCoin => {
3277 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3278 let c_gamma = builder
3279 .ins()
3280 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3281 let x1 = builder.ins().iadd(x0, c_gamma);
3282 let s30 = builder.ins().ushr_imm(x1, 30);
3283 let x2 = builder.ins().bxor(x1, s30);
3284 let c_m1 = builder
3285 .ins()
3286 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3287 let x3 = builder.ins().imul(x2, c_m1);
3288 let s27 = builder.ins().ushr_imm(x3, 27);
3289 let x4 = builder.ins().bxor(x3, s27);
3290 let c_m2 = builder
3291 .ins()
3292 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3293 let x5 = builder.ins().imul(x4, c_m2);
3294 let s31 = builder.ins().ushr_imm(x5, 31);
3295 let h = builder.ins().bxor(x5, s31);
3296 let one = builder.ins().iconst(types::I64, 1);
3297 let result = builder.ins().band(h, one);
3298 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3299 }
3300 JitOp::CoinFlipConst(threshold) => {
3301 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3302 let thr = builder.ins().iconst(types::I64, *threshold as i64);
3303 let cmp =
3304 builder
3305 .ins()
3306 .icmp(ir::condcodes::IntCC::UnsignedLessThan, x, thr);
3307 let zero = builder.ins().iconst(types::I64, 0);
3308 let one = builder.ins().iconst(types::I64, 1);
3309 let result = builder.ins().select(cmp, one, zero);
3310 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3311 }
3312 JitOp::UnfairCoinConst(p_bits) => {
3313 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3314 let c_gamma = builder
3315 .ins()
3316 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3317 let x1 = builder.ins().iadd(x0, c_gamma);
3318 let s30 = builder.ins().ushr_imm(x1, 30);
3319 let x2 = builder.ins().bxor(x1, s30);
3320 let c_m1 = builder
3321 .ins()
3322 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3323 let x3 = builder.ins().imul(x2, c_m1);
3324 let s27 = builder.ins().ushr_imm(x3, 27);
3325 let x4 = builder.ins().bxor(x3, s27);
3326 let c_m2 = builder
3327 .ins()
3328 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3329 let x5 = builder.ins().imul(x4, c_m2);
3330 let s31 = builder.ins().ushr_imm(x5, 31);
3331 let h = builder.ins().bxor(x5, s31);
3332
3333 let fval = builder.ins().fcvt_from_uint(types::F64, h);
3334 let max_f = builder.ins().f64const(u64::MAX as f64);
3335 let unit = builder.ins().fdiv(fval, max_f);
3336 let p_f = builder.ins().f64const(f64::from_bits(*p_bits));
3337 let cmp =
3338 builder
3339 .ins()
3340 .fcmp(ir::condcodes::FloatCC::LessThan, unit, p_f);
3341 let zero = builder.ins().iconst(types::I64, 0);
3342 let one = builder.ins().iconst(types::I64, 1);
3343 let result = builder.ins().select(cmp, one, zero);
3344 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3345 }
3346 JitOp::ChanceConst(p_bits) => {
3347 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3348 let c_gamma = builder
3349 .ins()
3350 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3351 let x1 = builder.ins().iadd(x0, c_gamma);
3352 let s30 = builder.ins().ushr_imm(x1, 30);
3353 let x2 = builder.ins().bxor(x1, s30);
3354 let c_m1 = builder
3355 .ins()
3356 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3357 let x3 = builder.ins().imul(x2, c_m1);
3358 let s27 = builder.ins().ushr_imm(x3, 27);
3359 let x4 = builder.ins().bxor(x3, s27);
3360 let c_m2 = builder
3361 .ins()
3362 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3363 let x5 = builder.ins().imul(x4, c_m2);
3364 let s31 = builder.ins().ushr_imm(x5, 31);
3365 let h = builder.ins().bxor(x5, s31);
3366
3367 let fval = builder.ins().fcvt_from_uint(types::F64, h);
3368 let max_f = builder.ins().f64const(u64::MAX as f64);
3369 let unit = builder.ins().fdiv(fval, max_f);
3370 let p_f = builder.ins().f64const(f64::from_bits(*p_bits));
3371 let cmp =
3372 builder
3373 .ins()
3374 .fcmp(ir::condcodes::FloatCC::LessThan, unit, p_f);
3375 let zero_bits =
3376 builder.ins().iconst(types::I64, 0.0_f64.to_bits() as i64);
3377 let one_bits =
3378 builder.ins().iconst(types::I64, 1.0_f64.to_bits() as i64);
3379 let result = builder.ins().select(cmp, one_bits, zero_bits);
3380 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3381 }
3382 JitOp::ShuffleConst(feedback, size, min) => {
3383 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3384 let fb = builder.ins().iconst(types::I64, *feedback as i64);
3385 let sz = builder.ins().iconst(types::I64, *size as i64);
3386 let mn = builder.ins().iconst(types::I64, *min as i64);
3387 let call = builder.ins().call(shuffle_func_ref, &[val, fb, sz, mn]);
3388 let result = builder.inst_results(call)[0];
3389 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3390 }
3391
3392 JitOp::UnitInterval => {
3394 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3396 let fval = builder.ins().fcvt_from_uint(types::F64, val);
3397 let max_f = builder.ins().f64const(u64::MAX as f64);
3398 let result = builder.ins().fdiv(fval, max_f);
3399 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3400 }
3401 JitOp::F64ToU64 => {
3402 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3403 let result = builder.ins().fcvt_to_uint_sat(types::I64, fval);
3404 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3405 }
3406 JitOp::RoundToU64 => {
3407 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3408 let rounded = round_half_away(&mut builder, fval);
3409 let result = builder.ins().fcvt_to_uint_sat(types::I64, rounded);
3410 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3411 }
3412 JitOp::FloorToU64 => {
3413 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3414 let floored = builder.ins().floor(fval);
3415 let result = builder.ins().fcvt_to_uint_sat(types::I64, floored);
3416 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3417 }
3418 JitOp::CeilToU64 => {
3419 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3420 let ceiled = builder.ins().ceil(fval);
3421 let result = builder.ins().fcvt_to_uint_sat(types::I64, ceiled);
3422 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3423 }
3424 JitOp::ClampF64Const(min_bits, max_bits) => {
3425 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3426 let fmin = builder.ins().f64const(f64::from_bits(*min_bits));
3427 let fmax = builder.ins().f64const(f64::from_bits(*max_bits));
3428 let clamped = clamp_ir(&mut builder, fval, fmin, fmax);
3429 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], clamped);
3430 }
3431 JitOp::LerpConst(a_bits, b_bits) => {
3432 let t = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3434 let a = builder.ins().f64const(f64::from_bits(*a_bits));
3435 let b = builder.ins().f64const(f64::from_bits(*b_bits));
3436 let diff = builder.ins().fsub(b, a);
3437 let scaled = builder.ins().fmul(t, diff);
3438 let result = builder.ins().fadd(a, scaled);
3439 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3440 }
3441 JitOp::ScaleRangeConst(min_bits, range_bits) => {
3442 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3444 let fval = builder.ins().fcvt_from_uint(types::F64, val);
3445 let max_f = builder.ins().f64const(u64::MAX as f64);
3446 let t = builder.ins().fdiv(fval, max_f);
3447 let fmin = builder.ins().f64const(f64::from_bits(*min_bits));
3448 let frange = builder.ins().f64const(f64::from_bits(*range_bits));
3449 let scaled = builder.ins().fmul(t, frange);
3450 let result = builder.ins().fadd(fmin, scaled);
3451 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3452 }
3453 JitOp::QuantizeConst(step_bits) => {
3454 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3456 let step = builder.ins().f64const(f64::from_bits(*step_bits));
3457 let divided = builder.ins().fdiv(fval, step);
3458 let rounded = round_half_away(&mut builder, divided);
3459 let result = builder.ins().fmul(rounded, step);
3460 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3461 }
3462
3463 JitOp::LutSampleConst(lut_ptr, lut_len) => {
3464 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3466 let ptr_val = builder.ins().iconst(types::I64, *lut_ptr as i64);
3467 let len_val = builder.ins().iconst(types::I64, *lut_len as i64);
3468 let call = builder
3469 .ins()
3470 .call(lut_sample_func_ref, &[input, ptr_val, len_val]);
3471 let result = builder.inst_results(call)[0];
3472 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3473 }
3474 JitOp::DiscretizeConst(range_bits, buckets) => {
3475 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3480 let range = f64::from_bits(*range_bits);
3481 let fzero = builder.ins().f64const(0.0);
3482 let frange = builder.ins().f64const(range);
3483 let fbuckets = builder.ins().f64const(*buckets as f64);
3484 let clamped = clamp_ir(&mut builder, fval, fzero, frange);
3485 let divided = builder.ins().fdiv(clamped, frange);
3486 let scaled = builder.ins().fmul(divided, fbuckets);
3487 let as_u64 = builder.ins().fcvt_to_uint_sat(types::I64, scaled);
3488 let max_bucket =
3489 builder.ins().iconst(types::I64, (*buckets - 1) as i64);
3490 let result = builder.ins().umin(as_u64, max_bucket);
3491 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3492 }
3493
3494 JitOp::WeightedPickConst(
3495 values_ptr,
3496 biases_ptr,
3497 primaries_ptr,
3498 aliases_ptr,
3499 n,
3500 ) => {
3501 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3503 let v_ptr = builder.ins().iconst(types::I64, *values_ptr as i64);
3504 let b_ptr = builder.ins().iconst(types::I64, *biases_ptr as i64);
3505 let p_ptr = builder.ins().iconst(types::I64, *primaries_ptr as i64);
3506 let a_ptr = builder.ins().iconst(types::I64, *aliases_ptr as i64);
3507 let n_val = builder.ins().iconst(types::I64, *n as i64);
3508 let call = builder.ins().call(
3509 weighted_pick_func_ref,
3510 &[input, v_ptr, b_ptr, p_ptr, a_ptr, n_val],
3511 );
3512 let result = builder.inst_results(call)[0];
3513 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3514 }
3515
3516 JitOp::MathUnary(idx) => {
3517 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3518 let func_ref = math_unary_refs[*idx as usize];
3519 let call = builder.ins().call(func_ref, &[input]);
3520 let result = builder.inst_results(call)[0];
3521 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3522 }
3523
3524 JitOp::MathBinary(idx) => {
3525 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3526 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3527 let func_ref = math_binary_refs[*idx as usize];
3528 let call = builder.ins().call(func_ref, &[a, b]);
3529 let result = builder.inst_results(call)[0];
3530 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3531 }
3532
3533 JitOp::ToF64 => {
3534 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3535 let fval = builder.ins().fcvt_from_uint(types::F64, val);
3536 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], fval);
3537 }
3538
3539 JitOp::RegBinOp(lane, arith) => {
3541 let vt = reg_lane_type(*lane);
3542 let a = load_reg128(&mut builder, buffer_ptr, input_slots[0], vt);
3543 let b = load_reg128(&mut builder, buffer_ptr, input_slots[2], vt);
3544 let is_float = matches!(*lane, 4 | 5);
3545 let r = match (arith, is_float) {
3546 (0, false) => builder.ins().iadd(a, b),
3547 (1, false) => builder.ins().isub(a, b),
3548 (2, false) => builder.ins().imul(a, b),
3549 (0, true) => builder.ins().fadd(a, b),
3550 (1, true) => builder.ins().fsub(a, b),
3551 (2, true) => builder.ins().fmul(a, b),
3552 _ => unreachable!("RegBinOp arith index out of range"),
3553 };
3554 store_reg128(&mut builder, buffer_ptr, output_slots[0], r);
3555 }
3556 JitOp::RegCopy => {
3557 let v =
3558 load_reg128(&mut builder, buffer_ptr, input_slots[0], types::I64X2);
3559 store_reg128(&mut builder, buffer_ptr, output_slots[0], v);
3560 }
3561 JitOp::RegSplat(lane) => {
3562 let vt = reg_lane_type(*lane);
3563 let scalar = match *lane {
3564 0 => {
3566 let v = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3567 builder.ins().ireduce(types::I8, v)
3568 }
3569 1 => {
3570 let v = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3571 builder.ins().ireduce(types::I16, v)
3572 }
3573 2 => {
3574 let v = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3575 builder.ins().ireduce(types::I32, v)
3576 }
3577 3 => load_slot(&mut builder, buffer_ptr, input_slots[0]),
3578 4 => {
3580 let f = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3581 builder.ins().fdemote(types::F32, f)
3582 }
3583 5 => load_slot_f64(&mut builder, buffer_ptr, input_slots[0]),
3584 _ => unreachable!("RegSplat lane index out of range"),
3585 };
3586 let v = builder.ins().splat(vt, scalar);
3587 store_reg128(&mut builder, buffer_ptr, output_slots[0], v);
3588 }
3589
3590 JitOp::U64Add2 => {
3592 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3593 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3594 let result = builder.ins().iadd(a, b);
3595 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3596 }
3597 JitOp::U64Sub2 => {
3598 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3599 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3600 let result = builder.ins().isub(a, b);
3601 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3602 }
3603 JitOp::U64Mul2 => {
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 result = builder.ins().imul(a, b);
3607 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3608 }
3609 JitOp::U64Div2 => {
3610 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3611 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3612 let zero = builder.ins().iconst(types::I64, 0);
3615 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, b, zero);
3616 let div_block = builder.create_block();
3617 let merge_block = builder.create_block();
3618 builder.append_block_param(merge_block, types::I64);
3619 builder
3620 .ins()
3621 .brif(is_zero, merge_block, &[zero], div_block, &[]);
3622 builder.switch_to_block(div_block);
3623 builder.seal_block(div_block);
3624 let div_result = builder.ins().udiv(a, b);
3625 builder.ins().jump(merge_block, &[div_result]);
3626 builder.switch_to_block(merge_block);
3627 builder.seal_block(merge_block);
3628 let result = builder.block_params(merge_block)[0];
3629 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3630 }
3631 JitOp::U64Mod2 => {
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 zero = builder.ins().iconst(types::I64, 0);
3637 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, b, zero);
3638 let rem_block = builder.create_block();
3639 let merge_block = builder.create_block();
3640 builder.append_block_param(merge_block, types::I64);
3641 builder
3642 .ins()
3643 .brif(is_zero, merge_block, &[zero], rem_block, &[]);
3644 builder.switch_to_block(rem_block);
3645 builder.seal_block(rem_block);
3646 let rem_result = builder.ins().urem(a, b);
3647 builder.ins().jump(merge_block, &[rem_result]);
3648 builder.switch_to_block(merge_block);
3649 builder.seal_block(merge_block);
3650 let result = builder.block_params(merge_block)[0];
3651 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3652 }
3653 JitOp::U64And => {
3654 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3655 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3656 let result = builder.ins().band(a, b);
3657 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3658 }
3659 JitOp::U64Or => {
3660 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3661 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3662 let result = builder.ins().bor(a, b);
3663 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3664 }
3665 JitOp::U64Xor => {
3666 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3667 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3668 let result = builder.ins().bxor(a, b);
3669 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3670 }
3671 JitOp::U64Shl => {
3672 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3673 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3674 let result = builder.ins().ishl(a, b);
3675 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3676 }
3677 JitOp::U64Shr => {
3678 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3679 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3680 let result = builder.ins().ushr(a, b);
3681 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3682 }
3683 JitOp::U64Not => {
3684 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3685 let result = builder.ins().bnot(a);
3686 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3687 }
3688
3689 JitOp::F64Add => {
3691 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3692 let b = load_slot_f64(&mut builder, buffer_ptr, input_slots[1]);
3693 let result = builder.ins().fadd(a, b);
3694 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3695 }
3696 JitOp::F64Sub => {
3697 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3698 let b = load_slot_f64(&mut builder, buffer_ptr, input_slots[1]);
3699 let result = builder.ins().fsub(a, b);
3700 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3701 }
3702 JitOp::F64Mul => {
3703 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3704 let b = load_slot_f64(&mut builder, buffer_ptr, input_slots[1]);
3705 let result = builder.ins().fmul(a, b);
3706 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3707 }
3708 JitOp::F64Div => {
3709 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3710 let b = load_slot_f64(&mut builder, buffer_ptr, input_slots[1]);
3711 let zero = builder.ins().f64const(0.0);
3713 let is_zero =
3714 builder.ins().fcmp(ir::condcodes::FloatCC::Equal, b, zero);
3715 let div_result = builder.ins().fdiv(a, b);
3716 let result = builder.ins().select(is_zero, zero, div_result);
3717 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3718 }
3719 JitOp::F64Mod => {
3720 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3722 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3723 let call = builder
3724 .ins()
3725 .call(math_binary_refs[F64_MOD_HELPER], &[a, b]);
3726 let result = builder.inst_results(call)[0];
3727 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3728 }
3729
3730 JitOp::IsPositiveCheck { name_ptr, name_len } => {
3731 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3737 let zero = builder.ins().iconst(types::I64, 0);
3738 let is_zero =
3739 builder.ins().icmp(ir::condcodes::IntCC::Equal, val, zero);
3740 let fail_block = builder.create_block();
3741 let ok_block = builder.create_block();
3742 builder.ins().brif(is_zero, fail_block, &[], ok_block, &[]);
3743
3744 builder.switch_to_block(fail_block);
3745 builder.seal_block(fail_block);
3746 let np = builder.ins().iconst(types::I64, *name_ptr as i64);
3747 let nl = builder.ins().iconst(types::I64, *name_len as i64);
3748 let _ = builder.ins().call(is_positive_fail_ref, &[val, np, nl]);
3749 builder.ins().jump(ok_block, &[]);
3753
3754 builder.switch_to_block(ok_block);
3755 builder.seal_block(ok_block);
3756 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3757 }
3758
3759 JitOp::InRangeCheck(lo, hi) => {
3760 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3764 let lo_v = builder.ins().iconst(types::I64, *lo as i64);
3765 let hi_v = builder.ins().iconst(types::I64, *hi as i64);
3766 let below = builder.ins().icmp(
3767 ir::condcodes::IntCC::UnsignedLessThan,
3768 val,
3769 lo_v,
3770 );
3771 let above = builder.ins().icmp(
3772 ir::condcodes::IntCC::UnsignedGreaterThan,
3773 val,
3774 hi_v,
3775 );
3776 let out_of_range = builder.ins().bor(below, above);
3777
3778 let fail_block = builder.create_block();
3779 let ok_block = builder.create_block();
3780 builder
3781 .ins()
3782 .brif(out_of_range, fail_block, &[], ok_block, &[]);
3783
3784 builder.switch_to_block(fail_block);
3785 builder.seal_block(fail_block);
3786 let _ = builder.ins().call(in_range_fail_ref, &[val, lo_v, hi_v]);
3787 builder.ins().jump(ok_block, &[]);
3788
3789 builder.switch_to_block(ok_block);
3790 builder.seal_block(ok_block);
3791 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3792 }
3793
3794 JitOp::IsOneOfCheck {
3795 allowed,
3796 set_ptr,
3797 set_len,
3798 } => {
3799 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3805 let mut any_match = builder.ins().iconst(types::I8, 0);
3806 for allow in allowed.iter() {
3807 let c = builder.ins().iconst(types::I64, *allow as i64);
3808 let eq = builder.ins().icmp(ir::condcodes::IntCC::Equal, val, c);
3809 any_match = builder.ins().bor(any_match, eq);
3810 }
3811 let fail_block = builder.create_block();
3812 let ok_block = builder.create_block();
3813 builder
3817 .ins()
3818 .brif(any_match, ok_block, &[], fail_block, &[]);
3819
3820 builder.switch_to_block(fail_block);
3821 builder.seal_block(fail_block);
3822 let sp = builder.ins().iconst(types::I64, *set_ptr as i64);
3823 let sl = builder.ins().iconst(types::I64, *set_len as i64);
3824 let _ = builder.ins().call(is_one_of_fail_ref, &[val, sp, sl]);
3825 builder.ins().jump(ok_block, &[]);
3826
3827 builder.switch_to_block(ok_block);
3828 builder.seal_block(ok_block);
3829 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3830 }
3831
3832 JitOp::U64Cmp(cc) => {
3833 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3834 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3835 let cmp = builder.ins().icmp(*cc, a, b);
3836 let zero = builder.ins().iconst(types::I64, 0);
3837 let one = builder.ins().iconst(types::I64, 1);
3838 let result = builder.ins().select(cmp, one, zero);
3839 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3840 }
3841 JitOp::F64Cmp(cc) => {
3842 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3843 let b = load_slot_f64(
3844 &mut builder,
3845 buffer_ptr,
3846 if input_slots.len() > 1 {
3847 input_slots[1]
3848 } else {
3849 input_slots[0]
3850 },
3851 );
3852 let cmp = builder.ins().fcmp(*cc, a, b);
3853 let zero = builder.ins().iconst(types::I64, 0);
3854 let one = builder.ins().iconst(types::I64, 1);
3855 let result = builder.ins().select(cmp, one, zero);
3856 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3857 }
3858 JitOp::SelectU64 => {
3859 let cond = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3860 let a = load_slot(
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(
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(&mut builder, buffer_ptr, output_slots[0], result);
3885 }
3886 JitOp::SelectF64 => {
3887 let cond = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3888 let a = load_slot_f64(
3889 &mut builder,
3890 buffer_ptr,
3891 if input_slots.len() > 1 {
3892 input_slots[1]
3893 } else {
3894 input_slots[0]
3895 },
3896 );
3897 let b = load_slot_f64(
3898 &mut builder,
3899 buffer_ptr,
3900 if input_slots.len() > 2 {
3901 input_slots[2]
3902 } else {
3903 input_slots[0]
3904 },
3905 );
3906 let zero = builder.ins().iconst(types::I64, 0);
3907 let is_nonzero =
3908 builder
3909 .ins()
3910 .icmp(ir::condcodes::IntCC::NotEqual, cond, zero);
3911 let result = builder.ins().select(is_nonzero, a, b);
3912 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3913 }
3914
3915 JitOp::I64ToF64 => {
3916 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3917 let fval = builder.ins().fcvt_from_sint(types::F64, val);
3918 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], fval);
3919 }
3920 JitOp::ToBool => {
3921 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3922 let zero = builder.ins().iconst(types::I64, 0);
3923 let one = builder.ins().iconst(types::I64, 1);
3924 let cmp = builder
3925 .ins()
3926 .icmp(ir::condcodes::IntCC::NotEqual, val, zero);
3927 let result = builder.ins().select(cmp, one, zero);
3928 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3929 }
3930 JitOp::ConstU64(v) | JitOp::ConstF64(v) => {
3931 let result = builder.ins().iconst(types::I64, *v as i64);
3932 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3933 }
3934 JitOp::HashRangeConst(max) => {
3935 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3936 let c_gamma = builder
3937 .ins()
3938 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3939 let x1 = builder.ins().iadd(input, c_gamma);
3940 let s30 = builder.ins().ushr_imm(x1, 30);
3941 let x2 = builder.ins().bxor(x1, s30);
3942 let c_m1 = builder
3943 .ins()
3944 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3945 let x3 = builder.ins().imul(x2, c_m1);
3946 let s27 = builder.ins().ushr_imm(x3, 27);
3947 let x4 = builder.ins().bxor(x3, s27);
3948 let c_m2 = builder
3949 .ins()
3950 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3951 let x5 = builder.ins().imul(x4, c_m2);
3952 let s31 = builder.ins().ushr_imm(x5, 31);
3953 let h = builder.ins().bxor(x5, s31);
3954 if *max == 0 {
3955 let zero = builder.ins().iconst(types::I64, 0);
3956 store_slot(&mut builder, buffer_ptr, output_slots[0], zero);
3957 } else {
3958 let m = builder.ins().iconst(types::I64, *max as i64);
3959 let rem = builder.ins().urem(h, m);
3960 store_slot(&mut builder, buffer_ptr, output_slots[0], rem);
3961 }
3962 }
3963 JitOp::HashIntervalConst(min_bits, max_bits) => {
3964 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3965 let c_gamma = builder
3966 .ins()
3967 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3968 let x1 = builder.ins().iadd(input, c_gamma);
3969 let s30 = builder.ins().ushr_imm(x1, 30);
3970 let x2 = builder.ins().bxor(x1, s30);
3971 let c_m1 = builder
3972 .ins()
3973 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3974 let x3 = builder.ins().imul(x2, c_m1);
3975 let s27 = builder.ins().ushr_imm(x3, 27);
3976 let x4 = builder.ins().bxor(x3, s27);
3977 let c_m2 = builder
3978 .ins()
3979 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3980 let x5 = builder.ins().imul(x4, c_m2);
3981 let s31 = builder.ins().ushr_imm(x5, 31);
3982 let h = builder.ins().bxor(x5, s31);
3983
3984 let h_f = builder.ins().fcvt_from_uint(types::F64, h);
3985 let denom = builder.ins().f64const(u64::MAX as f64);
3986 let unit = builder.ins().fdiv(h_f, denom);
3987 let min_f = f64::from_bits(*min_bits);
3988 let max_f = f64::from_bits(*max_bits);
3989 let span = builder.ins().f64const(max_f - min_f);
3990 let min_val = builder.ins().f64const(min_f);
3991 let scaled = builder.ins().fmul(unit, span);
3992 let res_f = builder.ins().fadd(min_val, scaled);
3993 let res = builder
3994 .ins()
3995 .bitcast(types::I64, ir::MemFlags::new(), res_f);
3996 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3997 }
3998 JitOp::InvLerpConst(a_bits, b_bits) => {
3999 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4000 let in_f =
4001 builder
4002 .ins()
4003 .bitcast(types::F64, ir::MemFlags::new(), input);
4004 let a_f = f64::from_bits(*a_bits);
4005 let b_f = f64::from_bits(*b_bits);
4006 let a_val = builder.ins().f64const(a_f);
4007 let inv_span = builder.ins().f64const(1.0 / (b_f - a_f));
4011 let diff = builder.ins().fsub(in_f, a_val);
4012 let t = builder.ins().fmul(diff, inv_span);
4013 let zero = builder.ins().f64const(0.0);
4014 let one = builder.ins().f64const(1.0);
4015 let res_f = clamp_ir(&mut builder, t, zero, one);
4016 let res = builder
4017 .ins()
4018 .bitcast(types::I64, ir::MemFlags::new(), res_f);
4019 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4020 }
4021 JitOp::RemapConst(in_min_bits, in_max_bits, out_min_bits, out_max_bits) => {
4022 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4023 let in_f =
4024 builder
4025 .ins()
4026 .bitcast(types::F64, ir::MemFlags::new(), input);
4027 let in_min = f64::from_bits(*in_min_bits);
4028 let in_max = f64::from_bits(*in_max_bits);
4029 let out_min = f64::from_bits(*out_min_bits);
4030 let out_max = f64::from_bits(*out_max_bits);
4031 let in_span_val = builder.ins().f64const(in_max - in_min);
4035 let in_min_val = builder.ins().f64const(in_min);
4036 let out_min_val = builder.ins().f64const(out_min);
4037 let out_span_val = builder.ins().f64const(out_max - out_min);
4038 let diff = builder.ins().fsub(in_f, in_min_val);
4039 let t = builder.ins().fdiv(diff, in_span_val);
4040 let scaled = builder.ins().fmul(t, out_span_val);
4041 let res_f = builder.ins().fadd(out_min_val, scaled);
4042 let res = builder
4043 .ins()
4044 .bitcast(types::I64, ir::MemFlags::new(), res_f);
4045 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4046 }
4047 JitOp::EpochOffsetConst(base) => {
4048 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4049 let b = builder.ins().iconst(types::I64, *base as i64);
4050 let res = builder.ins().iadd(val, b);
4051 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4052 }
4053 JitOp::EpochScaleConst(factor) => {
4054 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4055 let f = builder.ins().iconst(types::I64, *factor as i64);
4056 let res = builder.ins().imul(val, f);
4057 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4058 }
4059 JitOp::ThreadId => {
4060 let call = builder.ins().call(thread_id_func_ref, &[]);
4061 let res = builder.inst_results(call)[0];
4062 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4063 }
4064 JitOp::CurrentEpochMillis => {
4065 let call = builder.ins().call(current_epoch_millis_func_ref, &[]);
4066 let res = builder.inst_results(call)[0];
4067 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4068 }
4069 JitOp::Perlin1dConst(perm_ptr, freq_bits) => {
4070 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4071 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
4072 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
4073 let call = builder.ins().call(perlin_1d_func_ref, &[input, p, fb]);
4074 let res = builder.inst_results(call)[0];
4075 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4076 }
4077 JitOp::Perlin2dConst(perm_ptr, freq_bits) => {
4078 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4079 let y = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4080 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
4081 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
4082 let call = builder.ins().call(perlin_2d_func_ref, &[x, y, p, fb]);
4083 let res = builder.inst_results(call)[0];
4084 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4085 }
4086 JitOp::Simplex2dConst(perm_ptr, freq_bits) => {
4087 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4088 let y = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4089 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
4090 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
4091 let call = builder.ins().call(simplex_2d_func_ref, &[x, y, p, fb]);
4092 let res = builder.inst_results(call)[0];
4093 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4094 }
4095 JitOp::FractalNoise1dConst(perm_ptr, freq_bits, octaves) => {
4096 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4097 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
4098 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
4099 let oct = builder.ins().iconst(types::I64, *octaves as i64);
4100 let call = builder
4101 .ins()
4102 .call(fractal_noise_1d_func_ref, &[input, p, fb, oct]);
4103 let res = builder.inst_results(call)[0];
4104 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4105 }
4106 JitOp::FractalNoise2dConst(perm_ptr, freq_bits, octaves) => {
4107 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4108 let y = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4109 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
4110 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
4111 let oct = builder.ins().iconst(types::I64, *octaves as i64);
4112 let call = builder
4113 .ins()
4114 .call(fractal_noise_2d_func_ref, &[x, y, p, fb, oct]);
4115 let res = builder.inst_results(call)[0];
4116 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4117 }
4118 JitOp::CycleWalkConst(range, seed, inc) => {
4119 let pos = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4120 let r = builder.ins().iconst(types::I64, *range as i64);
4121 let s = builder.ins().iconst(types::I64, *seed as i64);
4122 let i = builder.ins().iconst(types::I64, *inc as i64);
4123 let call = builder.ins().call(cycle_walk_func_ref, &[pos, r, s, i]);
4124 let res = builder.inst_results(call)[0];
4125 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4126 }
4127
4128 JitOp::VariadicSum => {
4129 if input_slots.is_empty() {
4130 let zero = builder.ins().iconst(types::I64, 0);
4131 store_slot(&mut builder, buffer_ptr, output_slots[0], zero);
4132 } else {
4133 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4134 for &slot in &input_slots[1..] {
4135 let v = load_slot(&mut builder, buffer_ptr, slot);
4136 acc = builder.ins().iadd(acc, v);
4137 }
4138 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4139 }
4140 }
4141 JitOp::VariadicProduct => {
4142 if input_slots.is_empty() {
4143 let one = builder.ins().iconst(types::I64, 1);
4144 store_slot(&mut builder, buffer_ptr, output_slots[0], one);
4145 } else {
4146 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4147 for &slot in &input_slots[1..] {
4148 let v = load_slot(&mut builder, buffer_ptr, slot);
4149 acc = builder.ins().imul(acc, v);
4150 }
4151 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4152 }
4153 }
4154 JitOp::VariadicMin => {
4155 if input_slots.is_empty() {
4156 let ident = builder.ins().iconst(types::I64, u64::MAX as i64);
4165 store_slot(&mut builder, buffer_ptr, output_slots[0], ident);
4166 } else {
4167 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4168 for &slot in &input_slots[1..] {
4169 let v = load_slot(&mut builder, buffer_ptr, slot);
4170 let cmp = builder.ins().icmp(
4171 ir::condcodes::IntCC::UnsignedLessThan,
4172 v,
4173 acc,
4174 );
4175 acc = builder.ins().select(cmp, v, acc);
4176 }
4177 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4178 }
4179 }
4180 JitOp::VariadicMax => {
4181 if input_slots.is_empty() {
4182 let zero = builder.ins().iconst(types::I64, 0);
4183 store_slot(&mut builder, buffer_ptr, output_slots[0], zero);
4184 } else {
4185 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4186 for &slot in &input_slots[1..] {
4187 let v = load_slot(&mut builder, buffer_ptr, slot);
4188 let cmp = builder.ins().icmp(
4189 ir::condcodes::IntCC::UnsignedGreaterThan,
4190 v,
4191 acc,
4192 );
4193 acc = builder.ins().select(cmp, v, acc);
4194 }
4195 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4196 }
4197 }
4198
4199 JitOp::CeilToMultiple => {
4200 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4201 let m = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4202 let zero = builder.ins().iconst(types::I64, 0);
4203 let one = builder.ins().iconst(types::I64, 1);
4204 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, m, zero);
4205 let calc_block = builder.create_block();
4206 let merge_block = builder.create_block();
4207 builder.append_block_param(merge_block, types::I64);
4208 builder
4209 .ins()
4210 .brif(is_zero, merge_block, &[val], calc_block, &[]);
4211 builder.switch_to_block(calc_block);
4212 builder.seal_block(calc_block);
4213 let div = div_ceil(&mut builder, val, m, one);
4216 let high = builder.ins().umulhi(div, m);
4217 let low = builder.ins().imul(div, m);
4218 let zero_hi = builder.ins().iconst(types::I64, 0);
4219 let overflows =
4220 builder
4221 .ins()
4222 .icmp(ir::condcodes::IntCC::NotEqual, high, zero_hi);
4223 let max = builder.ins().iconst(types::I64, -1);
4224 let mul = builder.ins().select(overflows, max, low);
4225 builder.ins().jump(merge_block, &[mul]);
4226 builder.switch_to_block(merge_block);
4227 builder.seal_block(merge_block);
4228 let result = builder.block_params(merge_block)[0];
4229 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4230 }
4231 JitOp::CheckedAdd => {
4232 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4233 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4234 let sum = builder.ins().iadd(a, b);
4235 let is_overflow =
4236 builder
4237 .ins()
4238 .icmp(ir::condcodes::IntCC::UnsignedLessThan, sum, a);
4239 let zero = builder.ins().iconst(types::I64, 0);
4240 let result = builder.ins().select(is_overflow, zero, sum);
4241 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4242 }
4243 JitOp::CheckedSub => {
4244 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4245 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4246 let is_lt =
4247 builder
4248 .ins()
4249 .icmp(ir::condcodes::IntCC::UnsignedLessThan, a, b);
4250 let diff = builder.ins().isub(a, b);
4251 let zero = builder.ins().iconst(types::I64, 0);
4252 let result = builder.ins().select(is_lt, zero, diff);
4253 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4254 }
4255 JitOp::CheckedMul => {
4256 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4257 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4258 let prod = builder.ins().imul(a, b);
4259 let zero = builder.ins().iconst(types::I64, 0);
4260 let a_is_zero =
4261 builder.ins().icmp(ir::condcodes::IntCC::Equal, a, zero);
4262 let div_block = builder.create_block();
4263 let merge_block = builder.create_block();
4264 builder.append_block_param(merge_block, types::I64);
4265 builder
4266 .ins()
4267 .brif(a_is_zero, merge_block, &[zero], div_block, &[]);
4268 builder.switch_to_block(div_block);
4269 builder.seal_block(div_block);
4270 let div = builder.ins().udiv(prod, a);
4271 let ok = builder.ins().icmp(ir::condcodes::IntCC::Equal, div, b);
4272 let mul_res = builder.ins().select(ok, prod, zero);
4273 builder.ins().jump(merge_block, &[mul_res]);
4274 builder.switch_to_block(merge_block);
4275 builder.seal_block(merge_block);
4276 let result = builder.block_params(merge_block)[0];
4277 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4278 }
4279 JitOp::MultiplesAtLeast => {
4280 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4281 let m = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4282 let zero = builder.ins().iconst(types::I64, 0);
4283 let one = builder.ins().iconst(types::I64, 1);
4284 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, m, zero);
4285 let calc_block = builder.create_block();
4286 let merge_block = builder.create_block();
4287 builder.append_block_param(merge_block, types::I64);
4288 builder
4289 .ins()
4290 .brif(is_zero, merge_block, &[zero], calc_block, &[]);
4291 builder.switch_to_block(calc_block);
4292 builder.seal_block(calc_block);
4293 let div = div_ceil(&mut builder, val, m, one);
4294 builder.ins().jump(merge_block, &[div]);
4295 builder.switch_to_block(merge_block);
4296 builder.seal_block(merge_block);
4297 let result = builder.block_params(merge_block)[0];
4298 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4299 }
4300
4301 JitOp::BlendConst(mix_bits) => {
4302 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4307 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4308 let fa = builder.ins().bitcast(types::F64, ir::MemFlags::new(), a);
4309 let fb = builder.ins().bitcast(types::F64, ir::MemFlags::new(), b);
4310 let mix_f64 = f64::from_bits(*mix_bits);
4311 let mix_val = builder.ins().f64const(mix_f64);
4312 let one = builder.ins().f64const(1.0);
4313 let one_minus_mix = builder.ins().fsub(one, mix_val);
4314 let a_part = builder.ins().fmul(fa, one_minus_mix);
4315 let b_part = builder.ins().fmul(fb, mix_val);
4316 let sum = builder.ins().fadd(a_part, b_part);
4317 let result =
4318 builder.ins().bitcast(types::I64, ir::MemFlags::new(), sum);
4319 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4320 }
4321 JitOp::LfsrStepConst(feedback) => {
4322 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4323 let feedback = builder.ins().iconst(types::I64, *feedback as i64);
4324 let one = builder.ins().iconst(types::I64, 1);
4325 let zero = builder.ins().iconst(types::I64, 0);
4326 let shifted = builder.ins().ushr(val, one);
4327 let lsb = builder.ins().band(val, one);
4328 let is_odd =
4329 builder
4330 .ins()
4331 .icmp(ir::condcodes::IntCC::NotEqual, lsb, zero);
4332 let fb_mask = builder.ins().select(is_odd, feedback, zero);
4333 let result = builder.ins().bxor(shifted, fb_mask);
4334 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4335 }
4336 JitOp::PcgConst(seed, stream) => {
4337 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4338 let s = builder.ins().iconst(types::I64, *seed as i64);
4339 let st = builder.ins().iconst(types::I64, *stream as i64);
4340 let call = builder.ins().call(pcg_func_ref, &[input, s, st]);
4341 let result = builder.inst_results(call)[0];
4342 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4343 }
4344 JitOp::PcgStreamConst(seed) => {
4345 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4346 let st = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4347 let s = builder.ins().iconst(types::I64, *seed as i64);
4348 let call = builder.ins().call(pcg_stream_func_ref, &[input, st, s]);
4349 let result = builder.inst_results(call)[0];
4350 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4351 }
4352 JitOp::NOfConst(n, m) => {
4353 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4354 let n_val = builder.ins().iconst(types::I64, *n as i64);
4355 let m_val = builder.ins().iconst(types::I64, *m as i64);
4356 let call = builder.ins().call(n_of_func_ref, &[input, n_val, m_val]);
4357 let result = builder.inst_results(call)[0];
4358 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4359 }
4360
4361 JitOp::SlotCall { kit, scratch_base } => {
4362 emit_slot_call(
4363 &mut builder,
4364 buffer_ptr,
4365 scratch_ptr,
4366 slot_call_ref,
4367 kit,
4368 *scratch_base,
4369 input_slots,
4370 output_slots,
4371 );
4372 }
4373
4374 JitOp::Convert {
4375 from,
4376 to,
4377 kit,
4378 scratch_base,
4379 } => {
4380 emit_conversion(
4381 &mut builder,
4382 buffer_ptr,
4383 input_slots[0],
4384 output_slots[0],
4385 *from,
4386 *to,
4387 |builder| {
4388 emit_slot_call(
4389 builder,
4390 buffer_ptr,
4391 scratch_ptr,
4392 slot_call_ref,
4393 kit,
4394 *scratch_base,
4395 input_slots,
4396 output_slots,
4397 )
4398 },
4399 );
4400 }
4401
4402 JitOp::U64ToStr { scratch_base }
4403 | JitOp::I64ToStr { scratch_base }
4404 | JitOp::F64ToStr { scratch_base } => {
4405 let func = match jit_op {
4408 JitOp::U64ToStr { .. } => u64_to_str_ref,
4409 JitOp::I64ToStr { .. } => i64_to_str_ref,
4410 _ => f64_to_str_ref,
4411 };
4412 let value = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4413 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4414 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4415 builder
4416 .ins()
4417 .call(func, &[scratch_ptr, base_v, buffer_ptr, out_v, value]);
4418 }
4419 JitOp::JsonToStr { scratch_base } => {
4420 let ptr = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4421 let len = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4422 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4423 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4424 builder.ins().call(
4425 json_to_str_ref,
4426 &[scratch_ptr, base_v, buffer_ptr, out_v, ptr, len],
4427 );
4428 }
4429 JitOp::StrConcat { scratch_base } => {
4430 let n_words = input_slots.len();
4433 let frame = builder.create_sized_stack_slot(ir::StackSlotData::new(
4434 ir::StackSlotKind::ExplicitSlot,
4435 (n_words.max(1) * 8) as u32,
4436 3,
4437 ));
4438 for (k, &s) in input_slots.iter().enumerate() {
4439 let v = load_slot(&mut builder, buffer_ptr, s);
4440 builder.ins().stack_store(v, frame, (k * 8) as i32);
4441 }
4442 let pairs_ptr = builder.ins().stack_addr(types::I64, frame, 0);
4443 let n_v = builder.ins().iconst(types::I64, (n_words / 2) as i64);
4444 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4445 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4446 builder.ins().call(
4447 str_concat_ref,
4448 &[scratch_ptr, base_v, buffer_ptr, out_v, pairs_ptr, n_v],
4449 );
4450 }
4451
4452 JitOp::VecProduce { kind, scratch_base } => {
4453 let func = func_of(&vec_producer_refs, *kind);
4457 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4458 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4459 let words = load_words(&mut builder, buffer_ptr, input_slots, 4);
4460 let mut args = vec![scratch_ptr, base_v, buffer_ptr, out_v];
4461 args.extend(words);
4462 builder.ins().call(func, &args);
4463 }
4464 JitOp::VecReduce(kind) => {
4465 let func = func_of(&vec_reducer_refs, *kind);
4466 let words = load_words(&mut builder, buffer_ptr, input_slots, 4);
4467 let call = builder.ins().call(func, &words);
4468 let bits = builder.inst_results(call)[0];
4469 store_slot(&mut builder, buffer_ptr, output_slots[0], bits);
4470 }
4471 JitOp::RegLane(kind) => {
4472 let func = func_of(®_lane_refs, *kind);
4473 let words = load_words(&mut builder, buffer_ptr, input_slots, 3);
4474 let call = builder.ins().call(func, &words);
4475 let word = builder.inst_results(call)[0];
4476 store_slot(&mut builder, buffer_ptr, output_slots[0], word);
4477 }
4478 JitOp::RegProduce(kind) => {
4479 let func = func_of(®_producer_refs, *kind);
4480 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4481 let words = load_words(&mut builder, buffer_ptr, input_slots, 4);
4482 let mut args = vec![buffer_ptr, out_v];
4483 args.extend(words);
4484 builder.ins().call(func, &args);
4485 }
4486 JitOp::RegDotF32 => {
4487 let a =
4491 load_reg128(&mut builder, buffer_ptr, input_slots[0], types::F32X4);
4492 let b =
4493 load_reg128(&mut builder, buffer_ptr, input_slots[2], types::F32X4);
4494 let p = builder.ins().fmul(a, b);
4495 let p0 = builder.ins().extractlane(p, 0);
4496 let p1 = builder.ins().extractlane(p, 1);
4497 let p2 = builder.ins().extractlane(p, 2);
4498 let p3 = builder.ins().extractlane(p, 3);
4499 let s01 = builder.ins().fadd(p0, p1);
4500 let s23 = builder.ins().fadd(p2, p3);
4501 let s = builder.ins().fadd(s01, s23);
4502 let wide = builder.ins().fpromote(types::F64, s);
4503 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], wide);
4504 }
4505 JitOp::RegShuffleConst(mask) => {
4506 let x =
4510 load_reg128(&mut builder, buffer_ptr, input_slots[0], types::I8X16);
4511 let imm = builder
4512 .func
4513 .dfg
4514 .immediates
4515 .push(ir::ConstantData::from(&mask[..]));
4516 let r = builder.ins().shuffle(x, x, imm);
4517 store_reg128(&mut builder, buffer_ptr, output_slots[0], r);
4518 }
4519
4520 JitOp::Fallback => {
4521 }
4524 }
4525 if let Some((inst, mark)) = tracker_store {
4526 let calls = (mark..builder.func.dfg.num_insts()).any(|i| {
4527 builder.func.dfg.insts[ir::Inst::from_u32(i as u32)]
4528 .opcode()
4529 .is_call()
4530 });
4531 if !calls {
4532 builder.func.layout.remove_inst(inst);
4533 }
4534 }
4535 }
4536 if let Some((at, head, _, _, clean_ptr)) = &dispatcher {
4539 let one = builder.ins().iconst(types::I8, 1);
4540 builder
4541 .ins()
4542 .store(ir::MemFlags::trusted(), one, *clean_ptr, unit_idx as i32);
4543 let i = builder.ins().stack_load(types::I64, *at, 0);
4544 let next = builder.ins().iadd_imm(i, 1);
4545 builder.ins().stack_store(next, *at, 0);
4546 builder.ins().jump(*head, &[]);
4547 }
4548 }
4549
4550 if let Some((_, head, exit, _, _)) = dispatcher {
4551 builder.seal_block(head);
4554 builder.switch_to_block(exit);
4555 builder.seal_block(exit);
4556 }
4557 builder.ins().return_(&[]);
4558 builder.finalize();
4559 }
4560 let fallible = ctx.func.layout.blocks().any(|block| {
4563 ctx.func
4564 .layout
4565 .block_insts(block)
4566 .any(|inst| ctx.func.dfg.insts[inst].opcode().is_call())
4567 });
4568
4569 module
4570 .define_function(func_id, &mut ctx)
4571 .map_err(|e| format!("define function: {e}"))?;
4572 module.clear_context(&mut ctx);
4573 defined.push((func_id, fallible));
4574 }
4575 module
4576 .finalize_definitions()
4577 .map_err(|e| format!("finalize: {e}"))?;
4578
4579 let entries: Vec<JitEntry> = defined
4580 .iter()
4581 .map(|&(func_id, fallible)| {
4582 let code_ptr = module.get_finalized_function(func_id);
4583 let straight_fn: NativeFn = unsafe { mem::transmute(code_ptr) };
4584 let dispatch_fn: NativeDispatchFn = unsafe { mem::transmute(code_ptr) };
4585 (straight_fn, dispatch_fn, fallible)
4586 })
4587 .collect();
4588 let kits: Vec<SlotKitRef> = functions
4590 .iter()
4591 .flat_map(|(steps, _)| steps.iter())
4592 .filter_map(|(op, _, _)| op.slot_kit().cloned())
4593 .collect();
4594 let any_fallible = defined.iter().any(|&(_, f)| f);
4595 let code = super::kernels::JitCode::new(module, kits, any_fallible);
4596 Ok((entries, code))
4597}
4598
4599#[allow(clippy::too_many_arguments)]
4608fn emit_slot_call(
4609 builder: &mut FunctionBuilder,
4610 buffer_ptr: ir::Value,
4611 scratch_ptr: ir::Value,
4612 slot_call_ref: ir::FuncRef,
4613 kit: &SlotKitRef,
4614 scratch_base: usize,
4615 input_slots: &[usize],
4616 output_slots: &[usize],
4617) {
4618 let n_in = input_slots.len();
4619 let n_out = output_slots.len();
4620 let frame = |builder: &mut FunctionBuilder, n: usize| {
4621 builder.create_sized_stack_slot(ir::StackSlotData::new(
4622 ir::StackSlotKind::ExplicitSlot,
4623 (n.max(1) * 8) as u32,
4624 3,
4625 ))
4626 };
4627 let in_frame = frame(builder, n_in);
4628 let out_frame = frame(builder, n_out);
4629 for (k, &s) in input_slots.iter().enumerate() {
4630 let v = load_slot(builder, buffer_ptr, s);
4631 builder.ins().stack_store(v, in_frame, (k * 8) as i32);
4632 }
4633 let kit_ptr = builder
4634 .ins()
4635 .iconst(types::I64, std::sync::Arc::as_ptr(&kit.0) as usize as i64);
4636 let in_ptr = builder.ins().stack_addr(types::I64, in_frame, 0);
4637 let n_in_v = builder.ins().iconst(types::I64, n_in as i64);
4638 let out_ptr = builder.ins().stack_addr(types::I64, out_frame, 0);
4639 let n_out_v = builder.ins().iconst(types::I64, n_out as i64);
4640 let base_v = builder.ins().iconst(types::I64, scratch_base as i64);
4641 let n_sc_v = builder.ins().iconst(types::I64, kit.0.scratch.len() as i64);
4642 builder.ins().call(
4643 slot_call_ref,
4644 &[
4645 kit_ptr,
4646 in_ptr,
4647 n_in_v,
4648 out_ptr,
4649 n_out_v,
4650 scratch_ptr,
4651 base_v,
4652 n_sc_v,
4653 ],
4654 );
4655 for (k, &s) in output_slots.iter().enumerate() {
4656 let v = builder
4657 .ins()
4658 .stack_load(types::I64, out_frame, (k * 8) as i32);
4659 store_slot(builder, buffer_ptr, s, v);
4660 }
4661}
4662
4663fn emit_conversion(
4676 builder: &mut FunctionBuilder,
4677 buffer_ptr: ir::Value,
4678 in_slot: usize,
4679 out_slot: usize,
4680 from: crate::ast::PortType,
4681 to: crate::ast::PortType,
4682 slow: impl FnOnce(&mut FunctionBuilder),
4683) {
4684 use Scalar::{Bool, F32, F64, Signed, Unsigned};
4685 use ir::condcodes::{FloatCC, IntCC};
4686 let (Some(src), Some(dst)) = (Scalar::of(from), Scalar::of(to)) else {
4687 slow(builder);
4688 return;
4689 };
4690 let raw = load_slot(builder, buffer_ptr, in_slot);
4691
4692 let float_of = |builder: &mut FunctionBuilder| -> ir::Value {
4694 match src {
4695 F32 => {
4696 let bits = builder.ins().ireduce(types::I32, raw);
4697 let x = builder.ins().bitcast(types::F32, ir::MemFlags::new(), bits);
4698 builder.ins().fpromote(types::F64, x)
4699 }
4700 _ => builder.ins().bitcast(types::F64, ir::MemFlags::new(), raw),
4701 }
4702 };
4703 let store_float = |builder: &mut FunctionBuilder, x: ir::Value| {
4706 let word = if dst == F32 {
4707 let bits = builder.ins().bitcast(types::I32, ir::MemFlags::new(), x);
4708 builder.ins().uextend(types::I64, bits)
4709 } else {
4710 builder.ins().bitcast(types::I64, ir::MemFlags::new(), x)
4711 };
4712 store_slot(builder, buffer_ptr, out_slot, word);
4713 };
4714
4715 match (src, dst) {
4716 (_, Bool) => {
4719 let truth = match src {
4720 F32 | F64 => {
4721 let x = float_of(builder);
4726 let zero = builder.ins().f64const(0.0);
4727 let nonzero = builder.ins().fcmp(FloatCC::NotEqual, x, zero);
4728 let ordered = builder.ins().fcmp(FloatCC::Ordered, x, x);
4729 builder.ins().band(nonzero, ordered)
4730 }
4731 _ => builder.ins().icmp_imm(IntCC::NotEqual, raw, 0),
4732 };
4733 let word = builder.ins().uextend(types::I64, truth);
4734 store_slot(builder, buffer_ptr, out_slot, word);
4735 }
4736 (s, d) if s.int_range().is_some() && d.int_range().is_some() => {
4739 let (smin, smax) = s.int_range().expect("an integer");
4740 let (dmin, dmax) = d.int_range().expect("an integer");
4741 let mut fits = Vec::new();
4742 if dmin > smin {
4743 fits.push(builder.ins().icmp_imm(
4745 IntCC::SignedGreaterThanOrEqual,
4746 raw,
4747 dmin as i64,
4748 ));
4749 }
4750 if dmax < smax {
4751 let cc = if matches!(s, Signed(_)) {
4752 IntCC::SignedLessThanOrEqual
4753 } else {
4754 IntCC::UnsignedLessThanOrEqual
4755 };
4756 fits.push(builder.ins().icmp_imm(cc, raw, dmax as u64 as i64));
4757 }
4758 branch_on(
4759 builder,
4760 fits,
4761 |b| {
4762 store_slot(b, buffer_ptr, out_slot, raw);
4763 },
4764 slow,
4765 );
4766 }
4767 (s, F32 | F64) if s.int_range().is_some() => {
4769 let ty = if dst == F32 { types::F32 } else { types::F64 };
4770 let x = if matches!(s, Signed(_)) {
4771 builder.ins().fcvt_from_sint(ty, raw)
4772 } else {
4773 builder.ins().fcvt_from_uint(ty, raw)
4774 };
4775 store_float(builder, x);
4776 }
4777 (F32, F64) => {
4778 let x = float_of(builder);
4779 store_float(builder, x);
4780 }
4781 (F64, F32) => {
4782 let x = float_of(builder);
4783 let narrow = builder.ins().fdemote(types::F32, x);
4784 store_float(builder, narrow);
4785 }
4786 (F32 | F64, d) => {
4790 let (lo, hi) = match d {
4791 Unsigned(b) => (0.0, 2f64.powi(b as i32)),
4792 Signed(b) => (-(2f64.powi(b as i32 - 1)), 2f64.powi(b as i32 - 1)),
4793 _ => {
4794 slow(builder);
4795 return;
4796 }
4797 };
4798 let x = float_of(builder);
4799 let lo_v = builder.ins().f64const(lo);
4800 let hi_v = builder.ins().f64const(hi);
4801 let above = builder.ins().fcmp(FloatCC::GreaterThanOrEqual, x, lo_v);
4802 let below = builder.ins().fcmp(FloatCC::LessThan, x, hi_v);
4803 branch_on(
4804 builder,
4805 vec![above, below],
4806 |b| {
4807 let word = if matches!(d, Signed(_)) {
4808 b.ins().fcvt_to_sint_sat(types::I64, x)
4809 } else {
4810 b.ins().fcvt_to_uint_sat(types::I64, x)
4811 };
4812 store_slot(b, buffer_ptr, out_slot, word);
4813 },
4814 slow,
4815 );
4816 }
4817 _ => slow(builder),
4818 }
4819}
4820
4821fn branch_on(
4824 builder: &mut FunctionBuilder,
4825 conds: Vec<ir::Value>,
4826 fast: impl FnOnce(&mut FunctionBuilder),
4827 slow: impl FnOnce(&mut FunctionBuilder),
4828) {
4829 let mut conds = conds.into_iter();
4830 let Some(first) = conds.next() else {
4831 fast(builder);
4832 return;
4833 };
4834 let mut ok = first;
4835 for c in conds {
4836 ok = builder.ins().band(ok, c);
4837 }
4838 let fast_block = builder.create_block();
4839 let slow_block = builder.create_block();
4840 let done = builder.create_block();
4841 builder.ins().brif(ok, fast_block, &[], slow_block, &[]);
4842 builder.switch_to_block(fast_block);
4843 builder.seal_block(fast_block);
4844 fast(builder);
4845 builder.ins().jump(done, &[]);
4846 builder.switch_to_block(slow_block);
4847 builder.seal_block(slow_block);
4848 slow(builder);
4849 builder.ins().jump(done, &[]);
4850 builder.switch_to_block(done);
4851 builder.seal_block(done);
4852}
4853
4854fn load_slot(builder: &mut FunctionBuilder, buffer_ptr: ir::Value, slot: usize) -> ir::Value {
4855 let offset = (slot * 8) as i32;
4856 builder
4857 .ins()
4858 .load(types::I64, ir::MemFlags::trusted(), buffer_ptr, offset)
4859}
4860
4861fn store_slot(
4863 builder: &mut FunctionBuilder,
4864 buffer_ptr: ir::Value,
4865 slot: usize,
4866 value: ir::Value,
4867) -> ir::Inst {
4868 let offset = (slot * 8) as i32;
4869 builder
4870 .ins()
4871 .store(ir::MemFlags::trusted(), value, buffer_ptr, offset)
4872}
4873
4874fn reg_lane_type(lane: u8) -> ir::Type {
4877 match lane {
4878 0 => types::I8X16,
4879 1 => types::I16X8,
4880 2 => types::I32X4,
4881 3 => types::I64X2,
4882 4 => types::F32X4,
4883 5 => types::F64X2,
4884 _ => unreachable!("register lane index out of range"),
4885 }
4886}
4887
4888fn load_reg128(
4893 builder: &mut FunctionBuilder,
4894 buffer_ptr: ir::Value,
4895 first_slot: usize,
4896 vt: ir::Type,
4897) -> ir::Value {
4898 let offset = (first_slot * 8) as i32;
4899 builder
4900 .ins()
4901 .load(vt, ir::MemFlags::new(), buffer_ptr, offset)
4902}
4903
4904fn store_reg128(
4906 builder: &mut FunctionBuilder,
4907 buffer_ptr: ir::Value,
4908 first_slot: usize,
4909 value: ir::Value,
4910) {
4911 let offset = (first_slot * 8) as i32;
4912 builder
4913 .ins()
4914 .store(ir::MemFlags::new(), value, buffer_ptr, offset);
4915}
4916
4917fn round_half_away(builder: &mut FunctionBuilder, x: ir::Value) -> ir::Value {
4922 let t = builder.ins().trunc(x);
4923 let frac = builder.ins().fsub(x, t);
4924 let mag = builder.ins().fabs(frac);
4925 let half = builder.ins().f64const(0.5);
4926 let reaches = builder
4927 .ins()
4928 .fcmp(ir::condcodes::FloatCC::GreaterThanOrEqual, mag, half);
4929 let one = builder.ins().f64const(1.0);
4930 let step = builder.ins().fcopysign(one, x);
4931 let up = builder.ins().fadd(t, step);
4932 builder.ins().select(reaches, up, t)
4933}
4934
4935fn clamp_ir(
4940 builder: &mut FunctionBuilder,
4941 x: ir::Value,
4942 lo: ir::Value,
4943 hi: ir::Value,
4944) -> ir::Value {
4945 let below = builder.ins().fcmp(ir::condcodes::FloatCC::LessThan, x, lo);
4946 let above = builder
4947 .ins()
4948 .fcmp(ir::condcodes::FloatCC::GreaterThan, x, hi);
4949 let capped = builder.ins().select(above, hi, x);
4950 builder.ins().select(below, lo, capped)
4951}
4952
4953fn div_ceil(
4956 builder: &mut FunctionBuilder,
4957 val: ir::Value,
4958 m: ir::Value,
4959 one: ir::Value,
4960) -> ir::Value {
4961 let q = builder.ins().udiv(val, m);
4962 let r = builder.ins().urem(val, m);
4963 let zero = builder.ins().iconst(types::I64, 0);
4964 let inexact = builder.ins().icmp(ir::condcodes::IntCC::NotEqual, r, zero);
4965 let q1 = builder.ins().iadd(q, one);
4966 builder.ins().select(inexact, q1, q)
4967}
4968
4969fn func_of<K: PartialEq + Copy>(refs: &[(K, ir::FuncRef)], key: K) -> ir::FuncRef {
4971 refs.iter()
4972 .find(|(k, _)| *k == key)
4973 .map(|(_, r)| *r)
4974 .expect("every helper of the group is declared")
4975}
4976
4977fn load_words(
4980 builder: &mut FunctionBuilder,
4981 buffer_ptr: ir::Value,
4982 input_slots: &[usize],
4983 n: usize,
4984) -> Vec<ir::Value> {
4985 (0..n)
4986 .map(|k| match input_slots.get(k) {
4987 Some(&s) => load_slot(builder, buffer_ptr, s),
4988 None => builder.ins().iconst(types::I64, 0),
4989 })
4990 .collect()
4991}
4992
4993fn load_slot_f64(builder: &mut FunctionBuilder, buffer_ptr: ir::Value, slot: usize) -> ir::Value {
4995 let i64_val = load_slot(builder, buffer_ptr, slot);
4996 builder
4997 .ins()
4998 .bitcast(types::F64, ir::MemFlags::new(), i64_val)
4999}
5000
5001fn store_slot_f64(
5003 builder: &mut FunctionBuilder,
5004 buffer_ptr: ir::Value,
5005 slot: usize,
5006 value: ir::Value,
5007) {
5008 let i64_val = builder
5009 .ins()
5010 .bitcast(types::I64, ir::MemFlags::new(), value);
5011 store_slot(builder, buffer_ptr, slot, i64_val);
5012}
5013
5014#[cfg(test)]
5017mod tests {
5018 use super::*;
5019
5020 #[test]
5025 fn every_scalar_conversion_lowers_by_its_types() {
5026 let mut missing = Vec::new();
5027 for &from in crate::ast::PortType::ALL {
5028 for &to in crate::ast::PortType::ALL {
5029 if from == to || Scalar::of(from).is_none() || Scalar::of(to).is_none() {
5030 continue;
5031 }
5032 let Some(node) = crate::compile::assembly::boundary_adapter(from, to) else {
5033 continue;
5034 };
5035 match classify_node(node.as_ref()) {
5036 JitOp::Convert { from: f, to: t, .. } if f == from && t == to => {}
5037 other => missing.push(format!(
5038 "{} ({from:?} -> {to:?}) classified as {other:?}",
5039 node.meta().name
5040 )),
5041 }
5042 }
5043 }
5044 assert!(missing.is_empty(), "{}", missing.join("\n"));
5045 }
5046
5047 #[test]
5048 fn jit_identity() {
5049 let steps = vec![(JitOp::Identity, vec![0], vec![1])];
5050 let mut output_map = HashMap::new();
5051 output_map.insert("out".into(), 1);
5052 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5053 kernel.eval(&[42]);
5054 assert_eq!(kernel.get("out"), 42);
5055 }
5056
5057 #[test]
5058 fn jit_add_const() {
5059 let steps = vec![(JitOp::AddConst(100), vec![0], vec![1])];
5060 let mut output_map = HashMap::new();
5061 output_map.insert("out".into(), 1);
5062 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5063 kernel.eval(&[5]);
5064 assert_eq!(kernel.get("out"), 105);
5065 }
5066
5067 #[test]
5068 fn jit_mul_const() {
5069 let steps = vec![(JitOp::MulConst(7), vec![0], vec![1])];
5070 let mut output_map = HashMap::new();
5071 output_map.insert("out".into(), 1);
5072 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5073 kernel.eval(&[6]);
5074 assert_eq!(kernel.get("out"), 42);
5075 }
5076
5077 #[test]
5078 fn jit_mod_const() {
5079 let steps = vec![(JitOp::ModConst(100), vec![0], vec![1])];
5080 let mut output_map = HashMap::new();
5081 output_map.insert("out".into(), 1);
5082 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5083 kernel.eval(&[542]);
5084 assert_eq!(kernel.get("out"), 42);
5085 }
5086
5087 #[test]
5088 fn jit_hash() {
5089 let steps = vec![(JitOp::Hash, vec![0], vec![1])];
5090 let mut output_map = HashMap::new();
5091 output_map.insert("out".into(), 1);
5092 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5093
5094 kernel.eval(&[42]);
5095 let v1 = kernel.get("out");
5096
5097 let expected = xxhash_rust::xxh3::xxh3_64(&42u64.to_le_bytes());
5099 assert_eq!(v1, expected);
5100 }
5101
5102 #[test]
5103 fn jit_hash_deterministic() {
5104 let steps = vec![(JitOp::Hash, vec![0], vec![1])];
5105 let mut output_map = HashMap::new();
5106 output_map.insert("out".into(), 1);
5107 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5108
5109 kernel.eval(&[42]);
5110 let v1 = kernel.get("out");
5111 kernel.eval(&[42]);
5112 let v2 = kernel.get("out");
5113 assert_eq!(v1, v2);
5114 }
5115
5116 #[test]
5117 fn jit_chain_hash_mod() {
5118 let steps = vec![
5120 (JitOp::Hash, vec![0], vec![1]), (JitOp::ModConst(1_000_000), vec![1], vec![2]), ];
5123 let mut output_map = HashMap::new();
5124 output_map.insert("user_id".into(), 2);
5125 let mut kernel = compile_jit_raw(1, 3, steps, output_map, Vec::new()).unwrap();
5126
5127 kernel.eval(&[42]);
5128 let uid = kernel.get("user_id");
5129 assert!(uid < 1_000_000, "got {uid}");
5130 }
5131
5132 #[test]
5133 fn jit_clamp_const() {
5134 let steps = vec![(JitOp::ClampConst(10, 50), vec![0], vec![1])];
5135 let mut output_map = HashMap::new();
5136 output_map.insert("out".into(), 1);
5137 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5138
5139 kernel.eval(&[5]);
5140 assert_eq!(kernel.get("out"), 10); kernel.eval(&[30]);
5143 assert_eq!(kernel.get("out"), 30); kernel.eval(&[100]);
5146 assert_eq!(kernel.get("out"), 50); }
5148
5149 #[test]
5150 fn jit_interleave() {
5151 let steps = vec![(JitOp::Interleave, vec![0, 1], vec![2])];
5152 let mut output_map = HashMap::new();
5153 output_map.insert("out".into(), 2);
5154 let mut kernel = compile_jit_raw(2, 3, steps, output_map, Vec::new()).unwrap();
5155
5156 kernel.eval(&[0b101, 0b010]);
5157 assert_eq!(kernel.get("out"), 0b01_10_01);
5159 }
5160
5161 #[test]
5162 fn jit_mixed_radix() {
5163 let steps = vec![(
5165 JitOp::MixedRadixConst(vec![100, 1000, 0]),
5166 vec![0],
5167 vec![1, 2, 3],
5168 )];
5169 let mut output_map = HashMap::new();
5170 output_map.insert("d0".into(), 1);
5171 output_map.insert("d1".into(), 2);
5172 output_map.insert("d2".into(), 3);
5173 let mut kernel = compile_jit_raw(1, 4, steps, output_map, Vec::new()).unwrap();
5174
5175 kernel.eval(&[4_201_337]);
5177 assert_eq!(kernel.get("d0"), 37);
5178 assert_eq!(kernel.get("d1"), 13);
5179 assert_eq!(kernel.get("d2"), 42);
5180 }
5181
5182 #[test]
5183 fn jit_unit_interval() {
5184 let steps = vec![(JitOp::UnitInterval, 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(&[0]);
5190 let v = f64::from_bits(kernel.get("out"));
5191 assert!((v - 0.0).abs() < 1e-10);
5192
5193 kernel.eval(&[u64::MAX]);
5194 let v = f64::from_bits(kernel.get("out"));
5195 assert!((v - 1.0).abs() < 1e-10);
5196 }
5197
5198 #[test]
5199 fn jit_f64_to_u64() {
5200 let steps = vec![(JitOp::F64ToU64, vec![0], vec![1])];
5202 let mut output_map = HashMap::new();
5203 output_map.insert("out".into(), 1);
5204 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5205
5206 kernel.eval(&[3.7f64.to_bits()]);
5207 assert_eq!(kernel.get("out"), 3); }
5209
5210 #[test]
5211 fn jit_round_to_u64() {
5212 let steps = vec![(JitOp::RoundToU64, vec![0], vec![1])];
5213 let mut output_map = HashMap::new();
5214 output_map.insert("out".into(), 1);
5215 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5216
5217 kernel.eval(&[3.7f64.to_bits()]);
5218 assert_eq!(kernel.get("out"), 4);
5219
5220 kernel.eval(&[3.2f64.to_bits()]);
5221 assert_eq!(kernel.get("out"), 3);
5222 }
5223
5224 #[test]
5225 fn jit_clamp_f64() {
5226 let steps = vec![(
5227 JitOp::ClampF64Const(0.0f64.to_bits(), 1.0f64.to_bits()),
5228 vec![0],
5229 vec![1],
5230 )];
5231 let mut output_map = HashMap::new();
5232 output_map.insert("out".into(), 1);
5233 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5234
5235 kernel.eval(&[(-0.5f64).to_bits()]);
5236 assert_eq!(f64::from_bits(kernel.get("out")), 0.0);
5237
5238 kernel.eval(&[0.5f64.to_bits()]);
5239 assert_eq!(f64::from_bits(kernel.get("out")), 0.5);
5240
5241 kernel.eval(&[1.5f64.to_bits()]);
5242 assert_eq!(f64::from_bits(kernel.get("out")), 1.0);
5243 }
5244
5245 #[test]
5246 fn jit_lerp() {
5247 let steps = vec![(
5248 JitOp::LerpConst(10.0f64.to_bits(), 20.0f64.to_bits()),
5249 vec![0],
5250 vec![1],
5251 )];
5252 let mut output_map = HashMap::new();
5253 output_map.insert("out".into(), 1);
5254 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5255
5256 kernel.eval(&[0.0f64.to_bits()]);
5257 assert_eq!(f64::from_bits(kernel.get("out")), 10.0);
5258
5259 kernel.eval(&[1.0f64.to_bits()]);
5260 assert_eq!(f64::from_bits(kernel.get("out")), 20.0);
5261
5262 kernel.eval(&[0.5f64.to_bits()]);
5263 assert_eq!(f64::from_bits(kernel.get("out")), 15.0);
5264 }
5265
5266 #[test]
5267 fn jit_scale_range() {
5268 let steps = vec![(
5269 JitOp::ScaleRangeConst(10.0f64.to_bits(), 10.0f64.to_bits()),
5270 vec![0],
5271 vec![1],
5272 )];
5273 let mut output_map = HashMap::new();
5274 output_map.insert("out".into(), 1);
5275 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5276
5277 kernel.eval(&[0]);
5278 let v = f64::from_bits(kernel.get("out"));
5279 assert!((v - 10.0).abs() < 0.001);
5280
5281 kernel.eval(&[u64::MAX]);
5282 let v = f64::from_bits(kernel.get("out"));
5283 assert!((v - 20.0).abs() < 0.001);
5284 }
5285
5286 #[test]
5287 fn jit_quantize() {
5288 let steps = vec![(JitOp::QuantizeConst(10.0f64.to_bits()), vec![0], vec![1])];
5289 let mut output_map = HashMap::new();
5290 output_map.insert("out".into(), 1);
5291 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5292
5293 kernel.eval(&[13.0f64.to_bits()]);
5294 assert_eq!(f64::from_bits(kernel.get("out")), 10.0);
5295
5296 kernel.eval(&[17.0f64.to_bits()]);
5297 assert_eq!(f64::from_bits(kernel.get("out")), 20.0);
5298 }
5299
5300 #[test]
5301 fn jit_discretize() {
5302 let steps = vec![(
5303 JitOp::DiscretizeConst(100.0f64.to_bits(), 10),
5304 vec![0],
5305 vec![1],
5306 )];
5307 let mut output_map = HashMap::new();
5308 output_map.insert("out".into(), 1);
5309 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5310
5311 kernel.eval(&[0.0f64.to_bits()]);
5312 assert_eq!(kernel.get("out"), 0);
5313
5314 kernel.eval(&[55.0f64.to_bits()]);
5315 assert_eq!(kernel.get("out"), 5);
5316
5317 kernel.eval(&[99.0f64.to_bits()]);
5318 assert_eq!(kernel.get("out"), 9);
5319
5320 kernel.eval(&[200.0f64.to_bits()]);
5322 assert_eq!(kernel.get("out"), 9);
5323 }
5324
5325 #[test]
5326 fn jit_chain_unit_interval_lerp() {
5327 let steps = vec![
5329 (JitOp::UnitInterval, vec![0], vec![1]),
5330 (
5331 JitOp::LerpConst(100.0f64.to_bits(), 200.0f64.to_bits()),
5332 vec![1],
5333 vec![2],
5334 ),
5335 ];
5336 let mut output_map = HashMap::new();
5337 output_map.insert("out".into(), 2);
5338 let mut kernel = compile_jit_raw(1, 3, steps, output_map, Vec::new()).unwrap();
5339
5340 kernel.eval(&[0]);
5341 let v = f64::from_bits(kernel.get("out"));
5342 assert!((v - 100.0).abs() < 0.001);
5343
5344 kernel.eval(&[u64::MAX]);
5345 let v = f64::from_bits(kernel.get("out"));
5346 assert!((v - 200.0).abs() < 0.001);
5347 }
5348
5349 #[test]
5350 fn jit_multi_step_chain() {
5351 let steps = vec![
5353 (JitOp::AddConst(10), vec![0], vec![1]),
5354 (JitOp::MulConst(3), vec![1], vec![2]),
5355 (JitOp::ModConst(100), vec![2], vec![3]),
5356 ];
5357 let mut output_map = HashMap::new();
5358 output_map.insert("out".into(), 3);
5359 let mut kernel = compile_jit_raw(1, 4, steps, output_map, Vec::new()).unwrap();
5360
5361 kernel.eval(&[5]);
5362 assert_eq!(kernel.get("out"), 45);
5364 }
5365
5366 #[test]
5369 fn jit_is_positive_check_passes_positive() {
5370 let steps = vec![(
5371 JitOp::IsPositiveCheck {
5372 name_ptr: 0,
5373 name_len: 0,
5374 },
5375 vec![0],
5376 vec![1],
5377 )];
5378 let mut output_map = HashMap::new();
5379 output_map.insert("out".into(), 1);
5380 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5381 kernel.eval(&[42]);
5382 assert_eq!(kernel.get("out"), 42);
5383 kernel.eval(&[u64::MAX]);
5386 assert_eq!(kernel.get("out"), u64::MAX);
5387 }
5388
5389 #[test]
5390 fn jit_in_range_check_passes_interior() {
5391 let steps = vec![(JitOp::InRangeCheck(10, 100), vec![0], vec![1])];
5392 let mut output_map = HashMap::new();
5393 output_map.insert("out".into(), 1);
5394 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5395 kernel.eval(&[50]);
5396 assert_eq!(kernel.get("out"), 50);
5397 kernel.eval(&[10]);
5399 assert_eq!(kernel.get("out"), 10);
5400 kernel.eval(&[100]);
5401 assert_eq!(kernel.get("out"), 100);
5402 }
5403
5404 #[test]
5409 fn jit_is_one_of_check_passes_allowed_values() {
5410 let steps = vec![(
5411 JitOp::IsOneOfCheck {
5412 allowed: vec![1, 2, 3, 5, 8],
5413 set_ptr: 0,
5414 set_len: 0,
5415 },
5416 vec![0],
5417 vec![1],
5418 )];
5419 let mut output_map = HashMap::new();
5420 output_map.insert("out".into(), 1);
5421 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5422 for v in [1u64, 2, 3, 5, 8] {
5424 kernel.eval(&[v]);
5425 assert_eq!(kernel.get("out"), v);
5426 }
5427 }
5428
5429 #[test]
5430 fn jit_is_one_of_check_accepts_single_element_allow_list() {
5431 let steps = vec![(
5434 JitOp::IsOneOfCheck {
5435 allowed: vec![42],
5436 set_ptr: 0,
5437 set_len: 0,
5438 },
5439 vec![0],
5440 vec![1],
5441 )];
5442 let mut output_map = HashMap::new();
5443 output_map.insert("out".into(), 1);
5444 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5445 kernel.eval(&[42]);
5446 assert_eq!(kernel.get("out"), 42);
5447 }
5448
5449 fn extract_panic_msg(payload: Box<dyn std::any::Any + Send + 'static>) -> String {
5459 payload
5460 .downcast_ref::<String>()
5461 .cloned()
5462 .or_else(|| payload.downcast_ref::<&str>().map(|s| s.to_string()))
5463 .unwrap_or_else(|| "(non-string panic)".into())
5464 }
5465
5466 #[test]
5467 fn jit_is_positive_violation_is_catchable() {
5468 let steps = vec![(
5469 JitOp::IsPositiveCheck {
5470 name_ptr: 0,
5471 name_len: 0,
5472 },
5473 vec![0],
5474 vec![1],
5475 )];
5476 let mut output_map = HashMap::new();
5477 output_map.insert("out".into(), 1);
5478 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5479 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[0])))
5480 .expect_err("JIT violation should panic");
5481 assert!(extract_panic_msg(err).contains("must be > 0"));
5482 }
5483
5484 #[test]
5485 fn jit_in_range_violation_is_catchable() {
5486 let steps = vec![(JitOp::InRangeCheck(10, 100), vec![0], vec![1])];
5487 let mut output_map = HashMap::new();
5488 output_map.insert("out".into(), 1);
5489 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5490
5491 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[5])))
5492 .expect_err("below-range should panic");
5493 assert!(extract_panic_msg(err).contains("outside [10, 100]"));
5494
5495 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[500])))
5496 .expect_err("above-range should panic");
5497 assert!(extract_panic_msg(err).contains("outside [10, 100]"));
5498 }
5499
5500 #[test]
5501 fn jit_is_one_of_violation_is_catchable() {
5502 let steps = vec![(
5503 JitOp::IsOneOfCheck {
5504 allowed: vec![1, 3, 5],
5505 set_ptr: 0,
5506 set_len: 0,
5507 },
5508 vec![0],
5509 vec![1],
5510 )];
5511 let mut output_map = HashMap::new();
5512 output_map.insert("out".into(), 1);
5513 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5514 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[2])))
5515 .expect_err("disallowed value should panic");
5516 assert!(extract_panic_msg(err).contains("not in allowed set"));
5517 }
5518
5519 #[test]
5520 fn invoke_with_catch_restores_slot_after_foreign_panic() {
5521 let caught = std::panic::catch_unwind(|| {
5529 invoke_with_catch(|| panic!("foreign panic"));
5530 });
5531 assert!(caught.is_err(), "foreign panic should propagate out");
5532
5533 let steps = vec![(
5535 JitOp::IsPositiveCheck {
5536 name_ptr: 0,
5537 name_len: 0,
5538 },
5539 vec![0],
5540 vec![1],
5541 )];
5542 let mut output_map = HashMap::new();
5543 output_map.insert("out".into(), 1);
5544 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5545 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[0])))
5546 .expect_err("JIT violation should panic cleanly after foreign panic");
5547 assert!(extract_panic_msg(err).contains("must be > 0"));
5548
5549 kernel.eval(&[42]);
5551 assert_eq!(kernel.get("out"), 42);
5552 }
5553
5554 #[test]
5555 fn jit_kernel_survives_multiple_violations() {
5556 let steps = vec![(
5560 JitOp::IsPositiveCheck {
5561 name_ptr: 0,
5562 name_len: 0,
5563 },
5564 vec![0],
5565 vec![1],
5566 )];
5567 let mut output_map = HashMap::new();
5568 output_map.insert("out".into(), 1);
5569 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5570
5571 for _ in 0..3 {
5572 let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[0])))
5573 .expect_err("violation should still panic");
5574 }
5575 kernel.eval(&[42]);
5577 assert_eq!(kernel.get("out"), 42);
5578 }
5579
5580 const LADDER: &str = "input cycle: u64\ninput tenant_seed: u64\ninput operation_seed: u64\n\
5584 seeded_cycle := u64_add(cycle, tenant_seed)\nidentity_entropy := hash(seeded_cycle)\n\
5585 account_id := mul(identity_entropy, 10000000)\n\
5586 route_seed := u64_add(account_id, operation_seed)\nroute_entropy := hash(route_seed)\n\
5587 shard := mul(route_entropy, 64)\n\
5588 payload_seed := u64_add(identity_entropy, operation_seed)\n\
5589 payload_entropy := hash(payload_seed)\npayload_class := mul(payload_entropy, 8)\n\
5590 token_seed := u64_add(route_entropy, payload_entropy)\nevent_token := hash(token_seed)\n";
5591
5592 fn unit_counts(src: &str) -> (usize, usize) {
5594 let asm = || crate::dsl::compile::compile_polydat_to_assembler(src).unwrap();
5595 let pure = asm()
5596 .try_compile_pure_jit_raw()
5597 .unwrap()
5598 .core
5599 .cones
5600 .unit_count();
5601 let native = crate::dsl::compile::compile_polydat_with(
5602 src,
5603 crate::compile::select::Engine::Native(crate::compile::select::Provenance::Raw),
5604 )
5605 .unwrap()
5606 .plan()
5607 .native_segments;
5608 (pure, native)
5609 }
5610
5611 #[test]
5617 fn only_an_extern_splits_a_unit() {
5618 assert_eq!(unit_counts(LADDER), (1, 1));
5619 let extern_seed = LADDER.replace(
5620 "input operation_seed: u64",
5621 "extern operation_seed: u64 = 7",
5622 );
5623 assert_eq!(unit_counts(&extern_seed), (2, 2));
5624 let shape = "input x: u64\nMODE\nt := hash(x)\na := hash(t)\n\
5627 b := u64_add(t, mode)\nc := add(b, 1)\n";
5628 assert_eq!(
5629 unit_counts(&shape.replace("MODE", "input mode: u64")),
5630 (1, 1)
5631 );
5632 assert_eq!(
5633 unit_counts(&shape.replace("MODE", "extern mode: u64 = 3")),
5634 (2, 2)
5635 );
5636 }
5637}