crates/mimir-graph/src/csr.rs
//! CSR import: `(offsets, cols[, vals])` → [`Matrix`].
//!
//! GraphBLAS can take a CSR *without copying* (`GxB_Matrix_pack_CSR`), but it
//! takes ownership of the arrays and later frees them with its own `free`, so
//! they must come from `malloc` — a borrowed `&[u64]` (or an Arrow buffer)
//! cannot be handed over. [`Matrix::from_csr`] therefore does exactly one
//! `memcpy` into malloc'd arrays and packs them; the Arrow path
//! ([`crate::arrow`]) writes its offsets straight into a malloc'd array and
//! copies the column buffer once. [`Matrix::from_csr_import`] is the standard
//! `GrB_Matrix_import` (GraphBLAS copies internally) — kept for the bench.
//!
//! `pack` trusts its input (`GxB_FAST_IMPORT`), so everything is validated
//! here first: monotone offsets, `offsets[n] == cols.len()`, every column
//! `< ncols`, no duplicate column within a row.
use std::mem::size_of;
use std::ptr;
use graphblas_sys as sys;
use crate::error::{Error, Result, grb};
use crate::matrix::Matrix;
use crate::types::Scalar;
/// A `malloc`'d array GraphBLAS may take ownership of. Freed on drop unless
/// [`MallocBuf::into_raw`] handed it over.
pub(crate) struct MallocBuf<T: Copy> {
ptr: *mut T,
len: usize,
}
impl<T: Copy> MallocBuf<T> {
pub(crate) fn new(len: usize) -> Result<Self> {
// malloc(0) may return NULL; GraphBLAS wants a real allocation.
let bytes = len.max(1) * size_of::<T>();
let ptr = unsafe { libc::malloc(bytes) } as *mut T;
if ptr.is_null() {
return Err(Error::Alloc { bytes });
}
Ok(Self { ptr, len })
}
pub(crate) fn zeroed(len: usize) -> Result<Self> {
let bytes = len.max(1) * size_of::<T>();
let ptr = unsafe { libc::calloc(1, bytes) } as *mut T;
if ptr.is_null() {
return Err(Error::Alloc { bytes });
}
Ok(Self { ptr, len })
}
pub(crate) fn from_slice(s: &[T]) -> Result<Self> {
let b = Self::new(s.len())?;
unsafe { ptr::copy_nonoverlapping(s.as_ptr(), b.ptr, s.len()) };
Ok(b)
}
pub(crate) fn as_mut_slice(&mut self) -> &mut [T] {
unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
}
pub(crate) fn as_slice(&self) -> &[T] {
unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
}
/// Size in bytes as GraphBLAS wants it (at least one element).
pub(crate) fn bytes(&self) -> u64 {
(self.len.max(1) * size_of::<T>()) as u64
}
}
impl<T: Copy> Drop for MallocBuf<T> {
fn drop(&mut self) {
if !self.ptr.is_null() {
unsafe { libc::free(self.ptr as *mut libc::c_void) };
}
}
}
/// What a validation pass learnt about a CSR.
pub(crate) struct CsrShape {
/// Every row's columns strictly ascending.
pub(crate) sorted_rows: bool,
}
/// Check a CSR in O(n + nnz): ascending rows are checked by one comparison
/// per entry; only jumbled rows pay for the duplicate marker array.
pub(crate) fn validate(nrows: u64, ncols: u64, offsets: &[u64], cols: &[u64]) -> Result<CsrShape> {
if offsets.len() as u64 != nrows + 1 {
return Err(Error::invalid(format!(
"offsets has {} entries, want nrows+1 = {}",
offsets.len(),
nrows + 1
)));
}
if offsets[0] != 0 {
return Err(Error::invalid(format!(
"offsets[0] = {}, want 0",
offsets[0]
)));
}
if offsets[nrows as usize] != cols.len() as u64 {
return Err(Error::invalid(format!(
"offsets[nrows] = {} but cols has {}",
offsets[nrows as usize],
cols.len()
)));
}
// Fast tier: a strictly ascending row has no duplicates and needs no
// marker array (what a sorted edge scan produces).
let mut sorted_rows = true;
let mut jumbled: Vec<usize> = Vec::new();
for r in 0..nrows as usize {
let (a, b) = (offsets[r], offsets[r + 1]);
if b < a {
return Err(Error::invalid(format!(
"offsets not monotone at row {r}: {a} > {b}"
)));
}
let row = &cols[a as usize..b as usize];
if row.windows(2).all(|w| w[0] < w[1]) {
// ascending: the last column is the largest
if let Some(&last) = row.last()
&& last >= ncols
{
return Err(Error::invalid(format!(
"row {r}: column {last} >= ncols {ncols}"
)));
}
} else {
sorted_rows = false;
jumbled.push(r);
}
}
// Slow tier, jumbled rows only: per-column "last row seen" finds
// duplicates without sorting.
if !jumbled.is_empty() {
let mut mark = vec![u64::MAX; ncols as usize];
for r in jumbled {
for &c in &cols[offsets[r] as usize..offsets[r + 1] as usize] {
if c >= ncols {
return Err(Error::invalid(format!(
"row {r}: column {c} >= ncols {ncols}"
)));
}
if mark[c as usize] == r as u64 {
return Err(Error::invalid(format!("row {r}: duplicate column {c}")));
}
mark[c as usize] = r as u64;
}
}
}
Ok(CsrShape { sorted_rows })
}
/// Hand malloc'd CSR arrays to a fresh matrix (`GxB_Matrix_pack_CSR`, no copy).
/// `vals = None` packs an iso matrix whose every entry is `T::ONE`.
pub(crate) fn pack<T: Scalar>(
nrows: u64,
ncols: u64,
offsets: MallocBuf<u64>,
cols: MallocBuf<u64>,
vals: Option<MallocBuf<T>>,
sorted_rows: bool,
) -> Result<Matrix<T>> {
let m = Matrix::<T>::new(nrows, ncols)?;
let (iso, vals) = match vals {
Some(v) => (false, v),
None => {
let mut one = MallocBuf::<T>::new(1)?;
one.as_mut_slice()[0] = T::ONE;
(true, one)
}
};
let (ap_size, aj_size, ax_size) = (offsets.bytes(), cols.bytes(), vals.bytes());
// pack takes the pointers by reference and NULLs them on success; on
// failure they stay ours and the MallocBufs free them.
let mut ap = offsets.ptr;
let mut aj = cols.ptr;
let mut ax = vals.ptr as *mut libc::c_void;
let r = grb!(GxB_Matrix_pack_CSR(
m.raw,
&mut ap,
&mut aj,
&mut ax,
ap_size,
aj_size,
ax_size,
iso,
!sorted_rows,
ptr::null_mut()
));
// Whatever GraphBLAS took (pointer now NULL) must not be freed by us.
let mut offsets = offsets;
let mut cols = cols;
let mut vals = vals;
offsets.ptr = ap;
cols.ptr = aj;
vals.ptr = ax as *mut T;
r?;
Ok(m)
}
/// Iso build from coordinate arrays (`GxB_Matrix_build_Scalar`): every entry
/// `T::ONE`, duplicates collapse, any order, bounds checked by GraphBLAS.
///
/// The pointers must be valid for `n` reads each.
pub(crate) fn build_iso<T: Scalar>(
m: &Matrix<T>,
rows: *const u64,
cols: *const u64,
n: u64,
) -> Result<()> {
let mut s: sys::GrB_Scalar = ptr::null_mut();
grb!(GrB_Scalar_new(&mut s, T::grb_type()))?;
// setElement typecasts: 1 → true / 1 / 1.0 = T::ONE.
let r = grb!(GrB_Scalar_setElement_UINT64(s, 1))
.and_then(|_| grb!(GxB_Matrix_build_Scalar(m.raw, rows, cols, s, n)));
unsafe { sys::GrB_Scalar_free(&mut s) };
r
}
impl<T: Scalar> Matrix<T> {
/// Import an `nrows × ncols` CSR. `vals = None` → pattern (iso) matrix,
/// every entry `T::ONE`. One copy into malloc'd arrays, then a zero-copy
/// `GxB_Matrix_pack_CSR`; the input is validated first (see module docs).
pub fn from_csr(
nrows: u64,
ncols: u64,
offsets: &[u64],
cols: &[u64],
vals: Option<&[T]>,
) -> Result<Self> {
if let Some(v) = vals
&& v.len() != cols.len()
{
return Err(Error::invalid(format!(
"vals has {} entries, cols has {}",
v.len(),
cols.len()
)));
}
let shape = validate(nrows, ncols, offsets, cols)?;
let vals = vals.map(MallocBuf::from_slice).transpose()?;
pack(
nrows,
ncols,
MallocBuf::from_slice(offsets)?,
MallocBuf::from_slice(cols)?,
vals,
shape.sorted_rows,
)
}
/// Same input contract as [`Matrix::from_csr`], through the standard
/// `GrB_Matrix_import(.., GrB_CSR_FORMAT)`, which copies inside GraphBLAS
/// (and needs a full values array even for a pattern matrix).
pub fn from_csr_import(
nrows: u64,
ncols: u64,
offsets: &[u64],
cols: &[u64],
vals: Option<&[T]>,
) -> Result<Self> {
validate(nrows, ncols, offsets, cols)?;
let ones;
let vals = match vals {
Some(v) if v.len() == cols.len() => v,
Some(v) => {
return Err(Error::invalid(format!(
"vals has {} entries, cols has {}",
v.len(),
cols.len()
)));
}
None => {
ones = vec![T::ONE; cols.len()];
&ones[..]
}
};
crate::runtime::init()?;
let mut raw = ptr::null_mut();
let info = unsafe {
T::matrix_import(
&mut raw,
nrows,
ncols,
offsets.as_ptr(),
cols.as_ptr(),
vals.as_ptr(),
offsets.len() as u64,
cols.len() as u64,
vals.len() as u64,
)
};
crate::error::check("GrB_Matrix_import", info)?;
Ok(Matrix::from_raw(raw))
}
}