diff --git a/crates/specialized/rtx-cfd/src/kernels/cuda/e3_cg.cu b/crates/specialized/rtx-cfd/src/kernels/cuda/e3_cg.cu index f92ea06..78f7a8a 100644 --- a/crates/specialized/rtx-cfd/src/kernels/cuda/e3_cg.cu +++ b/crates/specialized/rtx-cfd/src/kernels/cuda/e3_cg.cu @@ -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]; +} diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device.rs index afb956e..9eafae2 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device.rs @@ -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 { + 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::>(); + let pick_f = |v: &[f32]| rows.iter().map(|&i| v[i as usize]).collect::>(); + 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). diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device_cg.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device_cg.rs index 6a4f4fb..e2377fb 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device_cg.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device_cg.rs @@ -15,6 +15,10 @@ const KERNELS: &str = include_str!("../../../../kernels/cuda/e3_cg.cu"); struct CgKernels { _module: Arc, + 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 = 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) { + 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) { + 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) { + 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) { + 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,15 +387,35 @@ impl DeviceCg { self.partial = rt.stream.alloc_zeros::(self.n_blocks).expect("alloc"); } self.cells = up_u(&to_u32(&fine.cells)); - self.top = up_u(&to_u32(&fine.top)); - self.bot = up_u(&to_u32(&fine.bot)); - self.ae = up_f(&fine.ae); - self.aw = up_f(&fine.aw); - self.an = up_f(&fine.an); - self.as_ = up_f(&fine.as_); - self.at = up_f(&fine.at); - self.ab = up_f(&fine.ab); - self.ap = up_f(&fine.ap); + // 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> = 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::>(); + 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::>(); + 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); + self.aw = up_f(&fine.aw); + self.an = up_f(&fine.an); + self.as_ = up_f(&fine.as_); + 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 {