rtx-backend-metal: GPU index_select / index_add via one-hot CSR SpMM
Override the Backend trait's host-round-trip defaults: gather is S @ X with S the [E x N] one-hot selection CSR; scatter-add is the adjoint S^T @ X, whose CSR is built directly by counting sort so duplicate indices land in one row and the spmm kernel (one thread per output element) accumulates them without atomics. CSR matrices are cached per thread keyed by the exact index list + dims, so a static graph topology (GNN message passing) builds each matrix once. Host fallback on degenerate shapes or any sparse-pipeline failure. 13 new parity tests vs CPU reference: duplicates, unreferenced rows, D=1/2/3, 15k x 5k x 64 gather/scatter, cache reuse, adjoint roundtrip. Verified on-device that the SpMM path (not the fallback) serves all 13. Co-Authored-By: Claude Fable 5 <[email protected]>
This commit is contained in:
quantum
co-authored by
Claude Fable 5
Claude Fable 5
parent
2e23d0f4c6
commit
9297976929
@@ -0,0 +1,282 @@
|
||||
//! Index operations (`index_select` / `index_add`) via sparse one-hot SpMM.
|
||||
//!
|
||||
//! Gathering rows `out[i] = x[indices[i]]` is exactly `S @ X` where `S` is the
|
||||
//! `[E x N]` one-hot selection matrix with `S[i, indices[i]] = 1`. Scatter-add
|
||||
//! (`index_add`) is the adjoint, `S^T @ X`, whose CSR form is built directly
|
||||
//! with a counting sort (so the GPU kernel — one thread per output element,
|
||||
//! looping a row's nonzeros — needs no atomics even with duplicate indices).
|
||||
//!
|
||||
//! CSR matrices are cached per thread, keyed by the exact index list and
|
||||
//! matrix dimensions, so repeated calls with a static graph topology (the GNN
|
||||
//! message-passing case) rebuild nothing.
|
||||
|
||||
use crate::MetalTensorPrimitive;
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
use rtx_metal::{MetalBuffer, MetalBufferUsage};
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
use rtx_metal::sparse::{spmm_csr, CsrMatrix};
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
use std::cell::RefCell;
|
||||
#[cfg(target_os = "macos")]
|
||||
use std::collections::HashMap;
|
||||
#[cfg(target_os = "macos")]
|
||||
use std::rc::Rc;
|
||||
|
||||
/// Maximum number of cached CSR matrices per thread before the cache is
|
||||
/// cleared. Each entry holds `O(nnz)` GPU memory, so keep this small.
|
||||
#[cfg(target_os = "macos")]
|
||||
const CSR_CACHE_CAP: usize = 32;
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
#[derive(PartialEq, Eq, Hash, Clone)]
|
||||
struct CsrCacheKey {
|
||||
/// False: selection matrix `S` (E x N). True: adjoint `S^T` (num_rows x E).
|
||||
transposed: bool,
|
||||
/// The dense dimension (N for select, num_rows for add).
|
||||
dim: usize,
|
||||
/// The exact index list (collision-proof; hashing is cheap relative to
|
||||
/// the SpMM itself and only rebuilt entries pay the CSR construction).
|
||||
indices: Vec<usize>,
|
||||
}
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
thread_local! {
|
||||
static CSR_CACHE: RefCell<HashMap<CsrCacheKey, Rc<CsrMatrix<f32>>>> =
|
||||
RefCell::new(HashMap::new());
|
||||
}
|
||||
|
||||
/// Get or build the `[E x N]` one-hot selection CSR for `indices`.
|
||||
#[cfg(target_os = "macos")]
|
||||
fn cached_select_csr(
|
||||
device: &rtx_metal::MetalDevice,
|
||||
indices: &[usize],
|
||||
num_src_rows: usize,
|
||||
) -> Option<Rc<CsrMatrix<f32>>> {
|
||||
let key = CsrCacheKey {
|
||||
transposed: false,
|
||||
dim: num_src_rows,
|
||||
indices: indices.to_vec(),
|
||||
};
|
||||
CSR_CACHE.with(|cache| {
|
||||
let mut cache = cache.borrow_mut();
|
||||
if let Some(csr) = cache.get(&key) {
|
||||
return Some(csr.clone());
|
||||
}
|
||||
let e = indices.len();
|
||||
let row_ptr: Vec<i32> = (0..=e as i32).collect();
|
||||
let col_indices: Vec<i32> = indices.iter().map(|&i| i as i32).collect();
|
||||
let values = vec![1.0f32; e];
|
||||
let csr =
|
||||
CsrMatrix::new(device, e, num_src_rows, &row_ptr, &col_indices, &values).ok()?;
|
||||
let csr = Rc::new(csr);
|
||||
if cache.len() >= CSR_CACHE_CAP {
|
||||
cache.clear();
|
||||
}
|
||||
cache.insert(key, csr.clone());
|
||||
Some(csr)
|
||||
})
|
||||
}
|
||||
|
||||
/// Get or build the `[num_rows x E]` transposed one-hot CSR (counting sort).
|
||||
#[cfg(target_os = "macos")]
|
||||
fn cached_scatter_csr(
|
||||
device: &rtx_metal::MetalDevice,
|
||||
indices: &[usize],
|
||||
num_rows: usize,
|
||||
) -> Option<Rc<CsrMatrix<f32>>> {
|
||||
let key = CsrCacheKey {
|
||||
transposed: true,
|
||||
dim: num_rows,
|
||||
indices: indices.to_vec(),
|
||||
};
|
||||
CSR_CACHE.with(|cache| {
|
||||
let mut cache = cache.borrow_mut();
|
||||
if let Some(csr) = cache.get(&key) {
|
||||
return Some(csr.clone());
|
||||
}
|
||||
let e = indices.len();
|
||||
// Counting sort: row r of S^T holds the input positions i with
|
||||
// indices[i] == r.
|
||||
let mut row_ptr = vec![0i32; num_rows + 1];
|
||||
for &r in indices {
|
||||
row_ptr[r + 1] += 1;
|
||||
}
|
||||
for r in 0..num_rows {
|
||||
row_ptr[r + 1] += row_ptr[r];
|
||||
}
|
||||
let mut next: Vec<i32> = row_ptr[..num_rows].to_vec();
|
||||
let mut col_indices = vec![0i32; e];
|
||||
for (i, &r) in indices.iter().enumerate() {
|
||||
col_indices[next[r] as usize] = i as i32;
|
||||
next[r] += 1;
|
||||
}
|
||||
let values = vec![1.0f32; e];
|
||||
let csr = CsrMatrix::new(device, num_rows, e, &row_ptr, &col_indices, &values).ok()?;
|
||||
let csr = Rc::new(csr);
|
||||
if cache.len() >= CSR_CACHE_CAP {
|
||||
cache.clear();
|
||||
}
|
||||
cache.insert(key, csr.clone());
|
||||
Some(csr)
|
||||
})
|
||||
}
|
||||
|
||||
/// Host fallback mirroring the `Backend::index_select` default body.
|
||||
fn host_index_select<const D: usize>(
|
||||
tensor: &MetalTensorPrimitive<D>,
|
||||
indices: &[usize],
|
||||
row_len: usize,
|
||||
out_shape: [usize; D],
|
||||
) -> MetalTensorPrimitive<D> {
|
||||
let src = tensor.to_vec();
|
||||
let mut out = Vec::with_capacity(indices.len() * row_len);
|
||||
for &idx in indices {
|
||||
out.extend_from_slice(&src[idx * row_len..(idx + 1) * row_len]);
|
||||
}
|
||||
super::creation::from_data(&out, out_shape, &tensor.device)
|
||||
}
|
||||
|
||||
/// Host fallback mirroring the `Backend::index_add` default body.
|
||||
fn host_index_add<const D: usize>(
|
||||
tensor: &MetalTensorPrimitive<D>,
|
||||
indices: &[usize],
|
||||
num_rows: usize,
|
||||
row_len: usize,
|
||||
out_shape: [usize; D],
|
||||
) -> MetalTensorPrimitive<D> {
|
||||
let src = tensor.to_vec();
|
||||
let mut out = vec![0.0f32; num_rows * row_len];
|
||||
for (i, &idx) in indices.iter().enumerate() {
|
||||
let dst = &mut out[idx * row_len..(idx + 1) * row_len];
|
||||
let row = &src[i * row_len..(i + 1) * row_len];
|
||||
for (d, &s) in dst.iter_mut().zip(row) {
|
||||
*d += s;
|
||||
}
|
||||
}
|
||||
super::creation::from_data(&out, out_shape, &tensor.device)
|
||||
}
|
||||
|
||||
/// Gather rows along dim 0: `out[i, ..] = tensor[indices[i], ..]`.
|
||||
///
|
||||
/// GPU path: one-hot CSR `[E x N]` times the dense tensor via `spmm_csr`.
|
||||
/// Falls back to the host implementation for degenerate shapes or any
|
||||
/// sparse-pipeline failure.
|
||||
///
|
||||
/// # Panics
|
||||
/// Panics if any index is `>= shape[0]`, if `D == 0`, or if the tensor is
|
||||
/// not contiguous.
|
||||
pub fn index_select<const D: usize>(
|
||||
tensor: &MetalTensorPrimitive<D>,
|
||||
indices: &[usize],
|
||||
) -> MetalTensorPrimitive<D> {
|
||||
assert!(D >= 1, "index_select requires at least one dimension");
|
||||
assert!(
|
||||
tensor.is_contiguous(),
|
||||
"index_select: tensor must be contiguous"
|
||||
);
|
||||
let shape = tensor.shape;
|
||||
let num_rows = shape[0];
|
||||
let row_len: usize = shape[1..].iter().product();
|
||||
for &idx in indices {
|
||||
assert!(
|
||||
idx < num_rows,
|
||||
"index_select: index {idx} out of range for {num_rows} rows"
|
||||
);
|
||||
}
|
||||
let mut out_shape = shape;
|
||||
out_shape[0] = indices.len();
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
if row_len > 0
|
||||
&& !indices.is_empty()
|
||||
&& num_rows <= i32::MAX as usize
|
||||
&& indices.len() < i32::MAX as usize
|
||||
{
|
||||
let device = tensor.device.metal_device();
|
||||
if let Some(csr) = cached_select_csr(device, indices, num_rows) {
|
||||
let out_numel = indices.len() * row_len;
|
||||
if let Ok(mut output) =
|
||||
MetalBuffer::new(device, out_numel, MetalBufferUsage::Shared)
|
||||
{
|
||||
if spmm_csr(device, &csr, tensor.data(), &mut output, row_len).is_ok() {
|
||||
return MetalTensorPrimitive::new(
|
||||
output,
|
||||
out_shape,
|
||||
tensor.device.clone(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
host_index_select(tensor, indices, row_len, out_shape)
|
||||
}
|
||||
|
||||
/// Scatter-add rows along dim 0 into a zero tensor with `num_rows` rows:
|
||||
/// `out[indices[i], ..] += tensor[i, ..]`. Adjoint of [`index_select`].
|
||||
///
|
||||
/// GPU path: transposed one-hot CSR `[num_rows x E]` (built by counting
|
||||
/// sort — duplicate indices land in the same CSR row, so the kernel
|
||||
/// accumulates them without atomics) times the dense tensor via `spmm_csr`.
|
||||
///
|
||||
/// # Panics
|
||||
/// Panics if `indices.len() != shape[0]`, if any index is `>= num_rows`,
|
||||
/// if `D == 0`, or if the tensor is not contiguous.
|
||||
pub fn index_add<const D: usize>(
|
||||
tensor: &MetalTensorPrimitive<D>,
|
||||
indices: &[usize],
|
||||
num_rows: usize,
|
||||
) -> MetalTensorPrimitive<D> {
|
||||
assert!(D >= 1, "index_add requires at least one dimension");
|
||||
assert!(
|
||||
tensor.is_contiguous(),
|
||||
"index_add: tensor must be contiguous"
|
||||
);
|
||||
let shape = tensor.shape;
|
||||
assert_eq!(
|
||||
indices.len(),
|
||||
shape[0],
|
||||
"index_add: indices.len() must equal the number of input rows"
|
||||
);
|
||||
let row_len: usize = shape[1..].iter().product();
|
||||
for &idx in indices {
|
||||
assert!(
|
||||
idx < num_rows,
|
||||
"index_add: index {idx} out of range for {num_rows} rows"
|
||||
);
|
||||
}
|
||||
let mut out_shape = shape;
|
||||
out_shape[0] = num_rows;
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
{
|
||||
if row_len > 0
|
||||
&& num_rows > 0
|
||||
&& indices.len() < i32::MAX as usize
|
||||
&& num_rows <= i32::MAX as usize
|
||||
{
|
||||
let device = tensor.device.metal_device();
|
||||
if let Some(csr) = cached_scatter_csr(device, indices, num_rows) {
|
||||
let out_numel = num_rows * row_len;
|
||||
if let Ok(mut output) =
|
||||
MetalBuffer::new(device, out_numel, MetalBufferUsage::Shared)
|
||||
{
|
||||
if spmm_csr(device, &csr, tensor.data(), &mut output, row_len).is_ok() {
|
||||
return MetalTensorPrimitive::new(
|
||||
output,
|
||||
out_shape,
|
||||
tensor.device.clone(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
host_index_add(tensor, indices, num_rows, row_len, out_shape)
|
||||
}
|
||||
@@ -10,6 +10,7 @@ pub mod conv;
|
||||
pub mod creation;
|
||||
pub mod device;
|
||||
pub mod gemm;
|
||||
pub mod index;
|
||||
pub mod normalization;
|
||||
pub mod reduction;
|
||||
pub mod shape;
|
||||
|
||||
Reference in New Issue
Block a user