1use std::collections::HashMap;
25
26use crate::ast::SlotShape;
27use crate::ast::{CompiledU64Op, PolydatNode};
28use crate::kernel::WireSource;
29
30#[cfg(feature = "jit")]
31use crate::compile::jit::{self, JitOp};
32
33enum HybridStep {
35 #[cfg(feature = "jit")]
38 Jit(JitSegment),
39 Closure(ClosureStep),
41}
42
43#[cfg(feature = "jit")]
44struct JitSegment {
45 code_fn: crate::compile::jit::NativeFn,
46 _module: crate::compile::jit::JitCode,
49 fallible: bool,
52 input_slots: Vec<usize>,
55 output_slots: Vec<usize>,
56 nodes: Vec<usize>,
59}
60
61#[cfg(feature = "jit")]
64struct PendingSegment {
65 step: usize,
66 batch: Vec<(JitOp, Vec<usize>, Vec<usize>)>,
67 input_slots: Vec<usize>,
68 output_slots: Vec<usize>,
69 nodes: Vec<usize>,
70}
71
72impl HybridStep {
73 fn input_slots(&self) -> &[usize] {
74 match self {
75 #[cfg(feature = "jit")]
76 HybridStep::Jit(seg) => &seg.input_slots,
77 HybridStep::Closure(cs) => &cs.input_slots,
78 }
79 }
80 fn output_slots(&self) -> &[usize] {
81 match self {
82 #[cfg(feature = "jit")]
83 HybridStep::Jit(seg) => &seg.output_slots,
84 HybridStep::Closure(cs) => &cs.output_slots,
85 }
86 }
87 fn accepts_none(&self) -> bool {
90 match self {
91 #[cfg(feature = "jit")]
92 HybridStep::Jit(_) => false,
93 HybridStep::Closure(cs) => cs.accepts_none,
94 }
95 }
96
97 #[cfg_attr(not(feature = "jit"), allow(unused_variables))]
100 fn failing_node(&self, buffer: &[u64], tracker: usize) -> usize {
101 match self {
102 #[cfg(feature = "jit")]
103 HybridStep::Jit(seg) => seg
104 .nodes
105 .get(buffer[tracker] as usize)
106 .copied()
107 .unwrap_or(usize::MAX),
108 HybridStep::Closure(cs) => cs.node,
109 }
110 }
111}
112
113enum ClosureOp {
117 U64(CompiledU64Op),
118 Slot(crate::ast::CompiledSlotOp),
119}
120
121struct ClosureStep {
122 op: ClosureOp,
123 input_slots: Vec<usize>,
124 output_slots: Vec<usize>,
125 scratch_range: (usize, usize),
127 accepts_none: bool,
129 node: usize,
131}
132
133type ResolvedOutput = (
136 usize,
137 crate::ast::PortType,
138 Option<std::sync::Arc<[usize]>>,
139 bool,
140);
141
142struct HybridCore {
148 engine: crate::compile::select::Engine,
153 buffer: Vec<u64>,
154 coord_count: usize,
155 steps: std::sync::Arc<Vec<HybridStep>>,
156 output_map: HashMap<String, usize>,
157 gather_buf: Vec<u64>,
158 scatter_buf: Vec<u64>,
159 scratch: Vec<crate::ast::ScratchBuf>,
163 ref_slots: Vec<bool>,
165 ref_scratch: Vec<(usize, usize)>,
167 output_types: HashMap<String, crate::ast::PortType>,
169 externs: crate::compile::externs::Externs,
171 traversals: std::sync::Arc<[crate::dsl::traversal::Traversal]>,
174 resolved_outputs: Vec<Option<ResolvedOutput>>,
177 _nodes: std::sync::Arc<Vec<Box<dyn PolydatNode>>>,
179 drive: crate::compile::Drive,
183 none: Vec<bool>,
185 ran: Vec<u64>,
188 epoch: u64,
194 all_ran: bool,
196 clean: Vec<bool>,
200 use_clean: bool,
202 plan: std::sync::Arc<crate::compile::Invalidation>,
205 volatile: std::sync::Arc<[bool]>,
207 side: std::sync::Arc<[bool]>,
209 slot_step: std::sync::Arc<[Option<usize>]>,
211 sites: std::sync::Arc<crate::compile::Attribution>,
213 cur_step: usize,
215 tracker: usize,
218 all: std::sync::Arc<[usize]>,
220 dirty: std::sync::Arc<[Vec<usize>]>,
226 any_none: bool,
230 volatile_steps: std::sync::Arc<[usize]>,
232}
233
234impl Clone for HybridCore {
235 fn clone(&self) -> Self {
236 let mut core = HybridCore {
237 engine: self.engine,
238 buffer: self.buffer.clone(),
239 coord_count: self.coord_count,
240 steps: self.steps.clone(),
241 output_map: self.output_map.clone(),
242 gather_buf: self.gather_buf.clone(),
243 scatter_buf: self.scatter_buf.clone(),
244 scratch: self.scratch.clone(),
245 ref_slots: self.ref_slots.clone(),
246 ref_scratch: self.ref_scratch.clone(),
247 output_types: self.output_types.clone(),
248 externs: self.externs.clone(),
249 traversals: self.traversals.clone(),
250 resolved_outputs: self.resolved_outputs.clone(),
251 _nodes: self._nodes.clone(),
252 drive: self.drive.clone(),
253 none: self.none.clone(),
254 ran: self.ran.clone(),
255 epoch: self.epoch,
256 all_ran: self.all_ran,
257 clean: self.clean.clone(),
258 use_clean: self.use_clean,
259 plan: self.plan.clone(),
260 volatile: self.volatile.clone(),
261 side: self.side.clone(),
262 slot_step: self.slot_step.clone(),
263 sites: self.sites.clone(),
264 cur_step: self.cur_step,
265 tracker: self.tracker,
266 all: self.all.clone(),
267 dirty: self.dirty.clone(),
268 any_none: self.any_none,
269 volatile_steps: self.volatile_steps.clone(),
270 };
271 core.republish_refs();
272 core
273 }
274}
275
276impl HybridCore {
277 crate::compile::shared_core_methods!();
278}
279
280impl HybridCore {
281 fn set_use_clean(&mut self, on: bool) {
285 self.use_clean = on;
286 let side = std::sync::Arc::clone(&self.side);
287 self.dirty = self
288 .plan
289 .input_dependents
290 .iter()
291 .map(|deps| {
292 if on {
293 deps.clone()
294 } else {
295 deps.iter().copied().filter(|&i| side[i]).collect()
296 }
297 })
298 .collect::<Vec<_>>()
299 .into();
300 }
301
302 fn program_identity(&self) -> usize {
308 std::sync::Arc::as_ptr(&self.steps) as *const () as usize
309 }
310
311 #[inline]
312 fn step_can_fail(&self, i: usize) -> bool {
313 match &self.steps[i] {
314 #[cfg(feature = "jit")]
315 HybridStep::Jit(seg) => seg.fallible,
316 _ => true,
317 }
318 }
319
320 #[inline]
326 fn failing_node(&self) -> usize {
327 self.steps[self.cur_step].failing_node(&self.buffer, self.tracker)
328 }
329
330 #[inline]
332 fn run_order(&mut self, order: &[usize]) {
333 let steps = &self.steps;
334 let none_free = !self.any_none;
335 for &i in order {
336 if self.all_ran || self.ran[i] == self.epoch {
337 continue;
338 }
339 let never = self.volatile[i];
340 if (self.use_clean || self.side[i]) && self.clean[i] && !never {
341 self.ran[i] = self.epoch;
342 continue;
343 }
344 self.cur_step = i;
345 run_hybrid_step(
346 &steps[i],
347 none_free,
348 &mut self.buffer,
349 &mut self.none,
350 &mut self.gather_buf,
351 &mut self.scatter_buf,
352 &mut self.scratch,
353 );
354 self.ran[i] = self.epoch;
355 self.clean[i] = !never;
356 }
357 }
358
359 #[inline]
365 fn run_fresh(&mut self) {
366 let steps = &self.steps;
367 for (i, step) in steps.iter().enumerate() {
368 if self.side[i] {
369 let never = self.volatile[i];
370 if self.clean[i] && !never {
371 continue;
372 }
373 self.clean[i] = !never;
374 }
375 self.cur_step = i;
376 run_hybrid_step(
377 step,
378 true,
379 &mut self.buffer,
380 &mut self.none,
381 &mut self.gather_buf,
382 &mut self.scatter_buf,
383 &mut self.scratch,
384 );
385 }
386 self.all_ran = true;
387 }
388
389 fn plan(&self) -> crate::EnginePlan {
391 let (native_segments, closure_steps) = self.engine_counts();
392 crate::EnginePlan {
393 native_segments,
394 closure_steps,
395 interpreted_nodes: 0,
396 }
397 }
398}
399
400#[inline]
403fn eval_all_hybrid_steps(core: &mut HybridCore) {
404 core.drive.stale = true;
405 core.eval_all();
406}
407
408impl HybridCore {
413 fn engine_counts(&self) -> (usize, usize) {
416 let closures = self
417 .steps
418 .iter()
419 .filter(|s| matches!(s, HybridStep::Closure(_)))
420 .count();
421 (self.steps.len() - closures, closures)
422 }
423}
424
425#[derive(Clone)]
430pub struct HybridKernelRaw {
431 core: HybridCore,
432}
433
434impl HybridKernelRaw {
435 crate::compile::kernel_accessors!(set_coords);
436 #[inline]
439 fn set_coords(&mut self, coords: &[u64]) {
440 for (i, &c) in coords
441 .iter()
442 .enumerate()
443 .take(self.core.externs.coordinate_slots())
444 {
445 if self.core.buffer[i] != c {
446 self.core.buffer[i] = c;
447 self.core.dirty_input(i);
448 }
449 }
450 }
451
452 #[inline]
454 pub fn eval(&mut self, coords: &[u64]) {
455 self.set_coords(coords);
456 eval_all_hybrid_steps(&mut self.core);
457 }
458
459 pub fn set_input(
463 &mut self,
464 name: &str,
465 value: crate::ast::Value,
466 ) -> Result<(), crate::kernel::WriteError> {
467 self.core.set_extern(name, value).map(|_| ())
468 }
469
470 pub fn set_input_at(
472 &mut self,
473 index: usize,
474 value: crate::ast::Value,
475 ) -> Result<(), crate::kernel::WriteError> {
476 self.core.set_extern_at(index, value).map(|_| ())
477 }
478
479 #[inline]
481 pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
482 self.core.guard_ref_slot(slot);
483 self.eval(coords);
484 self.core.buffer[slot]
485 }
486
487 pub fn engine_counts(&self) -> (usize, usize) {
490 self.core.engine_counts()
491 }
492
493 pub fn retain_nodes(&mut self, nodes: Vec<Box<dyn PolydatNode>>) {
495 self.core._nodes = std::sync::Arc::new(nodes);
496 }
497}
498
499#[derive(Clone)]
511pub struct HybridKernelPull {
512 core: HybridCore,
513 slot_provenance: Vec<crate::kernel::ProvMask>,
514 changed_mask: crate::kernel::ProvMask,
515 force_run: bool,
518}
519
520impl HybridKernelPull {
521 crate::compile::kernel_accessors!(set_inputs);
522 #[inline]
525 fn set_inputs(&mut self, coords: &[u64]) {
526 self.changed_mask.clear();
527 for (i, &c) in coords
528 .iter()
529 .enumerate()
530 .take(self.core.externs.coordinate_slots())
531 {
532 if self.core.buffer[i] != c {
533 self.core.buffer[i] = c;
534 self.changed_mask.set(i);
535 self.core.dirty_input(i);
536 }
537 }
538 }
539
540 #[inline]
542 pub fn eval(&mut self, coords: &[u64]) {
543 self.set_inputs(coords);
544 self.force_run = false;
545 eval_all_hybrid_steps(&mut self.core);
546 }
547
548 #[inline]
551 pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
552 self.core.guard_ref_slot(slot);
553 self.set_inputs(coords);
554 if !self.force_run
555 && slot < self.slot_provenance.len()
556 && !self.slot_provenance[slot].intersects(&self.changed_mask)
557 {
558 return self.core.buffer[slot];
559 }
560 self.force_run = false;
561 eval_all_hybrid_steps(&mut self.core);
562 self.core.buffer[slot]
563 }
564
565 pub fn set_input(
569 &mut self,
570 name: &str,
571 value: crate::ast::Value,
572 ) -> Result<(), crate::kernel::WriteError> {
573 self.core.set_extern(name, value)?;
574 self.force_run = true;
575 Ok(())
576 }
577
578 pub fn set_input_at(
580 &mut self,
581 index: usize,
582 value: crate::ast::Value,
583 ) -> Result<(), crate::kernel::WriteError> {
584 self.core.set_extern_at(index, value)?;
585 self.force_run = true;
586 Ok(())
587 }
588
589 pub fn engine_counts(&self) -> (usize, usize) {
592 self.core.engine_counts()
593 }
594
595 pub fn retain_nodes(&mut self, nodes: Vec<Box<dyn PolydatNode>>) {
597 self.core._nodes = std::sync::Arc::new(nodes);
598 }
599}
600
601#[derive(Clone)]
615pub struct HybridKernelPushPull {
616 core: HybridCore,
617 slot_provenance: Vec<crate::kernel::ProvMask>,
618 changed_mask: crate::kernel::ProvMask,
619 force_run: bool,
622}
623
624impl HybridKernelPushPull {
625 crate::compile::kernel_accessors!(set_inputs);
626 pub fn set_input(
631 &mut self,
632 name: &str,
633 value: crate::ast::Value,
634 ) -> Result<(), crate::kernel::WriteError> {
635 self.core.set_extern(name, value)?;
636 self.force_run = true;
637 Ok(())
638 }
639
640 pub fn set_input_at(
642 &mut self,
643 index: usize,
644 value: crate::ast::Value,
645 ) -> Result<(), crate::kernel::WriteError> {
646 self.core.set_extern_at(index, value)?;
647 self.force_run = true;
648 Ok(())
649 }
650
651 #[inline]
653 fn set_inputs(&mut self, coords: &[u64]) {
654 self.changed_mask.clear();
655 for (i, &c) in coords
656 .iter()
657 .enumerate()
658 .take(self.core.externs.coordinate_slots())
659 {
660 if self.core.buffer[i] != c {
661 self.core.buffer[i] = c;
662 self.changed_mask.set(i);
663 self.core.dirty_input(i);
664 }
665 }
666 }
667
668 #[inline]
670 pub fn eval(&mut self, coords: &[u64]) {
671 self.set_inputs(coords);
672 self.force_run = false;
673 self.core.drive.stale = true;
674 self.core.eval_all();
675 }
676
677 #[inline]
679 pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
680 self.core.guard_ref_slot(slot);
681 self.set_inputs(coords);
682 if !self.force_run
683 && slot < self.slot_provenance.len()
684 && !self.slot_provenance[slot].intersects(&self.changed_mask)
685 {
686 return self.core.buffer[slot];
687 }
688 self.force_run = false;
689 self.core.drive.stale = true;
690 self.core.eval_all();
691 self.core.buffer[slot]
692 }
693
694 pub fn engine_counts(&self) -> (usize, usize) {
697 self.core.engine_counts()
698 }
699
700 pub fn retain_nodes(&mut self, nodes: Vec<Box<dyn PolydatNode>>) {
702 self.core._nodes = std::sync::Arc::new(nodes);
703 }
704}
705
706pub type HybridKernel = HybridKernelPushPull;
711
712fn flatten_input_slots(
716 wiring: &[Vec<WireSource>],
717 nodes: &[Box<dyn PolydatNode>],
718 node_idx: usize,
719 port_offsets: &[Vec<usize>],
720 input_starts: &[usize],
721 input_widths: &[usize],
722) -> Vec<usize> {
723 let mut slots = Vec::new();
724 for source in &wiring[node_idx] {
725 let (start, w) = match source {
726 WireSource::Input(c) => (
727 input_starts.get(*c).copied().unwrap_or(*c),
728 input_widths.get(*c).copied().unwrap_or(1),
729 ),
730 WireSource::NodeOutput(u, p) => (
731 port_offsets[*u][*p],
732 nodes[*u].meta().outs[*p].typ.slot_width(),
733 ),
734 };
735 slots.extend(start..start + w);
736 }
737 slots
738}
739
740fn flatten_ref_output_starts(
743 nodes: &[Box<dyn PolydatNode>],
744 node_idx: usize,
745 port_offsets: &[Vec<usize>],
746) -> Vec<usize> {
747 nodes[node_idx]
748 .meta()
749 .outs
750 .iter()
751 .enumerate()
752 .filter(|(_, out)| out.typ.slot_color() == crate::ast::SlotColor::Ref2)
753 .map(|(p, _)| port_offsets[node_idx][p])
754 .collect()
755}
756
757fn flatten_output_slots(
759 nodes: &[Box<dyn PolydatNode>],
760 node_idx: usize,
761 port_offsets: &[Vec<usize>],
762) -> Vec<usize> {
763 let mut slots = Vec::new();
764 for (p, out) in nodes[node_idx].meta().outs.iter().enumerate() {
765 let start = port_offsets[node_idx][p];
766 slots.extend(start..start + out.typ.slot_width());
767 }
768 slots
769}
770
771fn refused(reason: String) -> crate::KernelError {
782 crate::KernelError::Refused {
783 engine: crate::compile::select::Engine::Native(crate::compile::select::Provenance::Auto),
784 reason,
785 }
786}
787
788#[cfg(feature = "jit")]
790#[allow(clippy::too_many_arguments)]
791pub(crate) fn build_hybrid(
792 nodes: &[Box<dyn PolydatNode>],
793 wiring: &[Vec<WireSource>],
794 coord_count: usize,
795 total_slots: usize,
796 port_offsets: &[Vec<usize>],
797 input_starts: &[usize],
798 input_widths: &[usize],
799 output_map: HashMap<String, usize>,
800 ref_slots: Vec<bool>,
801 input_types: &[crate::ast::PortType],
802 externs: crate::compile::externs::Externs,
803 constant: Vec<bool>,
804 volatile: Vec<bool>,
805 attribution: std::sync::Arc<crate::compile::Attribution>,
806) -> Result<HybridKernelPushPull, crate::KernelError> {
807 let mut steps: Vec<Option<HybridStep>> = Vec::new();
810 let mut pending: Vec<PendingSegment> = Vec::new();
811 let mut scratch: Vec<crate::ast::ScratchBuf> = Vec::new();
812 let mut ref_scratch: Vec<(usize, usize)> = Vec::new();
813 let mut max_inputs = 0usize;
814 let mut max_outputs = 0usize;
815 let graph = GraphView {
816 nodes,
817 wiring,
818 port_offsets,
819 input_types,
820 };
821
822 let classifications: Vec<(JitOp, Vec<usize>, Vec<usize>)> = nodes
824 .iter()
825 .enumerate()
826 .map(|(node_idx, node)| {
827 let wire_types: Vec<crate::ast::PortType> = wiring[node_idx]
831 .iter()
832 .map(|src| match src {
833 WireSource::Input(c) => input_types
834 .get(*c)
835 .copied()
836 .unwrap_or(crate::ast::PortType::U64),
837 WireSource::NodeOutput(j, p) => nodes[*j].meta().outs[*p].typ,
838 })
839 .collect();
840 let jit_op = jit::classify_node_typed(node.as_ref(), &wire_types);
841
842 let input_slots = flatten_input_slots(
843 wiring,
844 nodes,
845 node_idx,
846 port_offsets,
847 input_starts,
848 input_widths,
849 );
850 let output_slots = flatten_output_slots(nodes, node_idx, port_offsets);
851
852 max_inputs = max_inputs.max(input_slots.len());
853 max_outputs = max_outputs.max(output_slots.len());
854
855 (jit_op, input_slots, output_slots)
856 })
857 .collect();
858 let mut classifications = classifications;
868 let mut eligible = vec![false; nodes.len()];
869 for (node_idx, node) in nodes.iter().enumerate() {
870 if matches!(classifications[node_idx].0, JitOp::Fallback) {
871 continue;
872 }
873 if !crate::compile::none_rule_admits(
874 node.accepts_none_inputs(),
875 &wiring[node_idx],
876 &eligible,
877 ) {
878 classifications[node_idx].0 = JitOp::Fallback;
879 continue;
880 }
881 eligible[node_idx] = true;
882 }
883 let unset = externs.unset_slots();
889 if !unset.is_empty() {
890 let mut tainted = vec![false; nodes.len()];
891 for node_idx in 0..nodes.len() {
892 tainted[node_idx] = wiring[node_idx].iter().any(|src| match src {
893 WireSource::Input(c) => unset.contains(&input_starts[*c]),
894 WireSource::NodeOutput(j, _) => tainted[*j],
895 });
896 if tainted[node_idx] {
897 classifications[node_idx].0 = JitOp::Fallback;
898 }
899 }
900 }
901
902 let mut node_step = vec![usize::MAX; nodes.len()];
904 let order: Vec<usize> = (0..nodes.len())
912 .filter(|&k| constant[k])
913 .chain((0..nodes.len()).filter(|&k| !constant[k]))
914 .collect();
915 let mut rank = vec![0usize; nodes.len()];
916 for (pos, &k) in order.iter().enumerate() {
917 rank[k] = pos;
918 }
919 let is_side = |k: usize| matches!(nodes[k].purity(), crate::ast::Purity::SideChannel { .. });
928 let preds: Vec<Vec<usize>> = wiring
929 .iter()
930 .map(|w| {
931 w.iter()
932 .filter_map(|src| match src {
933 WireSource::NodeOutput(j, _) => Some(*j),
934 WireSource::Input(_) => None,
935 })
936 .collect()
937 })
938 .collect();
939 let inputs_read: Vec<Vec<usize>> = wiring
940 .iter()
941 .map(|w| {
942 w.iter()
943 .filter_map(|src| match src {
944 WireSource::Input(c) => Some(*c),
945 WireSource::NodeOutput(..) => None,
946 })
947 .collect()
948 })
949 .collect();
950 let fusible: Vec<bool> = (0..nodes.len())
951 .map(|k| !matches!(classifications[k].0, JitOp::Fallback) && !is_side(k))
952 .collect();
953 let class: Vec<u64> = (0..nodes.len())
954 .map(|k| constant[k] as u64 | (volatile[k] as u64) << 1)
955 .collect();
956 let plan = crate::compile::fusion_units::plan_units(
957 &preds,
958 &inputs_read,
959 &fusible,
960 &class,
961 &rank,
962 &|c| c & 1 == 1,
963 );
964 for members in plan.units {
965 let i = members[0];
966 if matches!(classifications[i].0, JitOp::Fallback) {
967 let (_, ref input_slots, ref output_slots) = classifications[i];
971 let step = closure_step_for(
972 &graph,
973 i,
974 input_slots.clone(),
975 output_slots.clone(),
976 &mut scratch,
977 &mut ref_scratch,
978 )?;
979 node_step[i] = steps.len();
980 steps.push(Some(HybridStep::Closure(step)));
981 } else {
982 for &k in &members {
986 let base = scratch.len();
987 classifications[k].0.place_scratch(base);
988 let elems = classifications[k].0.scratch_elems().to_vec();
989 ref_scratch.extend(crate::compile::assembly::scratch_pairs(
990 &nodes[k].meta().name,
991 &flatten_ref_output_starts(nodes, k, port_offsets),
992 &elems,
993 base,
994 ));
995 scratch.extend(elems.iter().map(|e| crate::ast::ScratchBuf::new(*e)));
996 }
997 let batch: Vec<(JitOp, Vec<usize>, Vec<usize>)> = members
1005 .iter()
1006 .map(|&k| classifications[k].clone())
1007 .collect();
1008 let written: std::collections::HashSet<usize> = batch
1009 .iter()
1010 .flat_map(|(_, _, o)| o.iter().copied())
1011 .collect();
1012 let mut input_slots: Vec<usize> = Vec::new();
1013 for (_, ins, _) in &batch {
1014 for &s in ins {
1015 if !written.contains(&s) && !input_slots.contains(&s) {
1016 input_slots.push(s);
1017 }
1018 }
1019 }
1020 let output_slots: Vec<usize> = batch
1021 .iter()
1022 .flat_map(|(_, _, o)| o.iter().copied())
1023 .collect();
1024 let segment = steps.len();
1025 for &k in &members {
1026 node_step[k] = segment;
1027 }
1028 steps.push(None);
1029 pending.push(PendingSegment {
1030 step: segment,
1031 batch,
1032 input_slots,
1033 output_slots,
1034 nodes: members,
1035 });
1036 }
1037 }
1038 let batches: Vec<&[jit::JitStep]> = pending.iter().map(|p| p.batch.as_slice()).collect();
1042 let (entries, code) = if batches.is_empty() {
1043 (Vec::new(), None)
1044 } else {
1045 let (entries, code) =
1046 jit::compile_jit_entries(&batches, Some(total_slots)).map_err(refused)?;
1047 (entries, Some(code))
1048 };
1049 for (p, (code_fn, fallible)) in pending.into_iter().zip(entries) {
1050 steps[p.step] = Some(HybridStep::Jit(JitSegment {
1051 code_fn,
1052 fallible,
1053 _module: code.clone().expect("a segment was compiled"),
1054 input_slots: p.input_slots,
1055 output_slots: p.output_slots,
1056 nodes: p.nodes,
1057 }));
1058 }
1059 let steps: Vec<HybridStep> = steps
1060 .into_iter()
1061 .map(|s| s.expect("every step is placed"))
1062 .collect();
1063
1064 let output_types = output_types_of(nodes, port_offsets, input_starts, input_types, &output_map);
1065 build_pushpull_from_steps(
1066 steps,
1067 scratch,
1068 ref_scratch,
1069 ref_slots,
1070 wiring,
1071 nodes,
1072 coord_count,
1073 total_slots,
1074 output_map,
1075 max_inputs,
1076 max_outputs,
1077 input_starts,
1078 input_widths,
1079 output_types,
1080 externs,
1081 constant,
1082 volatile,
1083 attribution,
1084 node_step,
1085 )
1086}
1087
1088fn output_types_of(
1091 nodes: &[Box<dyn PolydatNode>],
1092 port_offsets: &[Vec<usize>],
1093 input_starts: &[usize],
1094 input_types: &[crate::ast::PortType],
1095 output_map: &HashMap<String, usize>,
1096) -> HashMap<String, crate::ast::PortType> {
1097 let mut slot_types: HashMap<usize, crate::ast::PortType> = HashMap::new();
1098 for (start, ty) in input_starts.iter().zip(input_types) {
1099 slot_types.insert(*start, *ty);
1100 }
1101 for (node_idx, node) in nodes.iter().enumerate() {
1102 for (p, out) in node.meta().outs.iter().enumerate() {
1103 slot_types.insert(port_offsets[node_idx][p], out.typ);
1104 }
1105 }
1106 output_map
1107 .iter()
1108 .map(|(name, slot)| {
1109 (
1110 name.clone(),
1111 slot_types
1112 .get(slot)
1113 .copied()
1114 .unwrap_or(crate::ast::PortType::U64),
1115 )
1116 })
1117 .collect()
1118}
1119
1120#[cfg(not(feature = "jit"))]
1122#[allow(clippy::too_many_arguments)]
1123pub(crate) fn build_hybrid(
1124 nodes: &[Box<dyn PolydatNode>],
1125 wiring: &[Vec<WireSource>],
1126 coord_count: usize,
1127 total_slots: usize,
1128 port_offsets: &[Vec<usize>],
1129 input_starts: &[usize],
1130 input_widths: &[usize],
1131 output_map: HashMap<String, usize>,
1132 ref_slots: Vec<bool>,
1133 input_types: &[crate::ast::PortType],
1134 externs: crate::compile::externs::Externs,
1135 constant: Vec<bool>,
1136 volatile: Vec<bool>,
1137 attribution: std::sync::Arc<crate::compile::Attribution>,
1138) -> Result<HybridKernelPushPull, crate::KernelError> {
1139 let mut steps: Vec<HybridStep> = Vec::new();
1140 let mut scratch: Vec<crate::ast::ScratchBuf> = Vec::new();
1141 let mut ref_scratch: Vec<(usize, usize)> = Vec::new();
1142 let mut max_inputs = 0usize;
1143 let mut max_outputs = 0usize;
1144 let graph = GraphView {
1145 nodes,
1146 wiring,
1147 port_offsets,
1148 input_types,
1149 };
1150
1151 for node_idx in 0..nodes.len() {
1152 let input_slots = flatten_input_slots(
1153 wiring,
1154 nodes,
1155 node_idx,
1156 port_offsets,
1157 input_starts,
1158 input_widths,
1159 );
1160 let output_slots = flatten_output_slots(nodes, node_idx, port_offsets);
1161
1162 max_inputs = max_inputs.max(input_slots.len());
1163 max_outputs = max_outputs.max(output_slots.len());
1164
1165 let step = closure_step_for(
1166 &graph,
1167 node_idx,
1168 input_slots,
1169 output_slots,
1170 &mut scratch,
1171 &mut ref_scratch,
1172 )?;
1173 steps.push(HybridStep::Closure(step));
1174 }
1175 let node_step: Vec<usize> = (0..nodes.len()).collect();
1176
1177 let output_types = output_types_of(nodes, port_offsets, input_starts, input_types, &output_map);
1178 build_pushpull_from_steps(
1179 steps,
1180 scratch,
1181 ref_scratch,
1182 ref_slots,
1183 wiring,
1184 nodes,
1185 coord_count,
1186 total_slots,
1187 output_map,
1188 max_inputs,
1189 max_outputs,
1190 input_starts,
1191 input_widths,
1192 output_types,
1193 externs,
1194 constant,
1195 volatile,
1196 attribution,
1197 node_step,
1198 )
1199}
1200
1201#[derive(Clone, Copy)]
1206struct GraphView<'a> {
1207 nodes: &'a [Box<dyn PolydatNode>],
1208 wiring: &'a [Vec<WireSource>],
1209 port_offsets: &'a [Vec<usize>],
1210 input_types: &'a [crate::ast::PortType],
1211}
1212
1213fn closure_step_for(
1225 graph: &GraphView<'_>,
1226 node_idx: usize,
1227 input_slots: Vec<usize>,
1228 output_slots: Vec<usize>,
1229 scratch: &mut Vec<crate::ast::ScratchBuf>,
1230 ref_scratch: &mut Vec<(usize, usize)>,
1231) -> Result<ClosureStep, crate::KernelError> {
1232 let GraphView {
1233 nodes,
1234 wiring,
1235 port_offsets,
1236 input_types,
1237 } = *graph;
1238 let node = &nodes[node_idx];
1239 let scratch_start = scratch.len();
1240 let wire_types: Vec<crate::ast::PortType> = wiring[node_idx]
1241 .iter()
1242 .map(|src| match src {
1243 WireSource::Input(c) => input_types
1244 .get(*c)
1245 .copied()
1246 .unwrap_or(crate::ast::PortType::U64),
1247 WireSource::NodeOutput(j, p) => nodes[*j].meta().outs[*p].typ,
1248 })
1249 .collect();
1250 let op = if let Some(op) = node.compiled_u64() {
1251 ClosureOp::U64(op)
1252 } else if let Some(op) = crate::compile::assembly::identity_op(node.as_ref()) {
1253 ClosureOp::U64(op)
1254 } else if let Some(kit) = ref_copy_or_slot(node.as_ref(), &wire_types) {
1255 scratch.extend(kit.scratch.iter().map(|e| crate::ast::ScratchBuf::new(*e)));
1256 let starts = flatten_ref_output_starts(nodes, node_idx, port_offsets);
1257 ref_scratch.extend(crate::compile::assembly::scratch_pairs(
1258 &node.meta().name,
1259 &starts,
1260 &kit.scratch,
1261 scratch_start,
1262 ));
1263 ClosureOp::Slot(kit.op)
1264 } else {
1265 return Err(refused(format!(
1266 "node '{}' has no compiled form (docs/design/engines.md §8)",
1267 node.meta().name
1268 )));
1269 };
1270 Ok(ClosureStep {
1271 op,
1272 input_slots,
1273 output_slots,
1274 scratch_range: (scratch_start, scratch.len()),
1275 accepts_none: node.accepts_none_inputs(),
1276 node: node_idx,
1277 })
1278}
1279
1280#[allow(clippy::too_many_arguments)]
1286fn build_pushpull_from_steps(
1287 steps: Vec<HybridStep>,
1288 scratch: Vec<crate::ast::ScratchBuf>,
1289 ref_scratch: Vec<(usize, usize)>,
1290 ref_slots: Vec<bool>,
1291 wiring: &[Vec<WireSource>],
1292 nodes: &[Box<dyn PolydatNode>],
1293 coord_count: usize,
1294 total_slots: usize,
1295 output_map: HashMap<String, usize>,
1296 max_inputs: usize,
1297 max_outputs: usize,
1298 _input_starts: &[usize],
1299 input_widths: &[usize],
1300 output_types: HashMap<String, crate::ast::PortType>,
1301 externs: crate::compile::externs::Externs,
1302 constant: Vec<bool>,
1303 volatile: Vec<bool>,
1304 attribution: std::sync::Arc<crate::compile::Attribution>,
1305 node_step: Vec<usize>,
1306) -> Result<HybridKernelPushPull, crate::KernelError> {
1307 let step_count = steps.len();
1308 debug_assert_eq!(node_step.len(), nodes.len());
1309 debug_assert!(node_step.iter().all(|&s| s < step_count));
1310 let to_steps = |list: &[usize]| -> Vec<usize> {
1313 let mut v: Vec<usize> = list.iter().map(|&n| node_step[n]).collect();
1314 v.sort_unstable();
1315 v.dedup();
1316 v
1317 };
1318 let mut buffer = vec![0u64; total_slots + 1];
1320 let mut none = vec![false; total_slots];
1321 let any_none = externs.seed(&mut buffer, Some(&mut none));
1322
1323 let node_provenance = crate::kernel::PolydatProgram::compute_provenance(nodes, wiring);
1330 let input_dependents: Vec<Vec<usize>> =
1331 crate::kernel::PolydatProgram::compute_dependents(&node_provenance, input_widths.len())
1332 .iter()
1333 .map(|d| to_steps(d))
1334 .collect();
1335 let step_dependents: Vec<Vec<usize>> = input_widths
1336 .iter()
1337 .enumerate()
1338 .flat_map(|(i, w)| {
1339 std::iter::repeat_n(input_dependents.get(i).cloned().unwrap_or_default(), *w)
1340 })
1341 .collect();
1342
1343 let step_outs: Vec<&[usize]> = steps.iter().map(|s| s.output_slots()).collect();
1344 let slot_provenance =
1345 crate::compile::slot_provenance(coord_count, total_slots, &step_outs, &step_dependents);
1346
1347 debug_assert_eq!(constant.len(), nodes.len());
1352 debug_assert_eq!(volatile.len(), nodes.len());
1353 let mut step_constant = vec![true; step_count];
1354 let mut step_volatile = vec![false; step_count];
1355 let mut side = vec![false; step_count];
1356 for (n, node) in nodes.iter().enumerate() {
1357 let s = node_step[n];
1358 step_constant[s] &= constant[n];
1359 step_volatile[s] |= volatile[n];
1360 side[s] |= matches!(node.purity(), crate::ast::Purity::SideChannel { .. });
1361 }
1362 let volatile = step_volatile;
1363 let constants: Vec<usize> = (0..step_count).filter(|&i| step_constant[i]).collect();
1364 let step_inputs: Vec<&[usize]> = steps.iter().map(|s| s.input_slots()).collect();
1365 let step_outputs: Vec<&[usize]> = steps.iter().map(|s| s.output_slots()).collect();
1366 let plan = crate::compile::Invalidation::from_provenance(
1367 step_dependents.clone(),
1368 &step_inputs,
1369 &step_outputs,
1370 &output_map,
1371 total_slots,
1372 );
1373 let mut slot_step: Vec<Option<usize>> = vec![None; total_slots];
1374 for (i, outs) in step_outputs.iter().enumerate() {
1375 for &s in outs.iter() {
1376 slot_step[s] = Some(i);
1377 }
1378 }
1379 drop(step_inputs);
1380 drop(step_outputs);
1381
1382 let dirty: Vec<Vec<usize>> = plan.input_dependents.clone();
1383 let volatile_steps: Vec<usize> = (0..step_count).filter(|&i| volatile[i]).collect();
1384 let mut kernel = HybridKernelPushPull {
1385 core: HybridCore {
1386 engine: Engine::Native(Provenance::PushPull),
1387 buffer,
1388 coord_count,
1389 steps: std::sync::Arc::new(steps),
1390 output_map,
1391 gather_buf: vec![0u64; max_inputs.max(1)],
1392 scatter_buf: vec![0u64; max_outputs.max(1)],
1393 scratch,
1394 ref_slots,
1395 ref_scratch,
1396 output_types,
1397 externs,
1398 traversals: Vec::new().into(),
1399 resolved_outputs: Vec::new(),
1400 _nodes: std::sync::Arc::new(Vec::new()),
1401 drive: crate::compile::Drive {
1402 coords: Vec::new(),
1403 stale: true,
1404 },
1405 none,
1406 ran: vec![0; step_count],
1407 epoch: 0,
1408 all_ran: false,
1409 clean: vec![false; step_count],
1410 use_clean: true,
1411 plan: std::sync::Arc::new(plan),
1412 volatile: volatile.into(),
1413 side: side.into(),
1414 slot_step: slot_step.into(),
1415 sites: attribution,
1416 cur_step: 0,
1417 tracker: total_slots,
1418 all: (0..step_count).collect::<Vec<usize>>().into(),
1419 dirty: dirty.into(),
1420 any_none,
1421 volatile_steps: volatile_steps.into(),
1422 },
1423 slot_provenance,
1424 changed_mask: crate::kernel::ProvMask::all_below(coord_count), force_run: false,
1426 };
1427 kernel.core.begin_epoch();
1432 kernel.core.fold_steps(&constants)?;
1433 kernel.core.drive.stale = true;
1434 Ok(kernel)
1435}
1436
1437fn ref_copy_or_slot(
1441 node: &dyn PolydatNode,
1442 wire_types: &[crate::ast::PortType],
1443) -> Option<crate::ast::CompiledSlotKit> {
1444 let meta = node.meta();
1445 if (meta.name == "identity" || meta.name.starts_with("__port_"))
1446 && meta.outs.len() == 1
1447 && meta.outs[0].typ.slot_color() == crate::ast::SlotColor::Ref2
1448 {
1449 return crate::compile::assembly::ref_copy_kit(meta.outs[0].typ);
1450 }
1451 node.compiled_slot(
1452 wire_types,
1453 crate::compile::select::Engine::Native(crate::compile::select::Provenance::Auto),
1454 )
1455}
1456
1457impl HybridKernelRaw {
1460 fn mark_all_dirty(&mut self) {}
1462}
1463
1464impl HybridKernelPull {
1465 fn mark_all_dirty(&mut self) {
1467 self.changed_mask = crate::kernel::ProvMask::all_below(self.core.coord_count);
1468 self.force_run = true;
1469 }
1470}
1471
1472impl HybridKernelPushPull {
1473 fn mark_all_dirty(&mut self) {
1475 self.core.clean.fill(false);
1476 self.changed_mask = crate::kernel::ProvMask::all_below(self.core.coord_count);
1477 self.force_run = true;
1478 }
1479
1480 pub(crate) fn into_raw(self) -> HybridKernelRaw {
1483 let mut core = self.core;
1484 core.set_use_clean(false);
1485 core.engine = Engine::Native(Provenance::Raw);
1486 HybridKernelRaw { core }
1487 }
1488
1489 pub(crate) fn into_pull(self) -> HybridKernelPull {
1490 let mut core = self.core;
1491 core.set_use_clean(false);
1492 core.engine = Engine::Native(Provenance::Pull);
1493 let changed_mask = crate::kernel::ProvMask::all_below(core.coord_count);
1494 HybridKernelPull {
1495 core,
1496 slot_provenance: self.slot_provenance,
1497 changed_mask,
1498 force_run: false,
1499 }
1500 }
1501}
1502
1503use crate::compile::select::{Engine, Provenance};
1504
1505crate::compile::impl_kernel_trait!(HybridKernelRaw);
1506crate::compile::impl_kernel_trait!(HybridKernelPull);
1507crate::compile::impl_kernel_trait!(HybridKernelPushPull);
1508crate::compile::impl_slot_kernel!(HybridKernelRaw);
1509crate::compile::impl_slot_kernel!(HybridKernelPull);
1510crate::compile::impl_slot_kernel!(HybridKernelPushPull);
1511
1512#[inline(always)]
1527fn run_hybrid_step(
1528 step: &HybridStep,
1529 none_free: bool,
1530 buffer: &mut [u64],
1531 none: &mut [bool],
1532 gather: &mut [u64],
1533 scatter: &mut [u64],
1534 scratch: &mut [crate::ast::ScratchBuf],
1535) {
1536 if !none_free && !step.accepts_none() && step.input_slots().iter().any(|&s| none[s]) {
1537 for &s in step.output_slots() {
1538 none[s] = true;
1539 }
1540 return;
1541 }
1542 match step {
1543 #[cfg(feature = "jit")]
1544 HybridStep::Jit(seg) => {
1545 let code_fn = seg.code_fn;
1549 let buf_const = buffer.as_ptr();
1550 let buf_mut = buffer.as_mut_ptr();
1551 let sc = scratch.as_mut_ptr();
1552 if seg.fallible {
1553 crate::compile::jit::invoke_with_catch(move || unsafe {
1554 (code_fn)(buf_const, buf_mut, sc);
1555 });
1556 } else {
1557 unsafe { (code_fn)(buf_const, buf_mut, sc) };
1558 }
1559 }
1560 HybridStep::Closure(cs) => {
1561 for (i, &slot) in cs.input_slots.iter().enumerate() {
1562 gather[i] = buffer[slot];
1563 }
1564 match &cs.op {
1565 ClosureOp::U64(op) => op(
1566 &gather[..cs.input_slots.len()],
1567 &mut scatter[..cs.output_slots.len()],
1568 ),
1569 ClosureOp::Slot(op) => op(
1570 &gather[..cs.input_slots.len()],
1571 &mut scatter[..cs.output_slots.len()],
1572 &mut scratch[cs.scratch_range.0..cs.scratch_range.1],
1573 ),
1574 }
1575 for (i, &slot) in cs.output_slots.iter().enumerate() {
1576 buffer[slot] = scatter[i];
1577 }
1578 }
1579 }
1580 if !none_free {
1581 for &s in step.output_slots() {
1582 none[s] = false;
1583 }
1584 }
1585}