embedded3 PERF-3 P1-5 (c) part 1: scatter kernels (e3_scat_{f64,f32,u32,i32}) in the CG module; the refresh's seven f64 coefficient arrays and z links, and the V-cycle fine level's f32 arrays and maps, scattered on the changed rows into persistent buffers (DeviceVcycle::refresh_fine_rows) — slab CSV byte-identical, band check passed, device moving/cg green; ny 124 rebuild block 1,895 → 1,748 ms per step. The first form redirected the V-cycle's own e3_cg_scatter_f64 launch (a replace-all on the launch name) and failed every gate — the symbol table is the first thing to read when adding to a module
CI / Build (macos-latest) (push) Waiting to run
CI / Test (macos-latest) (push) Blocked by required conditions
CI / Test (ubuntu-latest) (push) Blocked by required conditions
CI / Python Bindings (maturin) (macos-latest) (push) Blocked by required conditions
CI / Python Bindings (maturin) (ubuntu-latest) (push) Blocked by required conditions
CI / WASM Build + Size Check (push) Blocked by required conditions
CI / Distributed Training Tests (push) Blocked by required conditions
CI / CI Success (push) Blocked by required conditions
CI / Build CPU-Only (Explicit) (push) Failing after 4s
CI / Format Check (push) Failing after 4s
Documentation / Build User Guide (push) Successful in 4s
Documentation / Build API Documentation (push) Failing after 30s
CI / Clippy Check (push) Failing after 1m0s
CI / Build (ubuntu-latest) (push) Failing after 2m53s
Performance Benchmarks / Run Benchmarks (push) Successful in 3m27s

Co-Authored-By: Claude Fable 5.1 <[email protected]>
This commit is contained in:
Omar Sobh
2026-09-20 18:09:02 -05:00
co-authored by Claude Fable 5.1
parent e3958fc9c5
commit 77fda442c6
3 changed files with 188 additions and 10 deletions
@@ -198,3 +198,27 @@ extern "C" __global__ void e3_cg_scatter_f64(
int g = cells[t];
z[g] = (double) x0[g];
}
/* P1-5 (c): scatter compact (index, value) lists into persistent buffers. */
extern "C" __global__ void e3_scat_f64(int n, const unsigned int* __restrict__ idx, const double* __restrict__ val, double* __restrict__ dst)
{
int t = blockIdx.x * blockDim.x + threadIdx.x;
if (t < n) dst[idx[t]] = val[t];
}
extern "C" __global__ void e3_scat_f32(int n, const unsigned int* __restrict__ idx, const float* __restrict__ val, float* __restrict__ dst)
{
int t = blockIdx.x * blockDim.x + threadIdx.x;
if (t < n) dst[idx[t]] = val[t];
}
extern "C" __global__ void e3_scat_u32(int n, const unsigned int* __restrict__ idx, const unsigned int* __restrict__ val, unsigned int* __restrict__ dst)
{
int t = blockIdx.x * blockDim.x + threadIdx.x;
if (t < n) dst[idx[t]] = val[t];
}
extern "C" __global__ void e3_scat_i32(int n, const unsigned int* __restrict__ idx, const int* __restrict__ val, int* __restrict__ dst)
{
int t = blockIdx.x * blockDim.x + threadIdx.x;
if (t < n) dst[idx[t]] = val[t];
}
@@ -203,6 +203,43 @@ impl DeviceVcycle {
self.fine_cells = level0.cells.clone();
}
/// P1-5 (c): as [`Self::refresh_fine`] with only the rows of `rows`
/// scattered into the persistent coefficient and neighbour buffers; the
/// index lists that change length are uploaded whole.
pub fn refresh_fine_rows(&mut self, level0: &LevelExport, rows: &[u32]) {
let rt = runtime();
let up_u = |v: &[u32]| -> CudaSlice<u32> {
rt.stream
.memcpy_stod(if v.is_empty() { &[0u32][..] } else { v })
.expect("upload")
};
let lv = &mut self.levels[0];
if lv.ae.len() != level0.ae.len() {
drop(lv);
self.refresh_fine(level0);
return;
}
lv.n_cells = level0.cells.len();
lv.n_red = level0.red.len();
lv.n_black = level0.black.len();
lv.cells = up_u(&level0.cells);
lv.red = up_u(&level0.red);
lv.black = up_u(&level0.black);
let pick_u = |v: &[u32]| rows.iter().map(|&i| v[i as usize]).collect::<Vec<u32>>();
let pick_f = |v: &[f32]| rows.iter().map(|&i| v[i as usize]).collect::<Vec<f32>>();
super::device_cg::scatter_u32(rows, &pick_u(&level0.top), &mut lv.top);
super::device_cg::scatter_u32(rows, &pick_u(&level0.bot), &mut lv.bot);
super::device_cg::scatter_u32(rows, &pick_u(&level0.coarse_of), &mut lv.coarse_of);
super::device_cg::scatter_f32(rows, &pick_f(&level0.ae), &mut lv.ae);
super::device_cg::scatter_f32(rows, &pick_f(&level0.aw), &mut lv.aw);
super::device_cg::scatter_f32(rows, &pick_f(&level0.an), &mut lv.an);
super::device_cg::scatter_f32(rows, &pick_f(&level0.as_), &mut lv.as_);
super::device_cg::scatter_f32(rows, &pick_f(&level0.at), &mut lv.at);
super::device_cg::scatter_f32(rows, &pick_f(&level0.ab), &mut lv.ab);
super::device_cg::scatter_f32(rows, &pick_f(&level0.ap), &mut lv.ap);
self.fine_cells = level0.cells.clone();
}
/// Zero the finest level's right-hand side (a gather writes only the
/// current cells; stale entries elsewhere would restrict into the
/// coarse levels).
@@ -15,6 +15,10 @@ const KERNELS: &str = include_str!("../../../../kernels/cuda/e3_cg.cu");
struct CgKernels {
_module: Arc<CudaModule>,
scat_f64: CudaFunction,
scat_f32: CudaFunction,
scat_u32: CudaFunction,
scat_i32: CudaFunction,
spmv: CudaFunction,
residual: CudaFunction,
dot_partial: CudaFunction,
@@ -32,11 +36,101 @@ struct CgKernels {
static CG: OnceLock<CgKernels> = OnceLock::new();
/// P1-5 (c): write `val[t]` to `dst[idx[t]]` on the device (the compact
/// lists uploaded, one launch). Empty lists are a no-op.
pub(crate) fn scatter_f64(idx: &[u32], val: &[f64], dst: &mut CudaSlice<f64>) {
if idx.is_empty() {
return;
}
let rt = runtime();
let k = kernels();
let d_idx = rt.stream.memcpy_stod(idx).expect("scatter idx");
let d_val = rt.stream.memcpy_stod(val).expect("scatter val");
let n = idx.len() as i32;
unsafe {
rt.stream
.launch_builder(&k.scat_f64)
.arg(&n)
.arg(&d_idx)
.arg(&d_val)
.arg(dst)
.launch(cfg(idx.len()))
.expect("scatter f64");
}
}
pub(crate) fn scatter_f32(idx: &[u32], val: &[f32], dst: &mut CudaSlice<f32>) {
if idx.is_empty() {
return;
}
let rt = runtime();
let k = kernels();
let d_idx = rt.stream.memcpy_stod(idx).expect("scatter idx");
let d_val = rt.stream.memcpy_stod(val).expect("scatter val");
let n = idx.len() as i32;
unsafe {
rt.stream
.launch_builder(&k.scat_f32)
.arg(&n)
.arg(&d_idx)
.arg(&d_val)
.arg(dst)
.launch(cfg(idx.len()))
.expect("scatter f32");
}
}
pub(crate) fn scatter_u32(idx: &[u32], val: &[u32], dst: &mut CudaSlice<u32>) {
if idx.is_empty() {
return;
}
let rt = runtime();
let k = kernels();
let d_idx = rt.stream.memcpy_stod(idx).expect("scatter idx");
let d_val = rt.stream.memcpy_stod(val).expect("scatter val");
let n = idx.len() as i32;
unsafe {
rt.stream
.launch_builder(&k.scat_u32)
.arg(&n)
.arg(&d_idx)
.arg(&d_val)
.arg(dst)
.launch(cfg(idx.len()))
.expect("scatter u32");
}
}
pub(crate) fn scatter_i32(idx: &[u32], val: &[i32], dst: &mut CudaSlice<i32>) {
if idx.is_empty() {
return;
}
let rt = runtime();
let k = kernels();
let d_idx = rt.stream.memcpy_stod(idx).expect("scatter idx");
let d_val = rt.stream.memcpy_stod(val).expect("scatter val");
let n = idx.len() as i32;
unsafe {
rt.stream
.launch_builder(&k.scat_i32)
.arg(&n)
.arg(&d_idx)
.arg(&d_val)
.arg(dst)
.launch(cfg(idx.len()))
.expect("scatter i32");
}
}
fn kernels() -> &'static CgKernels {
CG.get_or_init(|| {
let module = load_module(KERNELS, "e3_cg.cu", true);
let f = |name: &str| module.load_function(name).expect(name);
CgKernels {
scat_f64: f("e3_scat_f64"),
scat_f32: f("e3_scat_f32"),
scat_u32: f("e3_scat_u32"),
scat_i32: f("e3_scat_i32"),
spmv: f("e3_cg_spmv"),
residual: f("e3_cg_residual"),
dot_partial: f("e3_cg_dot_partial"),
@@ -293,6 +387,25 @@ impl DeviceCg {
self.partial = rt.stream.alloc_zeros::<f64>(self.n_blocks).expect("alloc");
}
self.cells = up_u(&to_u32(&fine.cells));
// P1-5 (c): with a patched level the changed rows are scattered into
// the persistent buffers (the same values everywhere else).
let rows_u32: Option<Vec<u32>> = match (patched, changed) {
(true, Some(rows)) if self.ae.len() == fine.ae.len() => Some(rows.iter().map(|&i| i as u32).collect()),
_ => None,
};
if let Some(rows) = rows_u32.as_deref() {
let pick_f = |v: &[f64]| rows.iter().map(|&i| v[i as usize]).collect::<Vec<f64>>();
let pick_u = |v: &[usize]| rows.iter().map(|&i| if v[i as usize] == usize::MAX { u32::MAX } else { v[i as usize] as u32 }).collect::<Vec<u32>>();
scatter_u32(rows, &pick_u(&fine.top), &mut self.top);
scatter_u32(rows, &pick_u(&fine.bot), &mut self.bot);
scatter_f64(rows, &pick_f(&fine.ae), &mut self.ae);
scatter_f64(rows, &pick_f(&fine.aw), &mut self.aw);
scatter_f64(rows, &pick_f(&fine.an), &mut self.an);
scatter_f64(rows, &pick_f(&fine.as_), &mut self.as_);
scatter_f64(rows, &pick_f(&fine.at), &mut self.at);
scatter_f64(rows, &pick_f(&fine.ab), &mut self.ab);
scatter_f64(rows, &pick_f(&fine.ap), &mut self.ap);
} else {
self.top = up_u(&to_u32(&fine.top));
self.bot = up_u(&to_u32(&fine.bot));
self.ae = up_f(&fine.ae);
@@ -302,6 +415,7 @@ impl DeviceCg {
self.at = up_f(&fine.at);
self.ab = up_f(&fine.ab);
self.ap = up_f(&fine.ap);
}
self.link_ptr = up_u(&link_ptr);
self.link_idx = up_u(&link_idx);
self.link_coef = up_f(if link_coef.is_empty() {
@@ -332,7 +446,10 @@ impl DeviceCg {
assert!(same(&fine_level.ae, &full.ae) && same(&fine_level.aw, &full.aw) && same(&fine_level.an, &full.an) && same(&fine_level.as_, &full.as_) && same(&fine_level.at, &full.at) && same(&fine_level.ab, &full.ab) && same(&fine_level.ap, &full.ap), "band export: coefficients differ");
}
let l_export = lap.elapsed();
self.vcycle.refresh_fine(&fine_level);
match rows_u32.as_deref() {
Some(rows) if self.fine_export.is_some() => self.vcycle.refresh_fine_rows(&fine_level, rows),
_ => self.vcycle.refresh_fine(&fine_level),
}
self.fine_export = Some(fine_level);
self.fine = Some(fine);
if profile {