crates/mimir-graph/src/matrix.rs
use std::marker::PhantomData;
use std::ptr;
use graphblas_sys as sys;
use crate::error::{Error, Result, check_value, grb};
use crate::runtime;
use crate::types::Scalar;
/// An owned `GrB_Matrix` of element type `T` (freed on drop).
pub struct Matrix<T: Scalar> {
pub(crate) raw: sys::GrB_Matrix,
_t: PhantomData<T>,
}
// A GrB_Matrix may move between threads; concurrent use from two threads is
// not allowed (pending work), hence no `Sync`.
unsafe impl<T: Scalar> Send for Matrix<T> {}
impl<T: Scalar> Drop for Matrix<T> {
fn drop(&mut self) {
if !self.raw.is_null() {
unsafe { sys::GrB_Matrix_free(&mut self.raw) };
}
}
}
impl<T: Scalar> Matrix<T> {
/// Empty `nrows × ncols` matrix.
pub fn new(nrows: u64, ncols: u64) -> Result<Self> {
runtime::init()?;
let mut raw = ptr::null_mut();
grb!(GrB_Matrix_new(&mut raw, T::grb_type(), nrows, ncols))?;
Ok(Self::from_raw(raw))
}
pub(crate) fn from_raw(raw: sys::GrB_Matrix) -> Self {
Self {
raw,
_t: PhantomData,
}
}
/// Build from coordinate tuples. Duplicates are combined (OR for bool,
/// sum otherwise). Indices are bounds-checked by GraphBLAS.
pub fn from_tuples(
nrows: u64,
ncols: u64,
rows: &[u64],
cols: &[u64],
vals: &[T],
) -> Result<Self> {
if rows.len() != cols.len() || rows.len() != vals.len() {
return Err(Error::invalid(format!(
"tuple arrays differ in length: rows {}, cols {}, vals {}",
rows.len(),
cols.len(),
vals.len()
)));
}
let m = Self::new(nrows, ncols)?;
let info = unsafe {
T::matrix_build(
m.raw,
rows.as_ptr(),
cols.as_ptr(),
vals.as_ptr(),
rows.len() as u64,
)
};
crate::error::check("GrB_Matrix_build", info)?;
Ok(m)
}
/// Build a pattern matrix (every entry = `T::ONE`) from `(row, col)` edges.
/// Duplicate edges collapse to one entry.
pub fn from_edges(nrows: u64, ncols: u64, rows: &[u64], cols: &[u64]) -> Result<Self> {
if rows.len() != cols.len() {
return Err(Error::invalid(format!(
"rows {} != cols {}",
rows.len(),
cols.len()
)));
}
let m = Self::new(nrows, ncols)?;
crate::csr::build_iso::<T>(&m, rows.as_ptr(), cols.as_ptr(), rows.len() as u64)?;
Ok(m)
}
pub fn nrows(&self) -> Result<u64> {
let mut n = 0;
grb!(GrB_Matrix_nrows(&mut n, self.raw))?;
Ok(n)
}
pub fn ncols(&self) -> Result<u64> {
let mut n = 0;
grb!(GrB_Matrix_ncols(&mut n, self.raw))?;
Ok(n)
}
/// Number of stored entries (finishes pending work).
pub fn nvals(&self) -> Result<u64> {
let mut n = 0;
grb!(GrB_Matrix_nvals(&mut n, self.raw))?;
Ok(n)
}
/// Entry `(i, j)`, `None` if not stored.
pub fn get(&self, i: u64, j: u64) -> Result<Option<T>> {
let mut x = T::default();
let present = check_value("GrB_Matrix_extractElement", unsafe {
T::matrix_extract_element(&mut x, self.raw, i, j)
})?;
Ok(present.then_some(x))
}
/// All entries as `(rows, cols, vals)`, row-major order not guaranteed.
pub fn to_tuples(&self) -> Result<(Vec<u64>, Vec<u64>, Vec<T>)> {
let n = self.nvals()?;
let mut i = vec![0u64; n as usize];
let mut j = vec![0u64; n as usize];
let mut x = vec![T::default(); n as usize];
let mut got = n;
let info = unsafe {
T::matrix_extract_tuples(
i.as_mut_ptr(),
j.as_mut_ptr(),
x.as_mut_ptr(),
&mut got,
self.raw,
)
};
crate::error::check("GrB_Matrix_extractTuples", info)?;
i.truncate(got as usize);
j.truncate(got as usize);
x.truncate(got as usize);
Ok((i, j, x))
}
/// Sorted `(row, col)` pairs of the stored entries.
pub fn edges(&self) -> Result<Vec<(u64, u64)>> {
let (i, j, _) = self.to_tuples()?;
let mut e: Vec<(u64, u64)> = i.into_iter().zip(j).collect();
e.sort_unstable();
Ok(e)
}
/// Sorted `(row, col)` pairs of the stored entries in the given ROWS
/// only (`rows` sorted, distinct): `A(rows, :)` extracted, so the cost
/// is the selected rows' entries, not the whole matrix — what
/// [`Matrix::edges`] plus a filter on the row costs at room scale.
pub fn row_edges(&self, rows: &[u64]) -> Result<Vec<(u64, u64)>> {
if rows.is_empty() {
return Ok(Vec::new());
}
let ncols = self.ncols()?;
let c = Self::new(rows.len() as u64, ncols)?;
grb!(GrB_Matrix_extract(
c.raw,
ptr::null_mut(),
ptr::null_mut(),
self.raw,
rows.as_ptr(),
rows.len() as u64,
sys::GrB_ALL,
ncols,
ptr::null_mut()
))?;
let (i, j, _) = c.to_tuples()?;
let mut e: Vec<(u64, u64)> = i.into_iter().zip(j).map(|(i, j)| (rows[i as usize], j)).collect();
e.sort_unstable();
Ok(e)
}
/// `Aᵀ` as a new matrix (materialised; worth caching for repeated
/// reverse walks, see `ops::reverse_reach`).
pub fn transpose(&self) -> Result<Self> {
let m = Self::new(self.ncols()?, self.nrows()?)?;
grb!(GrB_transpose(
m.raw,
ptr::null_mut(),
ptr::null_mut(),
self.raw,
ptr::null_mut()
))?;
Ok(m)
}
/// Deep copy.
pub fn dup(&self) -> Result<Self> {
let mut raw = ptr::null_mut();
grb!(GrB_Matrix_dup(&mut raw, self.raw))?;
Ok(Self::from_raw(raw))
}
/// Finish all pending work now (`GrB_Matrix_wait(.., GrB_MATERIALIZE)`).
pub fn wait(&self) -> Result<()> {
grb!(GrB_Matrix_wait(self.raw, sys::GrB_MATERIALIZE))
}
/// **Let GraphBLAS pick the cheapest format now.** A packed CSR stays CSR,
/// so a nearly empty n×n matrix keeps (n+1)·8 bytes of row pointers — tens
/// of MB per edge kind at room scale. Setting the sparsity control to AUTO
/// and materializing converts a mostly-empty matrix to hypersparse and
/// frees those pointers; a filled one stays as it is.
pub fn conform(&self) -> Result<()> {
grb!(GxB_Matrix_Option_set_INT32(
self.raw,
sys::GxB_SPARSITY_CONTROL,
sys::GxB_AUTO_SPARSITY as i32
))?;
self.wait()
}
/// Bytes this matrix holds (`GxB_Matrix_memoryUsage`).
pub fn memory_bytes(&self) -> Result<usize> {
let mut size = 0usize;
grb!(GxB_Matrix_memoryUsage(&mut size, self.raw))?;
Ok(size)
}
/// The raw handle, for calls this layer does not wrap. Ownership stays here.
pub fn as_raw(&self) -> sys::GrB_Matrix {
self.raw
}
}
impl<T: Scalar> std::fmt::Debug for Matrix<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Matrix")
.field("type", &std::any::type_name::<T>())
.field("nrows", &self.nrows().ok())
.field("ncols", &self.ncols().ok())
.field("nvals", &self.nvals().ok())
.finish()
}
}