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 {
91 match self {
92 #[cfg(feature = "jit")]
93 HybridStep::Jit(_) => false,
94 HybridStep::Closure(cs) => cs.accepts_none,
95 }
96 }
97
98 #[cfg_attr(not(feature = "jit"), allow(unused_variables))]
101 fn failing_node(&self, buffer: &[u64], tracker: usize) -> usize {
102 match self {
103 #[cfg(feature = "jit")]
104 HybridStep::Jit(seg) => seg
105 .nodes
106 .get(buffer[tracker] as usize)
107 .copied()
108 .unwrap_or(usize::MAX),
109 HybridStep::Closure(cs) => cs.node,
110 }
111 }
112}
113
114enum ClosureOp {
118 U64(CompiledU64Op),
119 Slot(crate::ast::CompiledSlotOp),
120}
121
122struct ClosureStep {
123 op: ClosureOp,
124 input_slots: Vec<usize>,
125 output_slots: Vec<usize>,
126 scratch_range: (usize, usize),
128 accepts_none: bool,
130 node: usize,
132}
133
134type ResolvedOutput = (
137 usize,
138 crate::ast::PortType,
139 Option<std::sync::Arc<[usize]>>,
140 bool,
141);
142
143struct HybridCore {
149 engine: crate::compile::select::Engine,
154 buffer: Vec<u64>,
155 coord_count: usize,
156 steps: std::sync::Arc<Vec<HybridStep>>,
157 output_map: HashMap<String, usize>,
158 gather_buf: Vec<u64>,
159 scatter_buf: Vec<u64>,
160 scratch: Vec<crate::ast::ScratchBuf>,
164 ref_slots: Vec<bool>,
166 ref_scratch: Vec<(usize, usize)>,
168 output_types: HashMap<String, crate::ast::PortType>,
170 externs: crate::compile::externs::Externs,
172 traversals: std::sync::Arc<[crate::dsl::traversal::Traversal]>,
175 resolved_outputs: Vec<Option<ResolvedOutput>>,
178 _nodes: std::sync::Arc<Vec<Box<dyn PolydatNode>>>,
180 drive: crate::compile::Drive,
184 none: Vec<bool>,
186 ran: Vec<u64>,
189 epoch: u64,
195 all_ran: bool,
197 clean: Vec<bool>,
201 use_clean: bool,
203 plan: std::sync::Arc<crate::compile::Invalidation>,
206 volatile: std::sync::Arc<[bool]>,
208 side: std::sync::Arc<[bool]>,
210 slot_step: std::sync::Arc<[Option<usize>]>,
212 sites: std::sync::Arc<crate::compile::Attribution>,
214 cur_step: usize,
216 tracker: usize,
219 all: std::sync::Arc<[usize]>,
221 dirty: std::sync::Arc<[Vec<usize>]>,
227 any_none: bool,
231 volatile_steps: std::sync::Arc<[usize]>,
233}
234
235impl Clone for HybridCore {
236 fn clone(&self) -> Self {
237 let mut core = HybridCore {
238 engine: self.engine,
239 buffer: self.buffer.clone(),
240 coord_count: self.coord_count,
241 steps: self.steps.clone(),
242 output_map: self.output_map.clone(),
243 gather_buf: self.gather_buf.clone(),
244 scatter_buf: self.scatter_buf.clone(),
245 scratch: self.scratch.clone(),
246 ref_slots: self.ref_slots.clone(),
247 ref_scratch: self.ref_scratch.clone(),
248 output_types: self.output_types.clone(),
249 externs: self.externs.clone(),
250 traversals: self.traversals.clone(),
251 resolved_outputs: self.resolved_outputs.clone(),
252 _nodes: self._nodes.clone(),
253 drive: self.drive.clone(),
254 none: self.none.clone(),
255 ran: self.ran.clone(),
256 epoch: self.epoch,
257 all_ran: self.all_ran,
258 clean: self.clean.clone(),
259 use_clean: self.use_clean,
260 plan: self.plan.clone(),
261 volatile: self.volatile.clone(),
262 side: self.side.clone(),
263 slot_step: self.slot_step.clone(),
264 sites: self.sites.clone(),
265 cur_step: self.cur_step,
266 tracker: self.tracker,
267 all: self.all.clone(),
268 dirty: self.dirty.clone(),
269 any_none: self.any_none,
270 volatile_steps: self.volatile_steps.clone(),
271 };
272 core.republish_refs();
273 core
274 }
275}
276
277impl HybridCore {
278 crate::compile::shared_core_methods!();
279}
280
281impl HybridCore {
282 fn set_use_clean(&mut self, on: bool) {
286 self.use_clean = on;
287 let side = std::sync::Arc::clone(&self.side);
288 self.dirty = self
289 .plan
290 .input_dependents
291 .iter()
292 .map(|deps| {
293 if on {
294 deps.clone()
295 } else {
296 deps.iter().copied().filter(|&i| side[i]).collect()
297 }
298 })
299 .collect::<Vec<_>>()
300 .into();
301 }
302
303 fn program_identity(&self) -> usize {
309 std::sync::Arc::as_ptr(&self.steps) as *const () as usize
310 }
311
312 #[inline]
313 fn step_can_fail(&self, i: usize) -> bool {
314 match &self.steps[i] {
315 #[cfg(feature = "jit")]
316 HybridStep::Jit(seg) => seg.fallible,
317 _ => true,
318 }
319 }
320
321 #[inline]
327 fn failing_node(&self) -> usize {
328 self.steps[self.cur_step].failing_node(&self.buffer, self.tracker)
329 }
330
331 #[inline]
333 fn run_order(&mut self, order: &[usize]) {
334 let steps = &self.steps;
335 let none_free = !self.any_none;
336 for &i in order {
337 if self.all_ran || self.ran[i] == self.epoch {
338 continue;
339 }
340 let never = self.volatile[i];
341 if (self.use_clean || self.side[i]) && self.clean[i] && !never {
342 self.ran[i] = self.epoch;
343 continue;
344 }
345 self.cur_step = i;
346 run_hybrid_step(
347 &steps[i],
348 none_free,
349 &mut self.buffer,
350 &mut self.none,
351 &mut self.gather_buf,
352 &mut self.scatter_buf,
353 &mut self.scratch,
354 );
355 self.ran[i] = self.epoch;
356 self.clean[i] = !never;
357 }
358 }
359
360 #[inline]
366 fn run_fresh(&mut self) {
367 let steps = &self.steps;
368 for (i, step) in steps.iter().enumerate() {
369 if self.side[i] {
370 let never = self.volatile[i];
371 if self.clean[i] && !never {
372 continue;
373 }
374 self.clean[i] = !never;
375 }
376 self.cur_step = i;
377 run_hybrid_step(
378 step,
379 true,
380 &mut self.buffer,
381 &mut self.none,
382 &mut self.gather_buf,
383 &mut self.scatter_buf,
384 &mut self.scratch,
385 );
386 }
387 self.all_ran = true;
388 }
389
390 fn plan(&self) -> crate::EnginePlan {
392 let (native_segments, closure_steps) = self.engine_counts();
393 crate::EnginePlan {
394 native_segments,
395 closure_steps,
396 interpreted_nodes: 0,
397 }
398 }
399}
400
401#[inline]
404fn eval_all_hybrid_steps(core: &mut HybridCore) {
405 core.drive.stale = true;
406 core.eval_all();
407}
408
409impl HybridCore {
414 fn engine_counts(&self) -> (usize, usize) {
417 let closures = self
418 .steps
419 .iter()
420 .filter(|s| matches!(s, HybridStep::Closure(_)))
421 .count();
422 (self.steps.len() - closures, closures)
423 }
424}
425
426#[derive(Clone)]
431pub struct HybridKernelRaw {
432 core: HybridCore,
433}
434
435impl HybridKernelRaw {
436 crate::compile::kernel_accessors!(set_coords);
437 #[inline]
440 fn set_coords(&mut self, coords: &[u64]) {
441 for (i, &c) in coords
442 .iter()
443 .enumerate()
444 .take(self.core.externs.coordinate_slots())
445 {
446 if self.core.buffer[i] != c {
447 self.core.buffer[i] = c;
448 self.core.dirty_input(i);
449 }
450 }
451 }
452
453 #[inline]
455 pub fn eval(&mut self, coords: &[u64]) {
456 self.set_coords(coords);
457 eval_all_hybrid_steps(&mut self.core);
458 }
459
460 pub fn set_input(
464 &mut self,
465 name: &str,
466 value: crate::ast::Value,
467 ) -> Result<(), crate::kernel::WriteError> {
468 self.core.set_extern(name, value).map(|_| ())
469 }
470
471 pub fn set_input_at(
473 &mut self,
474 index: usize,
475 value: crate::ast::Value,
476 ) -> Result<(), crate::kernel::WriteError> {
477 self.core.set_extern_at(index, value).map(|_| ())
478 }
479
480 #[inline]
482 pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
483 self.core.guard_ref_slot(slot);
484 self.eval(coords);
485 self.core.buffer[slot]
486 }
487
488 pub fn engine_counts(&self) -> (usize, usize) {
491 self.core.engine_counts()
492 }
493
494 pub fn retain_nodes(&mut self, nodes: Vec<Box<dyn PolydatNode>>) {
496 self.core._nodes = std::sync::Arc::new(nodes);
497 }
498}
499
500#[derive(Clone)]
512pub struct HybridKernelPull {
513 core: HybridCore,
514 slot_provenance: Vec<crate::kernel::ProvMask>,
515 changed_mask: crate::kernel::ProvMask,
516 force_run: bool,
519}
520
521impl HybridKernelPull {
522 crate::compile::kernel_accessors!(set_inputs);
523 #[inline]
526 fn set_inputs(&mut self, coords: &[u64]) {
527 self.changed_mask.clear();
528 for (i, &c) in coords
529 .iter()
530 .enumerate()
531 .take(self.core.externs.coordinate_slots())
532 {
533 if self.core.buffer[i] != c {
534 self.core.buffer[i] = c;
535 self.changed_mask.set(i);
536 self.core.dirty_input(i);
537 }
538 }
539 }
540
541 #[inline]
543 pub fn eval(&mut self, coords: &[u64]) {
544 self.set_inputs(coords);
545 self.force_run = false;
546 eval_all_hybrid_steps(&mut self.core);
547 }
548
549 #[inline]
552 pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
553 self.core.guard_ref_slot(slot);
554 self.set_inputs(coords);
555 if !self.force_run
556 && slot < self.slot_provenance.len()
557 && !self.slot_provenance[slot].intersects(&self.changed_mask)
558 {
559 return self.core.buffer[slot];
560 }
561 self.force_run = false;
562 eval_all_hybrid_steps(&mut self.core);
563 self.core.buffer[slot]
564 }
565
566 pub fn set_input(
570 &mut self,
571 name: &str,
572 value: crate::ast::Value,
573 ) -> Result<(), crate::kernel::WriteError> {
574 self.core.set_extern(name, value)?;
575 self.force_run = true;
576 Ok(())
577 }
578
579 pub fn set_input_at(
581 &mut self,
582 index: usize,
583 value: crate::ast::Value,
584 ) -> Result<(), crate::kernel::WriteError> {
585 self.core.set_extern_at(index, value)?;
586 self.force_run = true;
587 Ok(())
588 }
589
590 pub fn engine_counts(&self) -> (usize, usize) {
593 self.core.engine_counts()
594 }
595
596 pub fn retain_nodes(&mut self, nodes: Vec<Box<dyn PolydatNode>>) {
598 self.core._nodes = std::sync::Arc::new(nodes);
599 }
600}
601
602#[derive(Clone)]
616pub struct HybridKernelPushPull {
617 core: HybridCore,
618 slot_provenance: Vec<crate::kernel::ProvMask>,
619 changed_mask: crate::kernel::ProvMask,
620 force_run: bool,
623}
624
625impl HybridKernelPushPull {
626 crate::compile::kernel_accessors!(set_inputs);
627 pub fn set_input(
632 &mut self,
633 name: &str,
634 value: crate::ast::Value,
635 ) -> Result<(), crate::kernel::WriteError> {
636 self.core.set_extern(name, value)?;
637 self.force_run = true;
638 Ok(())
639 }
640
641 pub fn set_input_at(
643 &mut self,
644 index: usize,
645 value: crate::ast::Value,
646 ) -> Result<(), crate::kernel::WriteError> {
647 self.core.set_extern_at(index, value)?;
648 self.force_run = true;
649 Ok(())
650 }
651
652 #[inline]
654 fn set_inputs(&mut self, coords: &[u64]) {
655 self.changed_mask.clear();
656 for (i, &c) in coords
657 .iter()
658 .enumerate()
659 .take(self.core.externs.coordinate_slots())
660 {
661 if self.core.buffer[i] != c {
662 self.core.buffer[i] = c;
663 self.changed_mask.set(i);
664 self.core.dirty_input(i);
665 }
666 }
667 }
668
669 #[inline]
671 pub fn eval(&mut self, coords: &[u64]) {
672 self.set_inputs(coords);
673 self.force_run = false;
674 self.core.drive.stale = true;
675 self.core.eval_all();
676 }
677
678 #[inline]
680 pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
681 self.core.guard_ref_slot(slot);
682 self.set_inputs(coords);
683 if !self.force_run
684 && slot < self.slot_provenance.len()
685 && !self.slot_provenance[slot].intersects(&self.changed_mask)
686 {
687 return self.core.buffer[slot];
688 }
689 self.force_run = false;
690 self.core.drive.stale = true;
691 self.core.eval_all();
692 self.core.buffer[slot]
693 }
694
695 pub fn engine_counts(&self) -> (usize, usize) {
698 self.core.engine_counts()
699 }
700
701 pub fn retain_nodes(&mut self, nodes: Vec<Box<dyn PolydatNode>>) {
703 self.core._nodes = std::sync::Arc::new(nodes);
704 }
705}
706
707pub type HybridKernel = HybridKernelPushPull;
712
713fn flatten_input_slots(
717 wiring: &[Vec<WireSource>],
718 nodes: &[Box<dyn PolydatNode>],
719 node_idx: usize,
720 port_offsets: &[Vec<usize>],
721 input_starts: &[usize],
722 input_widths: &[usize],
723) -> Vec<usize> {
724 let mut slots = Vec::new();
725 for source in &wiring[node_idx] {
726 let (start, w) = match source {
727 WireSource::Input(c) => (
728 input_starts.get(*c).copied().unwrap_or(*c),
729 input_widths.get(*c).copied().unwrap_or(1),
730 ),
731 WireSource::NodeOutput(u, p) => (
732 port_offsets[*u][*p],
733 nodes[*u].meta().outs[*p].typ.slot_width(),
734 ),
735 };
736 slots.extend(start..start + w);
737 }
738 slots
739}
740
741fn flatten_ref_output_starts(
744 nodes: &[Box<dyn PolydatNode>],
745 node_idx: usize,
746 port_offsets: &[Vec<usize>],
747) -> Vec<usize> {
748 nodes[node_idx]
749 .meta()
750 .outs
751 .iter()
752 .enumerate()
753 .filter(|(_, out)| out.typ.slot_color() == crate::ast::SlotColor::Ref2)
754 .map(|(p, _)| port_offsets[node_idx][p])
755 .collect()
756}
757
758fn flatten_output_slots(
760 nodes: &[Box<dyn PolydatNode>],
761 node_idx: usize,
762 port_offsets: &[Vec<usize>],
763) -> Vec<usize> {
764 let mut slots = Vec::new();
765 for (p, out) in nodes[node_idx].meta().outs.iter().enumerate() {
766 let start = port_offsets[node_idx][p];
767 slots.extend(start..start + out.typ.slot_width());
768 }
769 slots
770}
771
772fn refused(reason: String) -> crate::KernelError {
783 crate::KernelError::Refused {
784 engine: crate::compile::select::Engine::Native(crate::compile::select::Provenance::Auto),
785 reason,
786 }
787}
788
789#[cfg(feature = "jit")]
791#[allow(clippy::too_many_arguments)]
792pub(crate) fn build_hybrid(
793 nodes: &[Box<dyn PolydatNode>],
794 wiring: &[Vec<WireSource>],
795 coord_count: usize,
796 total_slots: usize,
797 port_offsets: &[Vec<usize>],
798 input_starts: &[usize],
799 input_widths: &[usize],
800 output_map: HashMap<String, usize>,
801 ref_slots: Vec<bool>,
802 input_types: &[crate::ast::PortType],
803 externs: crate::compile::externs::Externs,
804 constant: Vec<bool>,
805 volatile: Vec<bool>,
806 attribution: std::sync::Arc<crate::compile::Attribution>,
807) -> Result<HybridKernelPushPull, crate::KernelError> {
808 let mut steps: Vec<Option<HybridStep>> = Vec::new();
811 let mut pending: Vec<PendingSegment> = Vec::new();
812 let mut scratch: Vec<crate::ast::ScratchBuf> = Vec::new();
813 let mut ref_scratch: Vec<(usize, usize)> = Vec::new();
814 let mut max_inputs = 0usize;
815 let mut max_outputs = 0usize;
816 let graph = GraphView {
817 nodes,
818 wiring,
819 port_offsets,
820 input_types,
821 };
822
823 let classifications: Vec<(JitOp, Vec<usize>, Vec<usize>)> = nodes
825 .iter()
826 .enumerate()
827 .map(|(node_idx, node)| {
828 let wire_types: Vec<crate::ast::PortType> = wiring[node_idx]
832 .iter()
833 .map(|src| match src {
834 WireSource::Input(c) => input_types
835 .get(*c)
836 .copied()
837 .unwrap_or(crate::ast::PortType::U64),
838 WireSource::NodeOutput(j, p) => nodes[*j].meta().outs[*p].typ,
839 })
840 .collect();
841 let jit_op = jit::classify_node_typed(node.as_ref(), &wire_types);
842
843 let input_slots = flatten_input_slots(
844 wiring,
845 nodes,
846 node_idx,
847 port_offsets,
848 input_starts,
849 input_widths,
850 );
851 let output_slots = flatten_output_slots(nodes, node_idx, port_offsets);
852
853 max_inputs = max_inputs.max(input_slots.len());
854 max_outputs = max_outputs.max(output_slots.len());
855
856 (jit_op, input_slots, output_slots)
857 })
858 .collect();
859 let mut classifications = classifications;
869 let mut eligible = vec![false; nodes.len()];
870 for (node_idx, node) in nodes.iter().enumerate() {
871 if matches!(classifications[node_idx].0, JitOp::Fallback) {
872 continue;
873 }
874 if !crate::compile::none_rule_admits(
875 node.accepts_none_inputs(),
876 &wiring[node_idx],
877 &eligible,
878 ) {
879 classifications[node_idx].0 = JitOp::Fallback;
880 continue;
881 }
882 eligible[node_idx] = true;
883 }
884 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())
911 .filter(|&k| constant[k])
912 .chain((0..nodes.len()).filter(|&k| !constant[k]))
913 .collect();
914 let mut rank = vec![0usize; nodes.len()];
915 for (pos, &k) in order.iter().enumerate() {
916 rank[k] = pos;
917 }
918 let is_side = |k: usize| matches!(nodes[k].purity(), crate::ast::Purity::SideChannel { .. });
931 let preds: Vec<Vec<usize>> = wiring
932 .iter()
933 .map(|w| {
934 w.iter()
935 .filter_map(|src| match src {
936 WireSource::NodeOutput(j, _) => Some(*j),
937 WireSource::Input(_) => None,
938 })
939 .collect()
940 })
941 .collect();
942 let inputs_read: Vec<Vec<usize>> = wiring
943 .iter()
944 .map(|w| {
945 w.iter()
946 .filter_map(|src| match src {
947 WireSource::Input(c) => Some(*c),
948 WireSource::NodeOutput(..) => None,
949 })
950 .collect()
951 })
952 .collect();
953 let fusible: Vec<bool> = (0..nodes.len())
954 .map(|k| !matches!(classifications[k].0, JitOp::Fallback) && !is_side(k))
955 .collect();
956 let class: Vec<u64> = (0..nodes.len())
957 .map(|k| constant[k] as u64 | (volatile[k] as u64) << 1)
958 .collect();
959 let unset_read = externs.unset_read_slots();
960 let reads: Vec<Vec<usize>> = inputs_read
961 .iter()
962 .map(|ins| {
963 ins.iter()
964 .map(|&c| input_starts[c])
965 .filter(|s| unset_read.binary_search(s).is_ok())
966 .collect()
967 })
968 .collect();
969 let class = crate::compile::fusion_units::refine_by_externs(&preds, &reads, &class);
970 let plan = crate::compile::fusion_units::plan_units(
971 &preds,
972 &inputs_read,
973 &fusible,
974 &class,
975 &rank,
976 &|c| c & 1 == 1,
977 );
978 for members in plan.units {
979 let i = members[0];
980 if matches!(classifications[i].0, JitOp::Fallback) {
981 let (_, ref input_slots, ref output_slots) = classifications[i];
985 let step = closure_step_for(
986 &graph,
987 i,
988 input_slots.clone(),
989 output_slots.clone(),
990 &mut scratch,
991 &mut ref_scratch,
992 )?;
993 node_step[i] = steps.len();
994 steps.push(Some(HybridStep::Closure(step)));
995 } else {
996 for &k in &members {
1000 let base = scratch.len();
1001 classifications[k].0.place_scratch(base);
1002 let elems = classifications[k].0.scratch_elems().to_vec();
1003 ref_scratch.extend(crate::compile::assembly::scratch_pairs(
1004 &nodes[k].meta().name,
1005 &flatten_ref_output_starts(nodes, k, port_offsets),
1006 &elems,
1007 base,
1008 ));
1009 scratch.extend(elems.iter().map(|e| crate::ast::ScratchBuf::new(*e)));
1010 }
1011 let batch: Vec<(JitOp, Vec<usize>, Vec<usize>)> = members
1019 .iter()
1020 .map(|&k| classifications[k].clone())
1021 .collect();
1022 let written: std::collections::HashSet<usize> = batch
1023 .iter()
1024 .flat_map(|(_, _, o)| o.iter().copied())
1025 .collect();
1026 let mut input_slots: Vec<usize> = Vec::new();
1027 for (_, ins, _) in &batch {
1028 for &s in ins {
1029 if !written.contains(&s) && !input_slots.contains(&s) {
1030 input_slots.push(s);
1031 }
1032 }
1033 }
1034 let output_slots: Vec<usize> = batch
1035 .iter()
1036 .flat_map(|(_, _, o)| o.iter().copied())
1037 .collect();
1038 let segment = steps.len();
1039 for &k in &members {
1040 node_step[k] = segment;
1041 }
1042 steps.push(None);
1043 pending.push(PendingSegment {
1044 step: segment,
1045 batch,
1046 input_slots,
1047 output_slots,
1048 nodes: members,
1049 });
1050 }
1051 }
1052 let batches: Vec<&[jit::JitStep]> = pending.iter().map(|p| p.batch.as_slice()).collect();
1056 let (entries, code) = if batches.is_empty() {
1057 (Vec::new(), None)
1058 } else {
1059 let (entries, code) =
1060 jit::compile_jit_entries(&batches, Some(total_slots)).map_err(refused)?;
1061 (entries, Some(code))
1062 };
1063 for (p, (code_fn, fallible)) in pending.into_iter().zip(entries) {
1064 steps[p.step] = Some(HybridStep::Jit(JitSegment {
1065 code_fn,
1066 fallible,
1067 _module: code.clone().expect("a segment was compiled"),
1068 input_slots: p.input_slots,
1069 output_slots: p.output_slots,
1070 nodes: p.nodes,
1071 }));
1072 }
1073 let steps: Vec<HybridStep> = steps
1074 .into_iter()
1075 .map(|s| s.expect("every step is placed"))
1076 .collect();
1077
1078 let output_types = output_types_of(nodes, port_offsets, input_starts, input_types, &output_map);
1079 build_pushpull_from_steps(
1080 steps,
1081 scratch,
1082 ref_scratch,
1083 ref_slots,
1084 wiring,
1085 nodes,
1086 coord_count,
1087 total_slots,
1088 output_map,
1089 max_inputs,
1090 max_outputs,
1091 input_starts,
1092 input_widths,
1093 output_types,
1094 externs,
1095 constant,
1096 volatile,
1097 attribution,
1098 node_step,
1099 )
1100}
1101
1102fn output_types_of(
1105 nodes: &[Box<dyn PolydatNode>],
1106 port_offsets: &[Vec<usize>],
1107 input_starts: &[usize],
1108 input_types: &[crate::ast::PortType],
1109 output_map: &HashMap<String, usize>,
1110) -> HashMap<String, crate::ast::PortType> {
1111 let mut slot_types: HashMap<usize, crate::ast::PortType> = HashMap::new();
1112 for (start, ty) in input_starts.iter().zip(input_types) {
1113 slot_types.insert(*start, *ty);
1114 }
1115 for (node_idx, node) in nodes.iter().enumerate() {
1116 for (p, out) in node.meta().outs.iter().enumerate() {
1117 slot_types.insert(port_offsets[node_idx][p], out.typ);
1118 }
1119 }
1120 output_map
1121 .iter()
1122 .map(|(name, slot)| {
1123 (
1124 name.clone(),
1125 slot_types
1126 .get(slot)
1127 .copied()
1128 .unwrap_or(crate::ast::PortType::U64),
1129 )
1130 })
1131 .collect()
1132}
1133
1134#[cfg(not(feature = "jit"))]
1136#[allow(clippy::too_many_arguments)]
1137pub(crate) fn build_hybrid(
1138 nodes: &[Box<dyn PolydatNode>],
1139 wiring: &[Vec<WireSource>],
1140 coord_count: usize,
1141 total_slots: usize,
1142 port_offsets: &[Vec<usize>],
1143 input_starts: &[usize],
1144 input_widths: &[usize],
1145 output_map: HashMap<String, usize>,
1146 ref_slots: Vec<bool>,
1147 input_types: &[crate::ast::PortType],
1148 externs: crate::compile::externs::Externs,
1149 constant: Vec<bool>,
1150 volatile: Vec<bool>,
1151 attribution: std::sync::Arc<crate::compile::Attribution>,
1152) -> Result<HybridKernelPushPull, crate::KernelError> {
1153 let mut steps: Vec<HybridStep> = Vec::new();
1154 let mut scratch: Vec<crate::ast::ScratchBuf> = Vec::new();
1155 let mut ref_scratch: Vec<(usize, usize)> = Vec::new();
1156 let mut max_inputs = 0usize;
1157 let mut max_outputs = 0usize;
1158 let graph = GraphView {
1159 nodes,
1160 wiring,
1161 port_offsets,
1162 input_types,
1163 };
1164
1165 for node_idx in 0..nodes.len() {
1166 let input_slots = flatten_input_slots(
1167 wiring,
1168 nodes,
1169 node_idx,
1170 port_offsets,
1171 input_starts,
1172 input_widths,
1173 );
1174 let output_slots = flatten_output_slots(nodes, node_idx, port_offsets);
1175
1176 max_inputs = max_inputs.max(input_slots.len());
1177 max_outputs = max_outputs.max(output_slots.len());
1178
1179 let step = closure_step_for(
1180 &graph,
1181 node_idx,
1182 input_slots,
1183 output_slots,
1184 &mut scratch,
1185 &mut ref_scratch,
1186 )?;
1187 steps.push(HybridStep::Closure(step));
1188 }
1189 let node_step: Vec<usize> = (0..nodes.len()).collect();
1190
1191 let output_types = output_types_of(nodes, port_offsets, input_starts, input_types, &output_map);
1192 build_pushpull_from_steps(
1193 steps,
1194 scratch,
1195 ref_scratch,
1196 ref_slots,
1197 wiring,
1198 nodes,
1199 coord_count,
1200 total_slots,
1201 output_map,
1202 max_inputs,
1203 max_outputs,
1204 input_starts,
1205 input_widths,
1206 output_types,
1207 externs,
1208 constant,
1209 volatile,
1210 attribution,
1211 node_step,
1212 )
1213}
1214
1215#[derive(Clone, Copy)]
1220struct GraphView<'a> {
1221 nodes: &'a [Box<dyn PolydatNode>],
1222 wiring: &'a [Vec<WireSource>],
1223 port_offsets: &'a [Vec<usize>],
1224 input_types: &'a [crate::ast::PortType],
1225}
1226
1227fn closure_step_for(
1239 graph: &GraphView<'_>,
1240 node_idx: usize,
1241 input_slots: Vec<usize>,
1242 output_slots: Vec<usize>,
1243 scratch: &mut Vec<crate::ast::ScratchBuf>,
1244 ref_scratch: &mut Vec<(usize, usize)>,
1245) -> Result<ClosureStep, crate::KernelError> {
1246 let GraphView {
1247 nodes,
1248 wiring,
1249 port_offsets,
1250 input_types,
1251 } = *graph;
1252 let node = &nodes[node_idx];
1253 let scratch_start = scratch.len();
1254 let wire_types: Vec<crate::ast::PortType> = wiring[node_idx]
1255 .iter()
1256 .map(|src| match src {
1257 WireSource::Input(c) => input_types
1258 .get(*c)
1259 .copied()
1260 .unwrap_or(crate::ast::PortType::U64),
1261 WireSource::NodeOutput(j, p) => nodes[*j].meta().outs[*p].typ,
1262 })
1263 .collect();
1264 let op = if let Some(op) = node.compiled_u64() {
1265 ClosureOp::U64(op)
1266 } else if let Some(op) = crate::compile::assembly::identity_op(node.as_ref()) {
1267 ClosureOp::U64(op)
1268 } else if let Some(kit) = ref_copy_or_slot(node.as_ref(), &wire_types) {
1269 scratch.extend(kit.scratch.iter().map(|e| crate::ast::ScratchBuf::new(*e)));
1270 let starts = flatten_ref_output_starts(nodes, node_idx, port_offsets);
1271 ref_scratch.extend(crate::compile::assembly::scratch_pairs(
1272 &node.meta().name,
1273 &starts,
1274 &kit.scratch,
1275 scratch_start,
1276 ));
1277 ClosureOp::Slot(kit.op)
1278 } else {
1279 return Err(refused(format!(
1280 "node '{}' has no compiled form (docs/design/engines.md §8)",
1281 node.meta().name
1282 )));
1283 };
1284 Ok(ClosureStep {
1285 op,
1286 input_slots,
1287 output_slots,
1288 scratch_range: (scratch_start, scratch.len()),
1289 accepts_none: node.accepts_none_inputs(),
1290 node: node_idx,
1291 })
1292}
1293
1294#[allow(clippy::too_many_arguments)]
1300fn build_pushpull_from_steps(
1301 steps: Vec<HybridStep>,
1302 scratch: Vec<crate::ast::ScratchBuf>,
1303 ref_scratch: Vec<(usize, usize)>,
1304 ref_slots: Vec<bool>,
1305 wiring: &[Vec<WireSource>],
1306 nodes: &[Box<dyn PolydatNode>],
1307 coord_count: usize,
1308 total_slots: usize,
1309 output_map: HashMap<String, usize>,
1310 max_inputs: usize,
1311 max_outputs: usize,
1312 _input_starts: &[usize],
1313 input_widths: &[usize],
1314 output_types: HashMap<String, crate::ast::PortType>,
1315 externs: crate::compile::externs::Externs,
1316 constant: Vec<bool>,
1317 volatile: Vec<bool>,
1318 attribution: std::sync::Arc<crate::compile::Attribution>,
1319 node_step: Vec<usize>,
1320) -> Result<HybridKernelPushPull, crate::KernelError> {
1321 let step_count = steps.len();
1322 debug_assert_eq!(node_step.len(), nodes.len());
1323 debug_assert!(node_step.iter().all(|&s| s < step_count));
1324 let to_steps = |list: &[usize]| -> Vec<usize> {
1327 let mut v: Vec<usize> = list.iter().map(|&n| node_step[n]).collect();
1328 v.sort_unstable();
1329 v.dedup();
1330 v
1331 };
1332 let mut buffer = vec![0u64; total_slots + 1];
1334 let mut none = vec![false; total_slots];
1335 let any_none = externs.seed(&mut buffer, Some(&mut none));
1336
1337 let node_provenance = crate::kernel::PolydatProgram::compute_provenance(nodes, wiring);
1344 let input_dependents: Vec<Vec<usize>> =
1345 crate::kernel::PolydatProgram::compute_dependents(&node_provenance, input_widths.len())
1346 .iter()
1347 .map(|d| to_steps(d))
1348 .collect();
1349 let step_dependents: Vec<Vec<usize>> = input_widths
1350 .iter()
1351 .enumerate()
1352 .flat_map(|(i, w)| {
1353 std::iter::repeat_n(input_dependents.get(i).cloned().unwrap_or_default(), *w)
1354 })
1355 .collect();
1356
1357 let step_outs: Vec<&[usize]> = steps.iter().map(|s| s.output_slots()).collect();
1358 let slot_provenance =
1359 crate::compile::slot_provenance(coord_count, total_slots, &step_outs, &step_dependents);
1360
1361 debug_assert_eq!(constant.len(), nodes.len());
1366 debug_assert_eq!(volatile.len(), nodes.len());
1367 let mut step_constant = vec![true; step_count];
1368 let mut step_volatile = vec![false; step_count];
1369 let mut side = vec![false; step_count];
1370 for (n, node) in nodes.iter().enumerate() {
1371 let s = node_step[n];
1372 step_constant[s] &= constant[n];
1373 step_volatile[s] |= volatile[n];
1374 side[s] |= matches!(node.purity(), crate::ast::Purity::SideChannel { .. });
1375 }
1376 let volatile = step_volatile;
1377 let constants: Vec<usize> = (0..step_count).filter(|&i| step_constant[i]).collect();
1378 let step_inputs: Vec<&[usize]> = steps.iter().map(|s| s.input_slots()).collect();
1379 let step_outputs: Vec<&[usize]> = steps.iter().map(|s| s.output_slots()).collect();
1380 let plan = crate::compile::Invalidation::from_provenance(
1381 step_dependents.clone(),
1382 &step_inputs,
1383 &step_outputs,
1384 &output_map,
1385 total_slots,
1386 );
1387 let mut slot_step: Vec<Option<usize>> = vec![None; total_slots];
1388 for (i, outs) in step_outputs.iter().enumerate() {
1389 for &s in outs.iter() {
1390 slot_step[s] = Some(i);
1391 }
1392 }
1393 drop(step_inputs);
1394 drop(step_outputs);
1395
1396 let dirty: Vec<Vec<usize>> = plan.input_dependents.clone();
1397 let volatile_steps: Vec<usize> = (0..step_count).filter(|&i| volatile[i]).collect();
1398 let mut kernel = HybridKernelPushPull {
1399 core: HybridCore {
1400 engine: Engine::Native(Provenance::PushPull),
1401 buffer,
1402 coord_count,
1403 steps: std::sync::Arc::new(steps),
1404 output_map,
1405 gather_buf: vec![0u64; max_inputs.max(1)],
1406 scatter_buf: vec![0u64; max_outputs.max(1)],
1407 scratch,
1408 ref_slots,
1409 ref_scratch,
1410 output_types,
1411 externs,
1412 traversals: Vec::new().into(),
1413 resolved_outputs: Vec::new(),
1414 _nodes: std::sync::Arc::new(Vec::new()),
1415 drive: crate::compile::Drive {
1416 coords: Vec::new(),
1417 stale: true,
1418 },
1419 none,
1420 ran: vec![0; step_count],
1421 epoch: 0,
1422 all_ran: false,
1423 clean: vec![false; step_count],
1424 use_clean: true,
1425 plan: std::sync::Arc::new(plan),
1426 volatile: volatile.into(),
1427 side: side.into(),
1428 slot_step: slot_step.into(),
1429 sites: attribution,
1430 cur_step: 0,
1431 tracker: total_slots,
1432 all: (0..step_count).collect::<Vec<usize>>().into(),
1433 dirty: dirty.into(),
1434 any_none,
1435 volatile_steps: volatile_steps.into(),
1436 },
1437 slot_provenance,
1438 changed_mask: crate::kernel::ProvMask::all_below(coord_count), force_run: false,
1440 };
1441 kernel.core.begin_epoch();
1446 kernel.core.fold_steps(&constants)?;
1447 kernel.core.drive.stale = true;
1448 Ok(kernel)
1449}
1450
1451fn ref_copy_or_slot(
1455 node: &dyn PolydatNode,
1456 wire_types: &[crate::ast::PortType],
1457) -> Option<crate::ast::CompiledSlotKit> {
1458 let meta = node.meta();
1459 if (meta.name == "identity" || meta.name.starts_with("__port_"))
1460 && meta.outs.len() == 1
1461 && meta.outs[0].typ.slot_color() == crate::ast::SlotColor::Ref2
1462 {
1463 return crate::compile::assembly::ref_copy_kit(meta.outs[0].typ);
1464 }
1465 node.compiled_slot(
1466 wire_types,
1467 crate::compile::select::Engine::Native(crate::compile::select::Provenance::Auto),
1468 )
1469}
1470
1471impl HybridKernelRaw {
1474 fn mark_all_dirty(&mut self) {}
1476}
1477
1478impl HybridKernelPull {
1479 fn mark_all_dirty(&mut self) {
1481 self.changed_mask = crate::kernel::ProvMask::all_below(self.core.coord_count);
1482 self.force_run = true;
1483 }
1484}
1485
1486impl HybridKernelPushPull {
1487 fn mark_all_dirty(&mut self) {
1489 self.core.clean.fill(false);
1490 self.changed_mask = crate::kernel::ProvMask::all_below(self.core.coord_count);
1491 self.force_run = true;
1492 }
1493
1494 pub(crate) fn into_raw(self) -> HybridKernelRaw {
1497 let mut core = self.core;
1498 core.set_use_clean(false);
1499 core.engine = Engine::Native(Provenance::Raw);
1500 HybridKernelRaw { core }
1501 }
1502
1503 pub(crate) fn into_pull(self) -> HybridKernelPull {
1504 let mut core = self.core;
1505 core.set_use_clean(false);
1506 core.engine = Engine::Native(Provenance::Pull);
1507 let changed_mask = crate::kernel::ProvMask::all_below(core.coord_count);
1508 HybridKernelPull {
1509 core,
1510 slot_provenance: self.slot_provenance,
1511 changed_mask,
1512 force_run: false,
1513 }
1514 }
1515}
1516
1517use crate::compile::select::{Engine, Provenance};
1518
1519crate::compile::impl_kernel_trait!(HybridKernelRaw);
1520crate::compile::impl_kernel_trait!(HybridKernelPull);
1521crate::compile::impl_kernel_trait!(HybridKernelPushPull);
1522crate::compile::impl_slot_kernel!(HybridKernelRaw);
1523crate::compile::impl_slot_kernel!(HybridKernelPull);
1524crate::compile::impl_slot_kernel!(HybridKernelPushPull);
1525
1526#[inline(always)]
1541fn run_hybrid_step(
1542 step: &HybridStep,
1543 none_free: bool,
1544 buffer: &mut [u64],
1545 none: &mut [bool],
1546 gather: &mut [u64],
1547 scatter: &mut [u64],
1548 scratch: &mut [crate::ast::ScratchBuf],
1549) {
1550 if !none_free && !step.accepts_none() && step.input_slots().iter().any(|&s| none[s]) {
1551 for &s in step.output_slots() {
1552 none[s] = true;
1553 }
1554 return;
1555 }
1556 match step {
1557 #[cfg(feature = "jit")]
1558 HybridStep::Jit(seg) => {
1559 let code_fn = seg.code_fn;
1563 let buf_const = buffer.as_ptr();
1564 let buf_mut = buffer.as_mut_ptr();
1565 let sc = scratch.as_mut_ptr();
1566 if seg.fallible {
1567 crate::compile::jit::invoke_with_catch(move || unsafe {
1568 (code_fn)(buf_const, buf_mut, sc);
1569 });
1570 } else {
1571 unsafe { (code_fn)(buf_const, buf_mut, sc) };
1572 }
1573 }
1574 HybridStep::Closure(cs) => {
1575 for (i, &slot) in cs.input_slots.iter().enumerate() {
1576 gather[i] = buffer[slot];
1577 }
1578 match &cs.op {
1579 ClosureOp::U64(op) => op(
1580 &gather[..cs.input_slots.len()],
1581 &mut scatter[..cs.output_slots.len()],
1582 ),
1583 ClosureOp::Slot(op) => op(
1584 &gather[..cs.input_slots.len()],
1585 &mut scatter[..cs.output_slots.len()],
1586 &mut scratch[cs.scratch_range.0..cs.scratch_range.1],
1587 ),
1588 }
1589 for (i, &slot) in cs.output_slots.iter().enumerate() {
1590 buffer[slot] = scatter[i];
1591 }
1592 }
1593 }
1594 if !none_free {
1595 for &s in step.output_slots() {
1596 none[s] = false;
1597 }
1598 }
1599}