polydat_core/compile/jit/
mod.rs

1// Copyright 2024-2026 Jonathan Shook
2// SPDX-License-Identifier: Apache-2.0
3
4//! Phase 3: Cranelift JIT compilation of Polydat Kernels.
5//!
6//! Generates native machine code from the DAG. The entire kernel
7//! becomes a single function over the state's slot buffer and
8//! scratch: `fn(coords: *const u64, buffer: *mut u64, scratch: *mut ScratchBuf)`
9//! (plus a clean-flag pointer on the provenance variant).
10//! Arithmetic is inlined over the buffer; a node with no named
11//! lowering runs its own slot kit through one helper call
12//! (`JitOp::SlotCall`), with the inputs gathered into the native
13//! frame and the outputs scattered back.
14//!
15//! The buffer is `Vec<u64>`. For f64 values, they are stored as their
16//! bit representation (`f64::to_bits()` / `f64::from_bits()`). The JIT
17//! uses Cranelift `bitcast` (free, no instruction emitted) to convert
18//! between i64 and f64 representations when crossing type boundaries.
19//!
20//! Feature-gated behind `jit`.
21//!
22//! Simple ops are fully inlined (hash is an inline splitmix64); ops
23//! with a body Cranelift cannot express (xxhash3, shuffle, interleave,
24//! the math functions) call an extern helper, and any other node with
25//! a kit calls the kit in place.
26
27#[cfg(feature = "jit")]
28mod codegen;
29#[cfg(feature = "jit")]
30pub mod host_isa;
31#[cfg(feature = "jit")]
32mod kernels;
33#[cfg(feature = "jit")]
34pub mod simd;
35
36#[cfg(feature = "jit")]
37pub use codegen::*;
38#[cfg(feature = "jit")]
39pub use kernels::*;
40
41#[cfg(all(test, feature = "jit"))]
42mod tests {
43    use super::*;
44
45    #[test]
46    fn test_inventory_tier_distribution() {
47        use crate::ast::PortType;
48        use crate::compile::assembly::WireRef;
49        use crate::dsl::factory::{ConstArg, build_node};
50        use crate::dsl::registry::registry;
51
52        let reg = registry();
53        let total = reg.len();
54
55        let mut p1_count = 0;
56        let mut p2_count = 0;
57        let mut p3_count = 0;
58        let mut unbuilt = 0;
59
60        for sig in &reg {
61            let mut consts = Vec::new();
62            for p in sig.params {
63                match p.slot_type {
64                    crate::ast::SlotType::ConstU64 => consts.push(ConstArg::Int(1)),
65                    crate::ast::SlotType::ConstF64 => consts.push(ConstArg::Float(1.0)),
66                    crate::ast::SlotType::ConstStr => consts.push(ConstArg::Str("test".into())),
67                    crate::ast::SlotType::ConstVecU64 => consts.push(ConstArg::Int(1)),
68                    crate::ast::SlotType::ConstVecF64 => consts.push(ConstArg::Float(1.0)),
69                    crate::ast::SlotType::ConstVec => consts.push(ConstArg::Int(1)),
70                    crate::ast::SlotType::Wire => {}
71                }
72            }
73            let wires = vec![WireRef::Input("cycle".to_string()); sig.wire_input_count().max(1)];
74            let wire_types = vec![PortType::U64; wires.len()];
75
76            let node_res = build_node(
77                &crate::dsl::factory::BuildContext::default(),
78                sig.name,
79                &wires,
80                &wire_types,
81                &consts,
82            );
83            if let Ok(node) = node_res {
84                // Every compiled form, not only the scalar one: a node
85                // whose form is a slot kit belongs in the P2 column.
86                let p2_eligible = crate::compile::node_tier(node.as_ref(), &wire_types)
87                    != crate::ast::CompileLevel::Phase1;
88                let p3_eligible = classify_node(node.as_ref()) != JitOp::Fallback;
89                if p3_eligible {
90                    p3_count += 1;
91                } else if p2_eligible {
92                    p2_count += 1;
93                } else {
94                    p1_count += 1;
95                }
96            } else {
97                unbuilt += 1;
98                p1_count += 1;
99            }
100        }
101
102        println!("\n=== COMPILER OPTIMIZATION INVENTORY SUMMARY ===");
103        println!("Total Registered Functions: {total}");
104        println!(
105            "Phase 3 (Full Native JIT):  {p3_count} ({:.1}%)",
106            (p3_count as f64 / total as f64) * 100.0
107        );
108        println!(
109            "Phase 2 (Captured Closure): {p2_count} ({:.1}%)",
110            (p2_count as f64 / total as f64) * 100.0
111        );
112        println!(
113            "Phase 1 (Interpreter Cones):{p1_count} ({:.1}%) (unbuilt fallback: {unbuilt})",
114            (p1_count as f64 / total as f64) * 100.0
115        );
116        println!("===============================================\n");
117
118        // The tier counts are a claim about the library, so one of
119        // them is asserted rather than printed. Asking only
120        // `compiled_u64` put 251 of 457 nodes in the P1 column and 2
121        // in P2, which said polydat interprets more than half of its
122        // own node library; it compiles all but a small tail. A
123        // proportion rather than a number, so adding nodes does not
124        // fail it, and no list of names, so a node that loses its
125        // compiled form is not excused by being on one.
126        let p1_share = p1_count as f64 / total as f64;
127        assert!(
128            p1_share < 0.15,
129            "{p1_count} of {total} registered nodes reach no compiled form \
130             ({:.1}%); the tier predicate has narrowed",
131            p1_share * 100.0
132        );
133    }
134}
polydat_core/compile/jit/mod.rs

polydat_core/compile/jit/
mod.rs
