Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 48 additions & 26 deletions native/lgj-abi/src/exports.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,8 @@ use std::cell::RefCell;
use std::panic::{catch_unwind, AssertUnwindSafe};

use lance_graph_mask_risc::{
execute, reference_execute, reference_scratch, scratch_words_for, ExecError, LaneRef, Planes,
Scratch, Value,
execute_into, reference_execute_into, scratch_words_for, tile_words_for, ExecError, Foreign,
LaneRef, Out, Planes, Scratch, Value,
};

use crate::abi::*;
Expand Down Expand Up @@ -1733,8 +1733,15 @@ fn exec_error_to_status(e: ExecError) -> i32 {
match e {
ExecError::LaneOutOfRange(_) => LGJ_ERR_INVALID_LANE,
ExecError::LaneKind { .. } => LGJ_ERR_LANE_KIND_MISMATCH,
// The lane is the right width but not in key order: a lane-shape
// precondition of the fold, reported as the same class of mismatch.
ExecError::LaneNotOrdered { .. } => LGJ_ERR_LANE_KIND_MISMATCH,
ExecError::LenMismatch { .. } => LGJ_ERR_MASK_LENGTH_MISMATCH,
ExecError::PlaneOutOfRange(_) | ExecError::PlaneTail(_) => LGJ_ERR_INVALID_HANDLE,
ExecError::PlaneOutOfRange(_)
| ExecError::PlaneTail(_)
| ExecError::ForeignOutOfRange(_) => LGJ_ERR_INVALID_HANDLE,
ExecError::ForeignLaneOutOfRange(_) => LGJ_ERR_INVALID_LANE,
ExecError::ForeignLaneKind { .. } => LGJ_ERR_LANE_KIND_MISMATCH,
ExecError::SumRowBound { .. } => LGJ_ERR_SUM_OVERFLOW,
ExecError::ScratchTooSmall { .. }
| ExecError::ScratchBufferTooSmall { .. }
Expand All @@ -1743,6 +1750,7 @@ fn exec_error_to_status(e: ExecError) -> i32 {
| ExecError::ScratchReadBeforeWrite { .. }
| ExecError::ScratchWords { .. }
| ExecError::BlendNeedsOut
| ExecError::TerminalNeedsOut { .. }
| ExecError::GateAliasesDst { .. }
| ExecError::RangeOutOfBounds { .. } => LGJ_ERR_ALLOCATION_FAILED,
}
Expand Down Expand Up @@ -1864,33 +1872,53 @@ fn plan_eval_impl(

match path {
Path::Simd => {
let need = match scratch_words_for(n_words, plan_lower::SLOTS as usize) {
// TILED: the scratch is `SLOTS × tile_words_for(rows)` words
// however many rows the resource holds, and the kept mask is
// written tile by tile into the destination resource's own words
// — the demanded sink, and the only population-sized write.
// Validation is total before the first tile, so an error leaves
// `dst_mask` byte-for-byte as it was, exactly as before.
let tile = tile_words_for(rows);
let need = match scratch_words_for(tile, plan_lower::SLOTS as usize) {
Some(n) => n,
None => return LGJ_ERR_LENGTH_OVERFLOW,
};
let mut g = match mask.write_mask() {
Some(g) => g,
None => return LGJ_ERR_WRONG_RESOURCE_KIND,
};
if g.words.len() != n_words {
return LGJ_ERR_MASK_LENGTH_MISMATCH;
}
PLAN_SCRATCH.with(|cell| {
let mut buf = cell.borrow_mut();
if buf.len() < need {
buf.resize(need, 0);
}
let mut scratch =
match Scratch::over(&mut buf[..need], n_words, plan_lower::SLOTS as usize) {
match Scratch::over(&mut buf[..need], tile, plan_lower::SLOTS as usize) {
Ok(s) => s,
Err(e) => return exec_error_to_status(e),
};
let count = match execute(&program, &planes, &mut scratch, None) {
Ok(Value::Count(c)) => c as u64,
match execute_into(
&program,
&planes,
&Foreign::NONE,
&mut scratch,
Out::Mask(&mut g.words),
) {
Ok(Value::Mask(_)) => {}
// The lowering emits exactly one terminal and it is
// `Count`; any other value means this file built a
// `Keep`; any other value means this file built a
// program it did not intend to.
Ok(_) => return LGJ_ERR_ALLOCATION_FAILED,
Err(e) => return exec_error_to_status(e),
};
let words = match scratch.slot(plan_lower::ACC_SLOT) {
Some(w) => w,
None => return LGJ_ERR_ALLOCATION_FAILED,
};
publish(&mask, words, count, out_count)
}
let count = kernels::simd_popcount(&g.words);
drop(g);
// SAFETY: non-null (checked by the caller); written only on success.
unsafe { *out_count = count };
LGJ_OK
})
}
// Fork A: the scalar symbol runs mask-risc's row-at-a-time oracle, so
Expand All @@ -1900,20 +1928,14 @@ fn plan_eval_impl(
// bit-packing bug — which is why the allocation gate names
// `lgj_plan_eval` and only it.
Path::Scalar => {
let count = match reference_execute(&program, &planes, None) {
Ok(Value::Count(c)) => c as u64,
let mut kept = vec![0u64; n_words];
match reference_execute_into(&program, &planes, &Foreign::NONE, Out::Mask(&mut kept)) {
Ok(Value::Mask(_)) => {}
Ok(_) => return LGJ_ERR_ALLOCATION_FAILED,
Err(e) => return exec_error_to_status(e),
};
let slots = match reference_scratch(&program, &planes) {
Ok(s) => s,
Err(e) => return exec_error_to_status(e),
};
let words = match slots.get(plan_lower::ACC_SLOT as usize) {
Some(w) => w.as_slice(),
None => return LGJ_ERR_ALLOCATION_FAILED,
};
publish(&mask, words, count, out_count)
}
let count = kernels::simd_popcount(&kept);
publish(&mask, &kept, count, out_count)
}
}
}
Expand Down
20 changes: 15 additions & 5 deletions native/lgj-abi/src/exports/tests/lowering_convergence.rs
Original file line number Diff line number Diff line change
Expand Up @@ -156,12 +156,22 @@ fn cmp_of(o: &LgjOpDesc) -> Cmp {
/// function on both sides is the whole point of comparing LOWERINGS rather
/// than comparing EXECUTIONS.
fn run(p: &Program, planes: &Planes) -> usize {
let words = planes.n_rows.div_ceil(64);
let slots = p.scratch_slots as usize;
let mut buf = vec![0u64; scratch_words_for(words, slots).expect("sized")];
let mut scratch = Scratch::over(&mut buf, words, slots).expect("carves");
match execute(p, planes, &mut scratch, None).expect("runs") {
// The tiled default scratch. This crate's own lowering ends in `Keep`
// (the destination mask is the demanded sink; the count is its
// popcount), the quack lowering in `Count` — both are read here.
let mut scratch = Scratch::for_program(p, planes.n_rows).expect("addressable");
let mut kept = vec![0u64; planes.n_rows.div_ceil(64)];
match execute_into(
p,
planes,
&Foreign::NONE,
&mut scratch,
Out::Mask(&mut kept),
)
.expect("runs")
{
Value::Count(c) => c,
Value::Mask(_) => ndarray::simd::popcount_batch_u64(&kept) as usize,
other => panic!("not a count: {other:?}"),
}
}
Expand Down
6 changes: 5 additions & 1 deletion native/lgj-abi/src/plan_lower.rs
Original file line number Diff line number Diff line change
Expand Up @@ -96,9 +96,13 @@ pub(crate) fn lower_plan(ops: &[LgjOpDesc]) -> Option<Lowered> {
});
}

// `Keep`, not `Count`: the destination mask IS the demanded result, so
// the executor writes it tile by tile straight into the resource's own
// words (`Out::Mask`) and the count is a popcount over that sink. Under
// tiled execution no population-sized scratch exists to count from.
Some(Lowered::Program(Program::new(
program_ops,
Terminal::Count {
Terminal::Keep {
mask: Operand::Scratch(ACC_SLOT),
},
)))
Expand Down
Loading