use super::*;
#[test]
fn fixed_point_batch_reuses_ifds_scratch_across_prepared_solves() {
let solve_phase = Cell::new(0usize);
let prepare_calls = Cell::new(0usize);
let second_prepare_saw_retained_slots = Cell::new(false);
let second_solve_saw_retained_slot = Cell::new(false);
let dispatch = |_: &vyre::ir::Program,
inputs: &[&[u8]],
grid: Option<[u32; 3]>,
outputs: &mut Vec<Vec<u8>>| {
if inputs.len() < 6 {
return crate::fixed_point_closure::fallback_bitset_equal_dispatch(inputs, outputs);
}
match inputs.len() {
17 => {
assert_eq!(grid, Some([1, 1, 1]));
if prepare_calls.get() == 1 && outputs.len() == 4 {
second_prepare_saw_retained_slots.set(true);
}
prepare_calls.set(prepare_calls.get() + 1);
outputs.clear();
outputs.push(Vec::new());
outputs.push(Vec::new());
outputs.push(Vec::new());
outputs.push(Vec::new());
for word in [0u32, 1, 1] {
outputs[0].extend_from_slice(&word.to_le_bytes());
}
for word in [0u32, 0] {
outputs[1].extend_from_slice(&word.to_le_bytes());
}
outputs[2].extend_from_slice(&1u32.to_le_bytes());
outputs[3].extend_from_slice(&1u32.to_le_bytes());
}
7 => {
assert_eq!(grid, Some([2, 1, 1]));
if outputs.is_empty() {
outputs.push(Vec::with_capacity(8));
} else {
second_solve_saw_retained_slot.set(true);
}
let frontier = u32::from_le_bytes(inputs[5][0..4].try_into().unwrap());
let next = match solve_phase.get() {
0 => {
assert_eq!(frontier, 0b01);
solve_phase.set(1);
0b11u32
}
1 => {
assert_eq!(frontier, 0b11);
solve_phase.set(2);
0b11u32
}
2 => {
assert_eq!(frontier, 0b10);
solve_phase.set(3);
0b11u32
}
3 => {
assert_eq!(frontier, 0b11);
solve_phase.set(4);
0b11u32
}
call => panic!("unexpected IFDS dispatch call {call}"),
};
outputs[0].clear();
outputs[0].extend_from_slice(&next.to_le_bytes());
}
len => panic!("unexpected IFDS dispatch input count {len}"),
}
Ok(())
};
let mut batch = FixedPointBatch::new(&dispatch);
batch
.prepare_ifds(1, 2, 1, &[(0, 0, 1)], &[], &[], &[])
.expect("batch must prepare IFDS CSR once");
let prepared = batch
.prepare_ifds(1, 2, 1, &[(0, 0, 1)], &[], &[], &[])
.expect("batch must reuse IFDS CSR prepare scratch");
assert!(second_prepare_saw_retained_slots.get());
assert_eq!(batch.ifds_prepare_scratch().output_slot_count(), 4);
let mut result = Vec::with_capacity(2);
let result_ptr = result.as_ptr();
let result_capacity = result.capacity();
batch
.ifds_prepared_into(&prepared, &[(0, 0, 0)], 4, &mut result)
.expect("first batch IFDS solve must fill caller-owned result");
assert_eq!(result, vec![0, 1024]);
let frontier_capacity = batch.ifds_scratch().frontier_word_capacity();
let frontier_byte_capacity = batch.ifds_scratch().frontier_byte_capacity();
let next_capacity = batch.ifds_scratch().next_word_capacity();
let output_slots = batch.ifds_scratch().output_slot_count();
result.clear();
batch
.ifds_prepared_into(&prepared, &[(0, 1, 0)], 4, &mut result)
.expect("second batch IFDS solve must reuse caller-owned result");
assert_eq!(result, vec![0, 1024]);
assert_eq!(result.as_ptr(), result_ptr);
assert_eq!(result.capacity(), result_capacity);
assert!(second_solve_saw_retained_slot.get());
assert_eq!(
batch.ifds_scratch().frontier_word_capacity(),
frontier_capacity
);
assert_eq!(
batch.ifds_scratch().frontier_byte_capacity(),
frontier_byte_capacity
);
assert_eq!(batch.ifds_scratch().next_word_capacity(), next_capacity);
assert_eq!(batch.ifds_scratch().output_slot_count(), output_slots);
assert_eq!(output_slots, 1);
}
#[test]
fn fixed_point_batch_prepared_many_ifds_reuses_one_solve_scratch() {
let solve_phase = Cell::new(0usize);
let retained_slot_seen = Cell::new(0usize);
let dispatch = |_: &vyre::ir::Program,
inputs: &[&[u8]],
grid: Option<[u32; 3]>,
outputs: &mut Vec<Vec<u8>>| {
if inputs.len() < 6 {
return crate::fixed_point_closure::fallback_bitset_equal_dispatch(inputs, outputs);
}
match inputs.len() {
17 => {
assert_eq!(grid, Some([1, 1, 1]));
outputs.clear();
outputs.resize_with(4, Vec::new);
for word in [0u32, 1, 1] {
outputs[0].extend_from_slice(&word.to_le_bytes());
}
for word in [0u32, 0] {
outputs[1].extend_from_slice(&word.to_le_bytes());
}
outputs[2].extend_from_slice(&1u32.to_le_bytes());
outputs[3].extend_from_slice(&1u32.to_le_bytes());
}
7 => {
assert_eq!(grid, Some([2, 1, 1]));
if outputs.is_empty() {
outputs.push(Vec::with_capacity(8));
} else {
retained_slot_seen.set(retained_slot_seen.get() + 1);
}
let frontier = u32::from_le_bytes(inputs[5][0..4].try_into().unwrap());
let next = match solve_phase.get() {
0 => {
assert_eq!(frontier, 0b01);
solve_phase.set(1);
0b11u32
}
1 => {
assert_eq!(frontier, 0b11);
solve_phase.set(2);
0b11u32
}
2 => {
assert_eq!(frontier, 0b10);
solve_phase.set(3);
0b11u32
}
3 => {
assert_eq!(frontier, 0b11);
solve_phase.set(4);
0b11u32
}
call => panic!("unexpected prepared-many IFDS dispatch {call}"),
};
outputs[0].clear();
outputs[0].extend_from_slice(&next.to_le_bytes());
}
len => panic!("unexpected IFDS input count {len}"),
}
Ok(())
};
let mut batch = FixedPointBatch::new(&dispatch);
let prepared = batch
.prepare_ifds(1, 2, 1, &[(0, 0, 1)], &[], &[], &[])
.expect("batch must prepare IFDS CSR once");
let seed_a: &[(u32, u32, u32)] = &[(0, 0, 0)];
let seed_b: &[(u32, u32, u32)] = &[(0, 1, 0)];
let results = batch
.ifds_prepared_many(&prepared, &[seed_a, seed_b], 4)
.expect("prepared-many IFDS batch must converge");
assert_eq!(results, vec![vec![0, 1024], vec![0, 1024]]);
assert!(retained_slot_seen.get() >= 3);
assert_eq!(batch.ifds_scratch().output_slot_count(), 1);
}