1use std::collections::HashMap;
21
22use crate::ast::{CompiledSlotOp, CompiledU64Op, PortType, ScratchBuf, ScratchElem};
23
24#[derive(Default)]
28pub(crate) struct P2Extras {
29 pub(crate) output_types: HashMap<String, PortType>,
30 pub(crate) externs: crate::compile::externs::Externs,
33 pub(crate) input_dependents: Vec<Vec<usize>>,
36 pub(crate) attribution: std::sync::Arc<crate::compile::Attribution>,
38}
39
40pub(crate) enum StepOp {
45 U64(CompiledU64Op),
46 Slot(CompiledSlotOp),
47 Copy,
50}
51pub(crate) struct P2Step {
53 pub(crate) name: String,
55 pub(crate) op: StepOp,
56 pub(crate) input_slots: Vec<usize>,
57 pub(crate) output_slots: Vec<usize>,
58 pub(crate) scratch: Vec<ScratchElem>,
60 pub(crate) ref_output_starts: Vec<usize>,
65 pub(crate) accepts_none: bool,
68 pub(crate) volatile: bool,
71 pub(crate) constant: bool,
74 pub(crate) side: bool,
78}
79
80struct CompiledStep {
81 op: StepOp,
82 input_slots: Vec<usize>,
83 output_slots: Vec<usize>,
84 scratch_range: (usize, usize),
85 accepts_none: bool,
87 volatile: bool,
89 constant: bool,
91 side: bool,
94}
95
96type ResolvedOutput = (
99 usize,
100 crate::ast::PortType,
101 Option<std::sync::Arc<[usize]>>,
102 bool,
103);
104
105struct KernelCore {
111 engine: crate::compile::select::Engine,
116 buffer: Vec<u64>,
117 coord_count: usize,
118 steps: std::sync::Arc<[CompiledStep]>,
119 output_map: HashMap<String, usize>,
120 gather_buf: Vec<u64>,
121 scatter_buf: Vec<u64>,
122 scratch: Vec<ScratchBuf>,
126 ref_slots: Vec<bool>,
129 ref_scratch: Vec<(usize, usize)>,
132 output_types: HashMap<String, PortType>,
134 externs: crate::compile::externs::Externs,
136 traversals: std::sync::Arc<[crate::dsl::traversal::Traversal]>,
139 resolved_outputs: Vec<Option<ResolvedOutput>>,
142 drive: crate::compile::Drive,
146 none: Vec<bool>,
149 ran: Vec<u64>,
152 epoch: u64,
158 all_ran: bool,
160 clean: Vec<bool>,
164 use_clean: bool,
168 plan: std::sync::Arc<crate::compile::Invalidation>,
171 slot_step: std::sync::Arc<[Option<usize>]>,
173 sites: std::sync::Arc<crate::compile::Attribution>,
175 cur_step: usize,
177 all: std::sync::Arc<[usize]>,
179 dirty: std::sync::Arc<[Vec<usize>]>,
185 volatile_steps: std::sync::Arc<[usize]>,
187 any_none: bool,
191}
192
193impl Clone for KernelCore {
194 fn clone(&self) -> Self {
195 let mut core = KernelCore {
196 engine: self.engine,
197 buffer: self.buffer.clone(),
198 coord_count: self.coord_count,
199 steps: self.steps.clone(),
200 output_map: self.output_map.clone(),
201 gather_buf: self.gather_buf.clone(),
202 scatter_buf: self.scatter_buf.clone(),
203 scratch: self.scratch.clone(),
204 ref_slots: self.ref_slots.clone(),
205 ref_scratch: self.ref_scratch.clone(),
206 output_types: self.output_types.clone(),
207 externs: self.externs.clone(),
208 traversals: self.traversals.clone(),
209 resolved_outputs: self.resolved_outputs.clone(),
210 drive: self.drive.clone(),
211 none: self.none.clone(),
212 ran: self.ran.clone(),
213 epoch: self.epoch,
214 all_ran: self.all_ran,
215 clean: self.clean.clone(),
216 use_clean: self.use_clean,
217 plan: self.plan.clone(),
218 slot_step: self.slot_step.clone(),
219 sites: self.sites.clone(),
220 cur_step: self.cur_step,
221 all: self.all.clone(),
222 dirty: self.dirty.clone(),
223 volatile_steps: self.volatile_steps.clone(),
224 any_none: self.any_none,
225 };
226 core.republish_refs();
227 core
228 }
229}
230
231impl KernelCore {
232 crate::compile::shared_core_methods!();
233 #[inline]
236 fn step_can_fail(&self, _i: usize) -> bool {
237 true
238 }
239
240 #[inline]
245 fn failing_node(&self) -> usize {
246 self.cur_step
247 }
248
249 #[inline]
251 fn run_order(&mut self, order: &[usize]) {
252 let steps = &self.steps;
253 let none_free = !self.any_none;
254 for &i in order {
255 if self.all_ran || self.ran[i] == self.epoch {
256 continue;
257 }
258 let step = &steps[i];
259 if (self.use_clean || step.side) && self.clean[i] && !step.volatile {
262 self.ran[i] = self.epoch;
263 continue;
264 }
265 self.cur_step = i;
266 if none_free {
267 run_step_fast(
268 step,
269 &mut self.buffer,
270 &mut self.gather_buf,
271 &mut self.scatter_buf,
272 &mut self.scratch,
273 );
274 } else {
275 run_step(
276 step,
277 &mut self.buffer,
278 &mut self.none,
279 &mut self.gather_buf,
280 &mut self.scatter_buf,
281 &mut self.scratch,
282 );
283 }
284 self.ran[i] = self.epoch;
285 self.clean[i] = !step.volatile;
286 }
287 }
288
289 #[inline]
295 fn run_fresh(&mut self) {
296 let steps = &self.steps;
297 for (i, step) in steps.iter().enumerate() {
298 if step.side {
299 if self.clean[i] && !step.volatile {
300 continue;
301 }
302 self.clean[i] = !step.volatile;
303 }
304 self.cur_step = i;
305 run_step_fast(
306 step,
307 &mut self.buffer,
308 &mut self.gather_buf,
309 &mut self.scatter_buf,
310 &mut self.scratch,
311 );
312 }
313 self.all_ran = true;
314 }
315
316 fn plan(&self) -> crate::EnginePlan {
318 crate::EnginePlan {
319 closure_steps: self.steps.len(),
320 ..Default::default()
321 }
322 }
323}
324
325#[allow(clippy::too_many_arguments)]
328fn build_core(
329 coord_count: usize,
330 total_slots: usize,
331 steps: Vec<P2Step>,
332 output_map: HashMap<String, usize>,
333 ref_slots: Vec<bool>,
334 extras: P2Extras,
335 use_clean: bool,
336 engine: crate::compile::select::Engine,
337) -> Result<KernelCore, crate::KernelError> {
338 let P2Extras {
339 output_types,
340 externs,
341 input_dependents,
342 attribution,
343 } = extras;
344 let max_inputs = steps.iter().map(|s| s.input_slots.len()).max().unwrap_or(0);
345 let max_outputs = steps
346 .iter()
347 .map(|s| s.output_slots.len())
348 .max()
349 .unwrap_or(0);
350 let mut scratch: Vec<ScratchBuf> = Vec::new();
351 let mut ref_scratch: Vec<(usize, usize)> = Vec::new();
352 let compiled_steps: Vec<CompiledStep> = steps
353 .into_iter()
354 .map(|step| {
355 let start = scratch.len();
356 scratch.extend(step.scratch.iter().map(|e| ScratchBuf::new(*e)));
357 ref_scratch.extend(crate::compile::assembly::scratch_pairs(
358 &step.name,
359 &step.ref_output_starts,
360 &step.scratch,
361 start,
362 ));
363 CompiledStep {
364 op: step.op,
365 input_slots: step.input_slots,
366 output_slots: step.output_slots,
367 scratch_range: (start, scratch.len()),
368 accepts_none: step.accepts_none,
369 volatile: step.volatile,
370 constant: step.constant,
371 side: step.side,
372 }
373 })
374 .collect();
375 let mut slot_step: Vec<Option<usize>> = vec![None; total_slots];
376 for (i, step) in compiled_steps.iter().enumerate() {
377 for &s in &step.output_slots {
378 slot_step[s] = Some(i);
379 }
380 }
381 let step_inputs: Vec<&[usize]> = compiled_steps
382 .iter()
383 .map(|s| s.input_slots.as_slice())
384 .collect();
385 let step_outputs: Vec<&[usize]> = compiled_steps
386 .iter()
387 .map(|s| s.output_slots.as_slice())
388 .collect();
389 let plan = crate::compile::Invalidation::from_provenance(
390 input_dependents,
391 &step_inputs,
392 &step_outputs,
393 &output_map,
394 total_slots,
395 );
396 let dirty: Vec<Vec<usize>> = plan
397 .input_dependents
398 .iter()
399 .map(|deps| {
400 if use_clean {
401 deps.clone()
402 } else {
403 deps.iter()
404 .copied()
405 .filter(|&i| compiled_steps[i].side)
406 .collect()
407 }
408 })
409 .collect();
410 let volatile_steps: Vec<usize> = (0..compiled_steps.len())
411 .filter(|&i| compiled_steps[i].volatile)
412 .collect();
413 let mut buffer = vec![0u64; total_slots];
414 let mut none = vec![false; total_slots];
415 let any_none = externs.seed(&mut buffer, Some(&mut none));
416 let step_count = compiled_steps.len();
417 let constants: Vec<usize> = compiled_steps
418 .iter()
419 .enumerate()
420 .filter(|(_, s)| s.constant)
421 .map(|(i, _)| i)
422 .collect();
423 let mut core = KernelCore {
424 engine,
425 buffer,
426 coord_count,
427 steps: compiled_steps.into(),
428 output_map,
429 gather_buf: vec![0u64; max_inputs],
430 scatter_buf: vec![0u64; max_outputs],
431 scratch,
432 ref_slots,
433 ref_scratch,
434 output_types,
435 externs,
436 traversals: Vec::new().into(),
437 resolved_outputs: Vec::new(),
438 drive: crate::compile::Drive {
439 coords: Vec::new(),
440 stale: true,
441 },
442 none,
443 ran: vec![0; step_count],
444 epoch: 0,
445 all_ran: false,
446 clean: vec![false; step_count],
447 use_clean,
448 plan: std::sync::Arc::new(plan),
449 slot_step: slot_step.into(),
450 sites: attribution,
451 cur_step: 0,
452 all: (0..step_count).collect::<Vec<usize>>().into(),
453 dirty: dirty.into(),
454 volatile_steps: volatile_steps.into(),
455 any_none,
456 };
457 core.begin_epoch();
462 core.fold_steps(&constants)?;
463 core.drive.stale = true;
464 Ok(core)
465}
466
467fn compute_slot_provenance(
470 coord_count: usize,
471 total_slots: usize,
472 input_dependents: &[Vec<usize>],
473 steps: &[CompiledStep],
474) -> Vec<crate::kernel::ProvMask> {
475 let outs: Vec<&[usize]> = steps.iter().map(|s| s.output_slots.as_slice()).collect();
476 crate::compile::slot_provenance(coord_count, total_slots, &outs, input_dependents)
477}
478
479macro_rules! closure_writes {
488 () => {
489 pub fn set_input(
494 &mut self,
495 name: &str,
496 value: crate::ast::Value,
497 ) -> Result<(), crate::kernel::WriteError> {
498 let slot = self.core.set_extern(name, value)?;
499 self.mark_input_changed(slot);
500 Ok(())
501 }
502
503 pub fn set_input_at(
505 &mut self,
506 index: usize,
507 value: crate::ast::Value,
508 ) -> Result<(), crate::kernel::WriteError> {
509 let slot = self.core.set_extern_at(index, value)?;
510 self.mark_input_changed(slot);
511 Ok(())
512 }
513
514 fn mark_all_dirty(&mut self) {
518 for i in 0..self.core.coord_count {
519 self.mark_input_changed(i);
520 }
521 }
522 };
523}
524
525#[derive(Clone)]
530pub struct CompiledKernelRaw {
532 core: KernelCore,
533}
534
535impl CompiledKernelRaw {
536 pub(crate) fn new(
537 coord_count: usize,
538 total_slots: usize,
539 steps: Vec<P2Step>,
540 output_map: HashMap<String, usize>,
541 ref_slots: Vec<bool>,
542 extras: P2Extras,
543 ) -> Result<Self, crate::KernelError> {
544 Ok(Self {
545 core: build_core(
546 coord_count,
547 total_slots,
548 steps,
549 output_map,
550 ref_slots,
551 extras,
552 false,
553 Engine::Closures(Provenance::Raw),
554 )?,
555 })
556 }
557
558 fn mark_input_changed(&mut self, slot: usize) {
561 self.core.dirty_input(slot);
562 }
563
564 #[inline]
567 fn set_coords(&mut self, coords: &[u64]) {
568 for (i, &c) in coords.iter().enumerate().take(self.core.coord_count) {
569 if self.core.buffer[i] != c {
570 self.core.buffer[i] = c;
571 self.core.dirty_input(i);
572 }
573 }
574 }
575
576 #[inline]
578 pub fn eval(&mut self, coords: &[u64]) {
579 self.set_coords(coords);
580 self.core.drive.stale = true;
581 self.core.eval_all();
582 }
583
584 #[inline]
586 pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
587 self.core.guard_ref_slot(slot);
588 self.eval(coords);
589 self.core.buffer[slot]
590 }
591
592 crate::compile::kernel_accessors!(set_coords);
593 closure_writes!();
594}
595
596#[derive(Clone)]
602pub struct CompiledKernelPush {
605 core: KernelCore,
606}
607
608impl CompiledKernelPush {
609 pub(crate) fn new(
610 coord_count: usize,
611 total_slots: usize,
612 steps: Vec<P2Step>,
613 output_map: HashMap<String, usize>,
614 input_dependents: Vec<Vec<usize>>,
615 ref_slots: Vec<bool>,
616 extras: P2Extras,
617 ) -> Result<Self, crate::KernelError> {
618 let _ = input_dependents;
620 Ok(Self {
621 core: build_core(
622 coord_count,
623 total_slots,
624 steps,
625 output_map,
626 ref_slots,
627 extras,
628 true,
629 Engine::Closures(Provenance::Push),
630 )?,
631 })
632 }
633
634 #[inline]
635 fn set_coords(&mut self, coords: &[u64]) {
636 for (i, &c) in coords.iter().enumerate().take(self.core.coord_count) {
637 if self.core.buffer[i] != c {
638 self.core.buffer[i] = c;
639 self.core.dirty_input(i);
640 }
641 }
642 }
643
644 fn mark_input_changed(&mut self, slot: usize) {
646 self.core.dirty_input(slot);
647 }
648
649 #[inline]
651 pub fn eval(&mut self, coords: &[u64]) {
652 self.set_coords(coords);
653 self.core.drive.stale = true;
654 self.core.eval_all();
655 }
656
657 #[inline]
659 pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
660 self.core.guard_ref_slot(slot);
661 self.eval(coords);
662 self.core.buffer[slot]
663 }
664
665 crate::compile::kernel_accessors!(set_coords);
666 closure_writes!();
667}
668
669#[derive(Clone)]
676pub struct CompiledKernelPull {
679 core: KernelCore,
680 slot_provenance: Vec<crate::kernel::ProvMask>,
681 changed_mask: crate::kernel::ProvMask,
682 force_run: bool,
685}
686
687impl CompiledKernelPull {
688 pub(crate) fn new(
689 coord_count: usize,
690 total_slots: usize,
691 steps: Vec<P2Step>,
692 output_map: HashMap<String, usize>,
693 input_dependents: &[Vec<usize>],
694 ref_slots: Vec<bool>,
695 extras: P2Extras,
696 ) -> Result<Self, crate::KernelError> {
697 let core = build_core(
698 coord_count,
699 total_slots,
700 steps,
701 output_map,
702 ref_slots,
703 extras,
704 false,
705 Engine::Closures(Provenance::Pull),
706 )?;
707 let slot_provenance =
708 compute_slot_provenance(coord_count, total_slots, input_dependents, &core.steps);
709 Ok(Self {
710 core,
711 slot_provenance,
712 changed_mask: crate::kernel::ProvMask::all_below(coord_count), force_run: false,
714 })
715 }
716
717 #[inline]
720 fn set_coords(&mut self, coords: &[u64]) {
721 self.changed_mask.clear();
722 for (i, &c) in coords.iter().enumerate().take(self.core.coord_count) {
723 if self.core.buffer[i] != c {
724 self.core.buffer[i] = c;
725 self.changed_mask.set(i);
726 self.core.dirty_input(i);
727 }
728 }
729 }
730
731 fn mark_input_changed(&mut self, slot: usize) {
734 self.core.dirty_input(slot);
735 self.force_run = true;
736 }
737
738 #[inline]
740 pub fn eval(&mut self, coords: &[u64]) {
741 self.set_coords(coords);
742 self.force_run = false;
743 self.core.drive.stale = true;
744 self.core.eval_all();
745 }
746
747 #[inline]
750 pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
751 self.core.guard_ref_slot(slot);
752 self.set_coords(coords);
753 if !self.force_run
754 && slot < self.slot_provenance.len()
755 && !self.slot_provenance[slot].intersects(&self.changed_mask)
756 {
757 return self.core.buffer[slot];
758 }
759 self.force_run = false;
760 self.core.drive.stale = true;
761 self.core.eval_all();
762 self.core.buffer[slot]
763 }
764
765 crate::compile::kernel_accessors!(set_coords);
766 closure_writes!();
767}
768
769#[derive(Clone)]
775pub struct CompiledKernelPushPull {
777 core: KernelCore,
778 slot_provenance: Vec<crate::kernel::ProvMask>,
779 changed_mask: crate::kernel::ProvMask,
780 force_run: bool,
783}
784
785impl CompiledKernelPushPull {
786 pub(crate) fn new(
787 coord_count: usize,
788 total_slots: usize,
789 steps: Vec<P2Step>,
790 output_map: HashMap<String, usize>,
791 input_dependents: Vec<Vec<usize>>,
792 ref_slots: Vec<bool>,
793 extras: P2Extras,
794 ) -> Result<Self, crate::KernelError> {
795 let core = build_core(
796 coord_count,
797 total_slots,
798 steps,
799 output_map,
800 ref_slots,
801 extras,
802 true,
803 Engine::Closures(Provenance::PushPull),
804 )?;
805 let slot_provenance =
806 compute_slot_provenance(coord_count, total_slots, &input_dependents, &core.steps);
807 Ok(Self {
808 core,
809 slot_provenance,
810 changed_mask: crate::kernel::ProvMask::all_below(coord_count),
811 force_run: false,
812 })
813 }
814
815 #[inline]
816 fn set_coords(&mut self, coords: &[u64]) {
817 self.changed_mask.clear();
818 for (i, &c) in coords.iter().enumerate().take(self.core.coord_count) {
819 if self.core.buffer[i] != c {
820 self.core.buffer[i] = c;
821 self.changed_mask.set(i);
822 self.core.dirty_input(i);
823 }
824 }
825 }
826
827 fn mark_input_changed(&mut self, slot: usize) {
830 self.core.dirty_input(slot);
831 self.force_run = true;
832 }
833
834 #[inline]
836 pub fn eval(&mut self, coords: &[u64]) {
837 self.set_coords(coords);
838 self.force_run = false;
839 self.core.drive.stale = true;
840 self.core.eval_all();
841 }
842
843 #[inline]
845 pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
846 self.core.guard_ref_slot(slot);
847 self.set_coords(coords);
848 if !self.force_run
849 && slot < self.slot_provenance.len()
850 && !self.slot_provenance[slot].intersects(&self.changed_mask)
851 {
852 return self.core.buffer[slot];
853 }
854 self.force_run = false;
855 self.core.drive.stale = true;
856 self.core.eval_all();
857 self.core.buffer[slot]
858 }
859
860 crate::compile::kernel_accessors!(set_coords);
861 closure_writes!();
862}
863
864use crate::compile::select::{Engine, Provenance};
867
868crate::compile::impl_kernel_trait!(CompiledKernelRaw);
869crate::compile::impl_kernel_trait!(CompiledKernelPush);
870crate::compile::impl_kernel_trait!(CompiledKernelPull);
871crate::compile::impl_kernel_trait!(CompiledKernelPushPull);
872crate::compile::impl_slot_kernel!(CompiledKernelRaw);
873crate::compile::impl_slot_kernel!(CompiledKernelPush);
874crate::compile::impl_slot_kernel!(CompiledKernelPull);
875crate::compile::impl_slot_kernel!(CompiledKernelPushPull);
876
877#[inline(always)]
881fn run_step(
882 step: &CompiledStep,
883 buffer: &mut [u64],
884 none: &mut [bool],
885 gather: &mut [u64],
886 scatter: &mut [u64],
887 scratch: &mut [ScratchBuf],
888) {
889 let mut any_none = false;
890 for (i, &s) in step.input_slots.iter().enumerate() {
891 gather[i] = buffer[s];
892 any_none |= none[s];
893 }
894 if any_none && !step.accepts_none {
895 for &s in &step.output_slots {
896 none[s] = true;
897 }
898 return;
899 }
900 if matches!(step.op, StepOp::Copy) {
901 for (&i, &o) in step.input_slots.iter().zip(&step.output_slots) {
902 buffer[o] = buffer[i];
903 none[o] = false;
904 }
905 return;
906 }
907 let (n_in, n_out) = (step.input_slots.len(), step.output_slots.len());
908 match &step.op {
909 StepOp::Copy => unreachable!(),
910 StepOp::U64(op) => op(&gather[..n_in], &mut scatter[..n_out]),
911 StepOp::Slot(op) => op(
912 &gather[..n_in],
913 &mut scatter[..n_out],
914 &mut scratch[step.scratch_range.0..step.scratch_range.1],
915 ),
916 }
917 for (i, &s) in step.output_slots.iter().enumerate() {
918 buffer[s] = scatter[i];
919 none[s] = false;
920 }
921}
922
923#[inline(always)]
925fn run_step_fast(
926 step: &CompiledStep,
927 buffer: &mut [u64],
928 gather: &mut [u64],
929 scatter: &mut [u64],
930 scratch: &mut [ScratchBuf],
931) {
932 if matches!(step.op, StepOp::Copy) {
933 for (&i, &o) in step.input_slots.iter().zip(&step.output_slots) {
934 buffer[o] = buffer[i];
935 }
936 return;
937 }
938 for (i, &s) in step.input_slots.iter().enumerate() {
939 gather[i] = buffer[s];
940 }
941 let (n_in, n_out) = (step.input_slots.len(), step.output_slots.len());
942 match &step.op {
943 StepOp::Copy => unreachable!(),
944 StepOp::U64(op) => op(&gather[..n_in], &mut scatter[..n_out]),
945 StepOp::Slot(op) => op(
946 &gather[..n_in],
947 &mut scatter[..n_out],
948 &mut scratch[step.scratch_range.0..step.scratch_range.1],
949 ),
950 }
951 for (i, &s) in step.output_slots.iter().enumerate() {
952 buffer[s] = scatter[i];
953 }
954}