use std::collections::HashMap;
use crate::ast::{CompiledSlotOp, CompiledU64Op, PortType, ScratchBuf, ScratchElem};
#[derive(Default)]
pub(crate) struct P2Extras {
pub(crate) output_types: HashMap<String, PortType>,
pub(crate) externs: crate::compile::externs::Externs,
pub(crate) input_dependents: Vec<Vec<usize>>,
pub(crate) attribution: std::sync::Arc<crate::compile::Attribution>,
}
pub(crate) enum StepOp {
U64(CompiledU64Op),
Slot(CompiledSlotOp),
Copy,
}
pub(crate) struct P2Step {
pub(crate) name: String,
pub(crate) op: StepOp,
pub(crate) input_slots: Vec<usize>,
pub(crate) output_slots: Vec<usize>,
pub(crate) scratch: Vec<ScratchElem>,
pub(crate) ref_output_starts: Vec<usize>,
pub(crate) accepts_none: bool,
pub(crate) volatile: bool,
pub(crate) constant: bool,
pub(crate) side: bool,
}
struct CompiledStep {
op: StepOp,
input_slots: Vec<usize>,
output_slots: Vec<usize>,
scratch_range: (usize, usize),
accepts_none: bool,
volatile: bool,
constant: bool,
side: bool,
}
type ResolvedOutput = (
usize,
crate::ast::PortType,
Option<std::sync::Arc<[usize]>>,
bool,
);
struct KernelCore {
engine: crate::compile::select::Engine,
buffer: Vec<u64>,
coord_count: usize,
steps: std::sync::Arc<[CompiledStep]>,
output_map: HashMap<String, usize>,
gather_buf: Vec<u64>,
scatter_buf: Vec<u64>,
scratch: Vec<ScratchBuf>,
ref_slots: Vec<bool>,
ref_scratch: Vec<(usize, usize)>,
output_types: HashMap<String, PortType>,
externs: crate::compile::externs::Externs,
traversals: std::sync::Arc<[crate::dsl::traversal::Traversal]>,
resolved_outputs: Vec<Option<ResolvedOutput>>,
drive: crate::compile::Drive,
none: Vec<bool>,
ran: Vec<u64>,
epoch: u64,
all_ran: bool,
clean: Vec<bool>,
use_clean: bool,
plan: std::sync::Arc<crate::compile::Invalidation>,
slot_step: std::sync::Arc<[Option<usize>]>,
sites: std::sync::Arc<crate::compile::Attribution>,
cur_step: usize,
all: std::sync::Arc<[usize]>,
dirty: std::sync::Arc<[Vec<usize>]>,
volatile_steps: std::sync::Arc<[usize]>,
any_none: bool,
}
impl Clone for KernelCore {
fn clone(&self) -> Self {
let mut core = KernelCore {
engine: self.engine,
buffer: self.buffer.clone(),
coord_count: self.coord_count,
steps: self.steps.clone(),
output_map: self.output_map.clone(),
gather_buf: self.gather_buf.clone(),
scatter_buf: self.scatter_buf.clone(),
scratch: self.scratch.clone(),
ref_slots: self.ref_slots.clone(),
ref_scratch: self.ref_scratch.clone(),
output_types: self.output_types.clone(),
externs: self.externs.clone(),
traversals: self.traversals.clone(),
resolved_outputs: self.resolved_outputs.clone(),
drive: self.drive.clone(),
none: self.none.clone(),
ran: self.ran.clone(),
epoch: self.epoch,
all_ran: self.all_ran,
clean: self.clean.clone(),
use_clean: self.use_clean,
plan: self.plan.clone(),
slot_step: self.slot_step.clone(),
sites: self.sites.clone(),
cur_step: self.cur_step,
all: self.all.clone(),
dirty: self.dirty.clone(),
volatile_steps: self.volatile_steps.clone(),
any_none: self.any_none,
};
core.republish_refs();
core
}
}
impl KernelCore {
crate::compile::shared_core_methods!();
#[inline]
fn step_can_fail(&self, _i: usize) -> bool {
true
}
fn program_identity(&self) -> usize {
std::sync::Arc::as_ptr(&self.steps) as *const () as usize
}
#[inline]
fn failing_node(&self) -> usize {
self.cur_step
}
#[inline]
fn run_order(&mut self, order: &[usize]) {
let steps = &self.steps;
let none_free = !self.any_none;
for &i in order {
if self.all_ran || self.ran[i] == self.epoch {
continue;
}
let step = &steps[i];
if (self.use_clean || step.side) && self.clean[i] && !step.volatile {
self.ran[i] = self.epoch;
continue;
}
self.cur_step = i;
if none_free {
run_step_fast(
step,
&mut self.buffer,
&mut self.gather_buf,
&mut self.scatter_buf,
&mut self.scratch,
);
} else {
run_step(
step,
&mut self.buffer,
&mut self.none,
&mut self.gather_buf,
&mut self.scatter_buf,
&mut self.scratch,
);
}
self.ran[i] = self.epoch;
self.clean[i] = !step.volatile;
}
}
#[inline]
fn run_fresh(&mut self) {
let steps = &self.steps;
for (i, step) in steps.iter().enumerate() {
if step.side {
if self.clean[i] && !step.volatile {
continue;
}
self.clean[i] = !step.volatile;
}
self.cur_step = i;
run_step_fast(
step,
&mut self.buffer,
&mut self.gather_buf,
&mut self.scatter_buf,
&mut self.scratch,
);
}
self.all_ran = true;
}
fn plan(&self) -> crate::EnginePlan {
crate::EnginePlan {
closure_steps: self.steps.len(),
..Default::default()
}
}
}
#[allow(clippy::too_many_arguments)]
fn build_core(
coord_count: usize,
total_slots: usize,
steps: Vec<P2Step>,
output_map: HashMap<String, usize>,
ref_slots: Vec<bool>,
extras: P2Extras,
use_clean: bool,
engine: crate::compile::select::Engine,
) -> Result<KernelCore, crate::KernelError> {
let P2Extras {
output_types,
externs,
input_dependents,
attribution,
} = extras;
let max_inputs = steps.iter().map(|s| s.input_slots.len()).max().unwrap_or(0);
let max_outputs = steps
.iter()
.map(|s| s.output_slots.len())
.max()
.unwrap_or(0);
let mut scratch: Vec<ScratchBuf> = Vec::new();
let mut ref_scratch: Vec<(usize, usize)> = Vec::new();
let compiled_steps: Vec<CompiledStep> = steps
.into_iter()
.map(|step| {
let start = scratch.len();
scratch.extend(step.scratch.iter().map(|e| ScratchBuf::new(*e)));
ref_scratch.extend(crate::compile::assembly::scratch_pairs(
&step.name,
&step.ref_output_starts,
&step.scratch,
start,
));
CompiledStep {
op: step.op,
input_slots: step.input_slots,
output_slots: step.output_slots,
scratch_range: (start, scratch.len()),
accepts_none: step.accepts_none,
volatile: step.volatile,
constant: step.constant,
side: step.side,
}
})
.collect();
let mut slot_step: Vec<Option<usize>> = vec![None; total_slots];
for (i, step) in compiled_steps.iter().enumerate() {
for &s in &step.output_slots {
slot_step[s] = Some(i);
}
}
let step_inputs: Vec<&[usize]> = compiled_steps
.iter()
.map(|s| s.input_slots.as_slice())
.collect();
let step_outputs: Vec<&[usize]> = compiled_steps
.iter()
.map(|s| s.output_slots.as_slice())
.collect();
let plan = crate::compile::Invalidation::from_provenance(
input_dependents,
&step_inputs,
&step_outputs,
&output_map,
total_slots,
);
let dirty: Vec<Vec<usize>> = plan
.input_dependents
.iter()
.map(|deps| {
if use_clean {
deps.clone()
} else {
deps.iter()
.copied()
.filter(|&i| compiled_steps[i].side)
.collect()
}
})
.collect();
let volatile_steps: Vec<usize> = (0..compiled_steps.len())
.filter(|&i| compiled_steps[i].volatile)
.collect();
let mut buffer = vec![0u64; total_slots];
let mut none = vec![false; total_slots];
let any_none = externs.seed(&mut buffer, Some(&mut none));
let step_count = compiled_steps.len();
let constants: Vec<usize> = compiled_steps
.iter()
.enumerate()
.filter(|(_, s)| s.constant)
.map(|(i, _)| i)
.collect();
let mut core = KernelCore {
engine,
buffer,
coord_count,
steps: compiled_steps.into(),
output_map,
gather_buf: vec![0u64; max_inputs],
scatter_buf: vec![0u64; max_outputs],
scratch,
ref_slots,
ref_scratch,
output_types,
externs,
traversals: Vec::new().into(),
resolved_outputs: Vec::new(),
drive: crate::compile::Drive {
coords: Vec::new(),
stale: true,
},
none,
ran: vec![0; step_count],
epoch: 0,
all_ran: false,
clean: vec![false; step_count],
use_clean,
plan: std::sync::Arc::new(plan),
slot_step: slot_step.into(),
sites: attribution,
cur_step: 0,
all: (0..step_count).collect::<Vec<usize>>().into(),
dirty: dirty.into(),
volatile_steps: volatile_steps.into(),
any_none,
};
core.begin_epoch();
core.fold_steps(&constants)?;
core.drive.stale = true;
Ok(core)
}
fn compute_slot_provenance(
coord_count: usize,
total_slots: usize,
input_dependents: &[Vec<usize>],
steps: &[CompiledStep],
) -> Vec<crate::kernel::ProvMask> {
let outs: Vec<&[usize]> = steps.iter().map(|s| s.output_slots.as_slice()).collect();
crate::compile::slot_provenance(coord_count, total_slots, &outs, input_dependents)
}
macro_rules! closure_writes {
() => {
pub fn set_input(
&mut self,
name: &str,
value: crate::ast::Value,
) -> Result<(), crate::kernel::WriteError> {
let slot = self.core.set_extern(name, value)?;
self.mark_input_changed(slot);
Ok(())
}
pub fn set_input_at(
&mut self,
index: usize,
value: crate::ast::Value,
) -> Result<(), crate::kernel::WriteError> {
let slot = self.core.set_extern_at(index, value)?;
self.mark_input_changed(slot);
Ok(())
}
fn mark_all_dirty(&mut self) {
for i in 0..self.core.coord_count {
self.mark_input_changed(i);
}
}
};
}
#[derive(Clone)]
pub struct CompiledKernelRaw {
core: KernelCore,
}
impl CompiledKernelRaw {
pub(crate) fn new(
coord_count: usize,
total_slots: usize,
steps: Vec<P2Step>,
output_map: HashMap<String, usize>,
ref_slots: Vec<bool>,
extras: P2Extras,
) -> Result<Self, crate::KernelError> {
Ok(Self {
core: build_core(
coord_count,
total_slots,
steps,
output_map,
ref_slots,
extras,
false,
Engine::Closures(Provenance::Raw),
)?,
})
}
fn mark_input_changed(&mut self, slot: usize) {
self.core.dirty_input(slot);
}
#[inline]
fn set_coords(&mut self, coords: &[u64]) {
for (i, &c) in coords
.iter()
.enumerate()
.take(self.core.externs.coordinate_slots())
{
if self.core.buffer[i] != c {
self.core.buffer[i] = c;
self.core.dirty_input(i);
}
}
}
#[inline]
pub fn eval(&mut self, coords: &[u64]) {
self.set_coords(coords);
self.core.drive.stale = true;
self.core.eval_all();
}
#[inline]
pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
self.core.guard_ref_slot(slot);
self.eval(coords);
self.core.buffer[slot]
}
crate::compile::kernel_accessors!(set_coords);
closure_writes!();
}
#[derive(Clone)]
pub struct CompiledKernelPush {
core: KernelCore,
}
impl CompiledKernelPush {
pub(crate) fn new(
coord_count: usize,
total_slots: usize,
steps: Vec<P2Step>,
output_map: HashMap<String, usize>,
input_dependents: Vec<Vec<usize>>,
ref_slots: Vec<bool>,
extras: P2Extras,
) -> Result<Self, crate::KernelError> {
let _ = input_dependents;
Ok(Self {
core: build_core(
coord_count,
total_slots,
steps,
output_map,
ref_slots,
extras,
true,
Engine::Closures(Provenance::Push),
)?,
})
}
#[inline]
fn set_coords(&mut self, coords: &[u64]) {
for (i, &c) in coords
.iter()
.enumerate()
.take(self.core.externs.coordinate_slots())
{
if self.core.buffer[i] != c {
self.core.buffer[i] = c;
self.core.dirty_input(i);
}
}
}
fn mark_input_changed(&mut self, slot: usize) {
self.core.dirty_input(slot);
}
#[inline]
pub fn eval(&mut self, coords: &[u64]) {
self.set_coords(coords);
self.core.drive.stale = true;
self.core.eval_all();
}
#[inline]
pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
self.core.guard_ref_slot(slot);
self.eval(coords);
self.core.buffer[slot]
}
crate::compile::kernel_accessors!(set_coords);
closure_writes!();
}
#[derive(Clone)]
pub struct CompiledKernelPull {
core: KernelCore,
slot_provenance: Vec<crate::kernel::ProvMask>,
changed_mask: crate::kernel::ProvMask,
force_run: bool,
}
impl CompiledKernelPull {
pub(crate) fn new(
coord_count: usize,
total_slots: usize,
steps: Vec<P2Step>,
output_map: HashMap<String, usize>,
input_dependents: &[Vec<usize>],
ref_slots: Vec<bool>,
extras: P2Extras,
) -> Result<Self, crate::KernelError> {
let core = build_core(
coord_count,
total_slots,
steps,
output_map,
ref_slots,
extras,
false,
Engine::Closures(Provenance::Pull),
)?;
let slot_provenance =
compute_slot_provenance(coord_count, total_slots, input_dependents, &core.steps);
Ok(Self {
core,
slot_provenance,
changed_mask: crate::kernel::ProvMask::all_below(coord_count), force_run: false,
})
}
#[inline]
fn set_coords(&mut self, coords: &[u64]) {
self.changed_mask.clear();
for (i, &c) in coords
.iter()
.enumerate()
.take(self.core.externs.coordinate_slots())
{
if self.core.buffer[i] != c {
self.core.buffer[i] = c;
self.changed_mask.set(i);
self.core.dirty_input(i);
}
}
}
fn mark_input_changed(&mut self, slot: usize) {
self.core.dirty_input(slot);
self.force_run = true;
}
#[inline]
pub fn eval(&mut self, coords: &[u64]) {
self.set_coords(coords);
self.force_run = false;
self.core.drive.stale = true;
self.core.eval_all();
}
#[inline]
pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
self.core.guard_ref_slot(slot);
self.set_coords(coords);
if !self.force_run
&& slot < self.slot_provenance.len()
&& !self.slot_provenance[slot].intersects(&self.changed_mask)
{
return self.core.buffer[slot];
}
self.force_run = false;
self.core.drive.stale = true;
self.core.eval_all();
self.core.buffer[slot]
}
crate::compile::kernel_accessors!(set_coords);
closure_writes!();
}
#[derive(Clone)]
pub struct CompiledKernelPushPull {
core: KernelCore,
slot_provenance: Vec<crate::kernel::ProvMask>,
changed_mask: crate::kernel::ProvMask,
force_run: bool,
}
impl CompiledKernelPushPull {
pub(crate) fn new(
coord_count: usize,
total_slots: usize,
steps: Vec<P2Step>,
output_map: HashMap<String, usize>,
input_dependents: Vec<Vec<usize>>,
ref_slots: Vec<bool>,
extras: P2Extras,
) -> Result<Self, crate::KernelError> {
let core = build_core(
coord_count,
total_slots,
steps,
output_map,
ref_slots,
extras,
true,
Engine::Closures(Provenance::PushPull),
)?;
let slot_provenance =
compute_slot_provenance(coord_count, total_slots, &input_dependents, &core.steps);
Ok(Self {
core,
slot_provenance,
changed_mask: crate::kernel::ProvMask::all_below(coord_count),
force_run: false,
})
}
#[inline]
fn set_coords(&mut self, coords: &[u64]) {
self.changed_mask.clear();
for (i, &c) in coords
.iter()
.enumerate()
.take(self.core.externs.coordinate_slots())
{
if self.core.buffer[i] != c {
self.core.buffer[i] = c;
self.changed_mask.set(i);
self.core.dirty_input(i);
}
}
}
fn mark_input_changed(&mut self, slot: usize) {
self.core.dirty_input(slot);
self.force_run = true;
}
#[inline]
pub fn eval(&mut self, coords: &[u64]) {
self.set_coords(coords);
self.force_run = false;
self.core.drive.stale = true;
self.core.eval_all();
}
#[inline]
pub fn eval_for_slot(&mut self, coords: &[u64], slot: usize) -> u64 {
self.core.guard_ref_slot(slot);
self.set_coords(coords);
if !self.force_run
&& slot < self.slot_provenance.len()
&& !self.slot_provenance[slot].intersects(&self.changed_mask)
{
return self.core.buffer[slot];
}
self.force_run = false;
self.core.drive.stale = true;
self.core.eval_all();
self.core.buffer[slot]
}
crate::compile::kernel_accessors!(set_coords);
closure_writes!();
}
use crate::compile::select::{Engine, Provenance};
crate::compile::impl_kernel_trait!(CompiledKernelRaw);
crate::compile::impl_kernel_trait!(CompiledKernelPush);
crate::compile::impl_kernel_trait!(CompiledKernelPull);
crate::compile::impl_kernel_trait!(CompiledKernelPushPull);
crate::compile::impl_slot_kernel!(CompiledKernelRaw);
crate::compile::impl_slot_kernel!(CompiledKernelPush);
crate::compile::impl_slot_kernel!(CompiledKernelPull);
crate::compile::impl_slot_kernel!(CompiledKernelPushPull);
#[inline(always)]
fn run_step(
step: &CompiledStep,
buffer: &mut [u64],
none: &mut [bool],
gather: &mut [u64],
scatter: &mut [u64],
scratch: &mut [ScratchBuf],
) {
let mut any_none = false;
for (i, &s) in step.input_slots.iter().enumerate() {
gather[i] = buffer[s];
any_none |= none[s];
}
if any_none && !step.accepts_none {
for &s in &step.output_slots {
none[s] = true;
}
return;
}
if matches!(step.op, StepOp::Copy) {
for (&i, &o) in step.input_slots.iter().zip(&step.output_slots) {
buffer[o] = buffer[i];
none[o] = false;
}
return;
}
let (n_in, n_out) = (step.input_slots.len(), step.output_slots.len());
match &step.op {
StepOp::Copy => unreachable!(),
StepOp::U64(op) => op(&gather[..n_in], &mut scatter[..n_out]),
StepOp::Slot(op) => op(
&gather[..n_in],
&mut scatter[..n_out],
&mut scratch[step.scratch_range.0..step.scratch_range.1],
),
}
for (i, &s) in step.output_slots.iter().enumerate() {
buffer[s] = scatter[i];
none[s] = false;
}
}
#[inline(always)]
fn run_step_fast(
step: &CompiledStep,
buffer: &mut [u64],
gather: &mut [u64],
scatter: &mut [u64],
scratch: &mut [ScratchBuf],
) {
if matches!(step.op, StepOp::Copy) {
for (&i, &o) in step.input_slots.iter().zip(&step.output_slots) {
buffer[o] = buffer[i];
}
return;
}
for (i, &s) in step.input_slots.iter().enumerate() {
gather[i] = buffer[s];
}
let (n_in, n_out) = (step.input_slots.len(), step.output_slots.len());
match &step.op {
StepOp::Copy => unreachable!(),
StepOp::U64(op) => op(&gather[..n_in], &mut scatter[..n_out]),
StepOp::Slot(op) => op(
&gather[..n_in],
&mut scatter[..n_out],
&mut scratch[step.scratch_range.0..step.scratch_range.1],
),
}
for (i, &s) in step.output_slots.iter().enumerate() {
buffer[s] = scatter[i];
}
}