rtx-cfd 3D Stage 1 item 2: mg3_vcycle.cu (seven-point red-black half-sweep, residual, coarsest; top/bot neighbour arrays carry the periodic wrap), LevelExport3/export_hierarchy3/vcycle_f32_reference3, DeviceVcycle3 (the K=1 sequence); gate 2 HELD: device = host f32 V-cycle to 4e-7 relative on 96×40×{1,8} (periodic and walls) and 378×62×62; 4.88 ms per V-cycle at 1.45 M cells incl. transfers
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 / Test (macos-latest) (push) Blocked by required conditions
CI / Build (macos-latest) (push) Waiting to run
CI / CI Success (push) Blocked by required conditions
CI / Clippy Check (push) Failing after 4s
CI / Build (ubuntu-latest) (push) Failing after 4s
CI / Build CPU-Only (Explicit) (push) Failing after 4s
Documentation / Build API Documentation (push) Failing after 4s
Documentation / Build User Guide (push) Successful in 4s
CI / Format Check (push) Failing after 17s
Performance Benchmarks / Run Benchmarks (push) Successful in 3m47s
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 / Test (macos-latest) (push) Blocked by required conditions
CI / Build (macos-latest) (push) Waiting to run
CI / CI Success (push) Blocked by required conditions
CI / Clippy Check (push) Failing after 4s
CI / Build (ubuntu-latest) (push) Failing after 4s
CI / Build CPU-Only (Explicit) (push) Failing after 4s
Documentation / Build API Documentation (push) Failing after 4s
Documentation / Build User Guide (push) Successful in 4s
CI / Format Check (push) Failing after 17s
Performance Benchmarks / Run Benchmarks (push) Successful in 3m47s
Co-Authored-By: Claude Fable 5.1 <[email protected]>
This commit is contained in:
co-authored by
Claude Fable 5.1
parent
c13b07b68d
commit
05ed96f383
@@ -0,0 +1,116 @@
|
|||||||
|
/**
|
||||||
|
* 3D Stage 1 (omni-cortex `docs/three_d_stage1_campaign.md`, gate 2): the
|
||||||
|
* V-cycle's maps for a MASKED, VARIABLE-COEFFICIENT seven-point operator.
|
||||||
|
* Per-cell arrays are [n] (one march); the index lists (cells, colours,
|
||||||
|
* children) drive the maps; `top`/`bot` give the neighbour above/below per
|
||||||
|
* cell (UINT_MAX = none; a zero coefficient is never read), so a periodic
|
||||||
|
* z is data. Restriction, prolongation and zero are the 2D kernels (they
|
||||||
|
* never touch the stencil).
|
||||||
|
*/
|
||||||
|
#define NONE 0xFFFFFFFFu
|
||||||
|
|
||||||
|
__device__ __forceinline__ float nb_sum3(
|
||||||
|
int g, int nx, const float* ae, const float* aw, const float* an, const float* as_,
|
||||||
|
const float* at, const float* ab, const unsigned int* top, const unsigned int* bot,
|
||||||
|
const float* x)
|
||||||
|
{
|
||||||
|
float s = 0.0f;
|
||||||
|
float e = ae[g]; if (e != 0.0f) s += e * x[g + 1];
|
||||||
|
float w = aw[g]; if (w != 0.0f) s += w * x[g - 1];
|
||||||
|
float nn = an[g]; if (nn != 0.0f) s += nn * x[g + nx];
|
||||||
|
float ss = as_[g]; if (ss != 0.0f) s += ss * x[g - nx];
|
||||||
|
float t = at[g]; if (t != 0.0f) s += t * x[top[g]];
|
||||||
|
float b = ab[g]; if (b != 0.0f) s += b * x[bot[g]];
|
||||||
|
return s;
|
||||||
|
}
|
||||||
|
|
||||||
|
extern "C" __global__ void mg3_rb_half(
|
||||||
|
int n_col, const unsigned int* __restrict__ col,
|
||||||
|
const float* __restrict__ ae, const float* __restrict__ aw,
|
||||||
|
const float* __restrict__ an, const float* __restrict__ as_,
|
||||||
|
const float* __restrict__ at, const float* __restrict__ ab,
|
||||||
|
const unsigned int* __restrict__ top, const unsigned int* __restrict__ bot,
|
||||||
|
const float* __restrict__ ap, const float* __restrict__ b,
|
||||||
|
float* __restrict__ x, int nx)
|
||||||
|
{
|
||||||
|
int t = blockIdx.x * blockDim.x + threadIdx.x;
|
||||||
|
if (t >= n_col) return;
|
||||||
|
int g = col[t];
|
||||||
|
float s = nb_sum3(g, nx, ae, aw, an, as_, at, ab, top, bot, x);
|
||||||
|
x[g] = (b[g] + s) / ap[g];
|
||||||
|
}
|
||||||
|
|
||||||
|
extern "C" __global__ void mg3_residual(
|
||||||
|
int n_cells, const unsigned int* __restrict__ cells,
|
||||||
|
const float* __restrict__ ae, const float* __restrict__ aw,
|
||||||
|
const float* __restrict__ an, const float* __restrict__ as_,
|
||||||
|
const float* __restrict__ at, const float* __restrict__ ab,
|
||||||
|
const unsigned int* __restrict__ top, const unsigned int* __restrict__ bot,
|
||||||
|
const float* __restrict__ ap, const float* __restrict__ b,
|
||||||
|
const float* __restrict__ x, float* __restrict__ r, int nx)
|
||||||
|
{
|
||||||
|
int t = blockIdx.x * blockDim.x + threadIdx.x;
|
||||||
|
if (t >= n_cells) return;
|
||||||
|
int g = cells[t];
|
||||||
|
float s = nb_sum3(g, nx, ae, aw, an, as_, at, ab, top, bot, x);
|
||||||
|
r[g] = b[g] - (ap[g] * x[g] - s);
|
||||||
|
}
|
||||||
|
|
||||||
|
/* b_c[c] = sum of r_f over the children of coarse cell c (fixed order). */
|
||||||
|
extern "C" __global__ void mg3_restrict(
|
||||||
|
int n_coarse, const unsigned int* __restrict__ coarse_cells,
|
||||||
|
const unsigned int* __restrict__ children_ptr, const unsigned int* __restrict__ children_idx,
|
||||||
|
const float* __restrict__ r_f, float* __restrict__ b_c)
|
||||||
|
{
|
||||||
|
int t = blockIdx.x * blockDim.x + threadIdx.x;
|
||||||
|
if (t >= n_coarse) return;
|
||||||
|
float s = 0.0f;
|
||||||
|
for (unsigned int p = children_ptr[t]; p < children_ptr[t + 1]; ++p) s += r_f[children_idx[p]];
|
||||||
|
b_c[coarse_cells[t]] = s;
|
||||||
|
}
|
||||||
|
|
||||||
|
/* x_f += 2 x_c[coarse_of[idx]] */
|
||||||
|
extern "C" __global__ void mg3_prolong(
|
||||||
|
int n_cells, const unsigned int* __restrict__ cells, const unsigned int* __restrict__ coarse_of,
|
||||||
|
float* __restrict__ x_f, const float* __restrict__ x_c)
|
||||||
|
{
|
||||||
|
int t = blockIdx.x * blockDim.x + threadIdx.x;
|
||||||
|
if (t >= n_cells) return;
|
||||||
|
int idx = cells[t];
|
||||||
|
x_f[idx] += 2.0f * x_c[coarse_of[idx]];
|
||||||
|
}
|
||||||
|
|
||||||
|
extern "C" __global__ void mg3_zero(int n_cells, const unsigned int* __restrict__ cells, float* __restrict__ x)
|
||||||
|
{
|
||||||
|
int t = blockIdx.x * blockDim.x + threadIdx.x;
|
||||||
|
if (t >= n_cells) return;
|
||||||
|
x[cells[t]] = 0.0f;
|
||||||
|
}
|
||||||
|
|
||||||
|
/* The coarsest level: one thread, `sweeps` symmetric RED-BLACK sweeps
|
||||||
|
* (red, black, black, red) from zero — the host's ordering. */
|
||||||
|
extern "C" __global__ void mg3_coarsest(
|
||||||
|
int n_cells, const unsigned int* __restrict__ cells,
|
||||||
|
int n_red, const unsigned int* __restrict__ red,
|
||||||
|
int n_black, const unsigned int* __restrict__ black,
|
||||||
|
const float* __restrict__ ae, const float* __restrict__ aw,
|
||||||
|
const float* __restrict__ an, const float* __restrict__ as_,
|
||||||
|
const float* __restrict__ at, const float* __restrict__ ab,
|
||||||
|
const unsigned int* __restrict__ top, const unsigned int* __restrict__ bot,
|
||||||
|
const float* __restrict__ ap, const float* __restrict__ b,
|
||||||
|
float* __restrict__ x, int nx, int sweeps)
|
||||||
|
{
|
||||||
|
if (blockIdx.x * blockDim.x + threadIdx.x != 0) return;
|
||||||
|
for (int t = 0; t < n_cells; ++t) x[cells[t]] = 0.0f;
|
||||||
|
for (int sw = 0; sw < sweeps; ++sw) {
|
||||||
|
for (int half = 0; half < 4; ++half) {
|
||||||
|
const unsigned int* list = (half == 0 || half == 3) ? red : black;
|
||||||
|
int n_list = (half == 0 || half == 3) ? n_red : n_black;
|
||||||
|
for (int t = 0; t < n_list; ++t) {
|
||||||
|
int g = list[t];
|
||||||
|
float s = nb_sum3(g, nx, ae, aw, an, as_, at, ab, top, bot, x);
|
||||||
|
x[g] = (b[g] + s) / ap[g];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,328 @@
|
|||||||
|
//! Gate 2: the seven-point V-cycle on the CUDA device (`mg3_vcycle.cu`),
|
||||||
|
//! one march, persistent buffers per operator; the 2D `poisson/device.rs`
|
||||||
|
//! K = 1 sequence with the z terms. One runtime per process.
|
||||||
|
|
||||||
|
use super::LevelExport3;
|
||||||
|
use cudarc::driver::{
|
||||||
|
CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg,
|
||||||
|
};
|
||||||
|
use cudarc::nvrtc::{CompileOptions, compile_ptx_with_opts};
|
||||||
|
use std::sync::{Arc, OnceLock};
|
||||||
|
|
||||||
|
const KERNELS: &str = include_str!("../../../../kernels/cuda/mg3_vcycle.cu");
|
||||||
|
|
||||||
|
pub(crate) struct Runtime3 {
|
||||||
|
pub(crate) ctx: Arc<CudaContext>,
|
||||||
|
pub(crate) stream: Arc<CudaStream>,
|
||||||
|
_module: Arc<CudaModule>,
|
||||||
|
f_half: CudaFunction,
|
||||||
|
f_res: CudaFunction,
|
||||||
|
f_restrict: CudaFunction,
|
||||||
|
f_prolong: CudaFunction,
|
||||||
|
f_zero: CudaFunction,
|
||||||
|
f_coarsest: CudaFunction,
|
||||||
|
}
|
||||||
|
|
||||||
|
static RUNTIME: OnceLock<Runtime3> = OnceLock::new();
|
||||||
|
|
||||||
|
pub(crate) fn runtime3() -> &'static Runtime3 {
|
||||||
|
RUNTIME.get_or_init(|| {
|
||||||
|
let ctx = CudaContext::new(0).expect("CUDA context (device 0)");
|
||||||
|
let stream = ctx.default_stream();
|
||||||
|
let arch = std::env::var("RTX_CUDA_ARCH").unwrap_or_else(|_| "sm_120".to_string());
|
||||||
|
let ptx = compile_ptx_with_opts(
|
||||||
|
KERNELS,
|
||||||
|
CompileOptions {
|
||||||
|
arch: Some(Box::leak(arch.into_boxed_str())),
|
||||||
|
..Default::default()
|
||||||
|
},
|
||||||
|
)
|
||||||
|
.expect("nvrtc: mg3_vcycle.cu");
|
||||||
|
let module = ctx.load_module(ptx).expect("mg3_vcycle module");
|
||||||
|
let f = |name: &str| module.load_function(name).expect(name);
|
||||||
|
Runtime3 {
|
||||||
|
f_half: f("mg3_rb_half"),
|
||||||
|
f_res: f("mg3_residual"),
|
||||||
|
f_restrict: f("mg3_restrict"),
|
||||||
|
f_prolong: f("mg3_prolong"),
|
||||||
|
f_zero: f("mg3_zero"),
|
||||||
|
f_coarsest: f("mg3_coarsest"),
|
||||||
|
ctx,
|
||||||
|
stream,
|
||||||
|
_module: module,
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) struct DevLevel3 {
|
||||||
|
pub(crate) n: usize,
|
||||||
|
pub(crate) nx: i32,
|
||||||
|
pub(crate) n_cells: usize,
|
||||||
|
pub(crate) n_red: usize,
|
||||||
|
pub(crate) n_black: usize,
|
||||||
|
pub(crate) cells: CudaSlice<u32>,
|
||||||
|
pub(crate) red: CudaSlice<u32>,
|
||||||
|
pub(crate) black: CudaSlice<u32>,
|
||||||
|
pub(crate) top: CudaSlice<u32>,
|
||||||
|
pub(crate) bot: CudaSlice<u32>,
|
||||||
|
pub(crate) coarse_of: CudaSlice<u32>,
|
||||||
|
pub(crate) children_ptr: CudaSlice<u32>,
|
||||||
|
pub(crate) children_idx: CudaSlice<u32>,
|
||||||
|
pub(crate) ae: CudaSlice<f32>,
|
||||||
|
pub(crate) aw: CudaSlice<f32>,
|
||||||
|
pub(crate) an: CudaSlice<f32>,
|
||||||
|
pub(crate) as_: CudaSlice<f32>,
|
||||||
|
pub(crate) at: CudaSlice<f32>,
|
||||||
|
pub(crate) ab: CudaSlice<f32>,
|
||||||
|
pub(crate) ap: CudaSlice<f32>,
|
||||||
|
pub(crate) b: CudaSlice<f32>,
|
||||||
|
pub(crate) x: CudaSlice<f32>,
|
||||||
|
pub(crate) r: CudaSlice<f32>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// One operator's hierarchy on the device.
|
||||||
|
pub struct DeviceVcycle3 {
|
||||||
|
pub(crate) levels: Vec<DevLevel3>,
|
||||||
|
sweeps: usize,
|
||||||
|
fine_cells: Vec<u32>,
|
||||||
|
r_f32: Vec<f32>,
|
||||||
|
z_f32: Vec<f32>,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn cfg(n_items: usize) -> LaunchConfig {
|
||||||
|
LaunchConfig {
|
||||||
|
grid_dim: ((n_items as u32).div_ceil(256).max(1), 1, 1),
|
||||||
|
block_dim: (256, 1, 1),
|
||||||
|
shared_mem_bytes: 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl DeviceVcycle3 {
|
||||||
|
pub fn new(levels: &[LevelExport3], sweeps: usize) -> Self {
|
||||||
|
let rt = runtime3();
|
||||||
|
let up_u = |v: &[u32]| -> CudaSlice<u32> {
|
||||||
|
rt.stream
|
||||||
|
.memcpy_stod(if v.is_empty() { &[0u32][..] } else { v })
|
||||||
|
.expect("upload")
|
||||||
|
};
|
||||||
|
let up_f = |v: &[f32]| -> CudaSlice<f32> { rt.stream.memcpy_stod(v).expect("upload") };
|
||||||
|
let dev: Vec<DevLevel3> = levels
|
||||||
|
.iter()
|
||||||
|
.map(|l| {
|
||||||
|
let n = l.nx * l.ny * l.nz;
|
||||||
|
DevLevel3 {
|
||||||
|
n,
|
||||||
|
nx: l.nx as i32,
|
||||||
|
n_cells: l.cells.len(),
|
||||||
|
n_red: l.red.len(),
|
||||||
|
n_black: l.black.len(),
|
||||||
|
cells: up_u(&l.cells),
|
||||||
|
red: up_u(&l.red),
|
||||||
|
black: up_u(&l.black),
|
||||||
|
top: up_u(&l.top),
|
||||||
|
bot: up_u(&l.bot),
|
||||||
|
coarse_of: up_u(&l.coarse_of),
|
||||||
|
children_ptr: up_u(&l.children_ptr),
|
||||||
|
children_idx: up_u(&l.children_idx),
|
||||||
|
ae: up_f(&l.ae),
|
||||||
|
aw: up_f(&l.aw),
|
||||||
|
an: up_f(&l.an),
|
||||||
|
as_: up_f(&l.as_),
|
||||||
|
at: up_f(&l.at),
|
||||||
|
ab: up_f(&l.ab),
|
||||||
|
ap: up_f(&l.ap),
|
||||||
|
b: rt.stream.alloc_zeros::<f32>(n).expect("alloc"),
|
||||||
|
x: rt.stream.alloc_zeros::<f32>(n).expect("alloc"),
|
||||||
|
r: rt.stream.alloc_zeros::<f32>(n).expect("alloc"),
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
let n0 = dev[0].n;
|
||||||
|
Self {
|
||||||
|
levels: dev,
|
||||||
|
sweeps,
|
||||||
|
fine_cells: levels[0].cells.clone(),
|
||||||
|
r_f32: vec![0.0; n0],
|
||||||
|
z_f32: vec![0.0; n0],
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn depth(&self) -> usize {
|
||||||
|
self.levels.len()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn half(&mut self, l: usize, colour: u8) {
|
||||||
|
let rt = runtime3();
|
||||||
|
let lv = &mut self.levels[l];
|
||||||
|
let (list, n_list) = if colour == 0 {
|
||||||
|
(&lv.red, lv.n_red)
|
||||||
|
} else {
|
||||||
|
(&lv.black, lv.n_black)
|
||||||
|
};
|
||||||
|
let n_list_i = n_list as i32;
|
||||||
|
unsafe {
|
||||||
|
rt.stream
|
||||||
|
.launch_builder(&rt.f_half)
|
||||||
|
.arg(&n_list_i)
|
||||||
|
.arg(list)
|
||||||
|
.arg(&lv.ae)
|
||||||
|
.arg(&lv.aw)
|
||||||
|
.arg(&lv.an)
|
||||||
|
.arg(&lv.as_)
|
||||||
|
.arg(&lv.at)
|
||||||
|
.arg(&lv.ab)
|
||||||
|
.arg(&lv.top)
|
||||||
|
.arg(&lv.bot)
|
||||||
|
.arg(&lv.ap)
|
||||||
|
.arg(&lv.b)
|
||||||
|
.arg(&mut lv.x)
|
||||||
|
.arg(&lv.nx)
|
||||||
|
.launch(cfg(n_list))
|
||||||
|
.expect("mg3_rb_half");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn smooth(&mut self, l: usize) {
|
||||||
|
for _ in 0..self.sweeps {
|
||||||
|
self.half(l, 0);
|
||||||
|
self.half(l, 1);
|
||||||
|
self.half(l, 1);
|
||||||
|
self.half(l, 0);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The V-cycle on the device with `r` already in `levels[0].b`; the
|
||||||
|
/// correction is left in `levels[0].x`. No transfers.
|
||||||
|
pub(crate) fn vcycle_on_device(&mut self) {
|
||||||
|
let rt = runtime3();
|
||||||
|
let depth = self.levels.len();
|
||||||
|
for l in 0..depth - 1 {
|
||||||
|
{
|
||||||
|
let lv = &mut self.levels[l];
|
||||||
|
let n_cells_i = lv.n_cells as i32;
|
||||||
|
unsafe {
|
||||||
|
rt.stream
|
||||||
|
.launch_builder(&rt.f_zero)
|
||||||
|
.arg(&n_cells_i)
|
||||||
|
.arg(&lv.cells)
|
||||||
|
.arg(&mut lv.x)
|
||||||
|
.launch(cfg(lv.n_cells))
|
||||||
|
.expect("mg3_zero");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
self.smooth(l);
|
||||||
|
{
|
||||||
|
let lv = &mut self.levels[l];
|
||||||
|
let n_cells_i = lv.n_cells as i32;
|
||||||
|
unsafe {
|
||||||
|
rt.stream
|
||||||
|
.launch_builder(&rt.f_res)
|
||||||
|
.arg(&n_cells_i)
|
||||||
|
.arg(&lv.cells)
|
||||||
|
.arg(&lv.ae)
|
||||||
|
.arg(&lv.aw)
|
||||||
|
.arg(&lv.an)
|
||||||
|
.arg(&lv.as_)
|
||||||
|
.arg(&lv.at)
|
||||||
|
.arg(&lv.ab)
|
||||||
|
.arg(&lv.top)
|
||||||
|
.arg(&lv.bot)
|
||||||
|
.arg(&lv.ap)
|
||||||
|
.arg(&lv.b)
|
||||||
|
.arg(&lv.x)
|
||||||
|
.arg(&mut lv.r)
|
||||||
|
.arg(&lv.nx)
|
||||||
|
.launch(cfg(lv.n_cells))
|
||||||
|
.expect("mg3_residual");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let (fine, coarse) = self.levels.split_at_mut(l + 1);
|
||||||
|
let (lf, lc) = (&fine[l], &mut coarse[0]);
|
||||||
|
let n_c_cells_i = lc.n_cells as i32;
|
||||||
|
unsafe {
|
||||||
|
rt.stream
|
||||||
|
.launch_builder(&rt.f_restrict)
|
||||||
|
.arg(&n_c_cells_i)
|
||||||
|
.arg(&lc.cells)
|
||||||
|
.arg(&lf.children_ptr)
|
||||||
|
.arg(&lf.children_idx)
|
||||||
|
.arg(&lf.r)
|
||||||
|
.arg(&mut lc.b)
|
||||||
|
.launch(cfg(lc.n_cells))
|
||||||
|
.expect("mg3_restrict");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
{
|
||||||
|
let lv = &mut self.levels[depth - 1];
|
||||||
|
let (n_cells_i, sw_i) = (lv.n_cells as i32, 50i32);
|
||||||
|
let (n_red_i, n_black_i) = (lv.n_red as i32, lv.n_black as i32);
|
||||||
|
unsafe {
|
||||||
|
rt.stream
|
||||||
|
.launch_builder(&rt.f_coarsest)
|
||||||
|
.arg(&n_cells_i)
|
||||||
|
.arg(&lv.cells)
|
||||||
|
.arg(&n_red_i)
|
||||||
|
.arg(&lv.red)
|
||||||
|
.arg(&n_black_i)
|
||||||
|
.arg(&lv.black)
|
||||||
|
.arg(&lv.ae)
|
||||||
|
.arg(&lv.aw)
|
||||||
|
.arg(&lv.an)
|
||||||
|
.arg(&lv.as_)
|
||||||
|
.arg(&lv.at)
|
||||||
|
.arg(&lv.ab)
|
||||||
|
.arg(&lv.top)
|
||||||
|
.arg(&lv.bot)
|
||||||
|
.arg(&lv.ap)
|
||||||
|
.arg(&lv.b)
|
||||||
|
.arg(&mut lv.x)
|
||||||
|
.arg(&lv.nx)
|
||||||
|
.arg(&sw_i)
|
||||||
|
.launch(LaunchConfig {
|
||||||
|
grid_dim: (1, 1, 1),
|
||||||
|
block_dim: (32, 1, 1),
|
||||||
|
shared_mem_bytes: 0,
|
||||||
|
})
|
||||||
|
.expect("mg3_coarsest");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for l in (0..depth - 1).rev() {
|
||||||
|
{
|
||||||
|
let (fine, coarse) = self.levels.split_at_mut(l + 1);
|
||||||
|
let (lf, lc) = (&mut fine[l], &coarse[0]);
|
||||||
|
let n_cells_i = lf.n_cells as i32;
|
||||||
|
unsafe {
|
||||||
|
rt.stream
|
||||||
|
.launch_builder(&rt.f_prolong)
|
||||||
|
.arg(&n_cells_i)
|
||||||
|
.arg(&lf.cells)
|
||||||
|
.arg(&lf.coarse_of)
|
||||||
|
.arg(&mut lf.x)
|
||||||
|
.arg(&lc.x)
|
||||||
|
.launch(cfg(lf.n_cells))
|
||||||
|
.expect("mg3_prolong");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
self.smooth(l);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `z = M⁻¹ r` on the active cells (upload, V-cycle, download).
|
||||||
|
pub fn apply(&mut self, r: &[f64], z: &mut [f64]) {
|
||||||
|
let rt = runtime3();
|
||||||
|
for (dst, &src) in self.r_f32.iter_mut().zip(r) {
|
||||||
|
*dst = src as f32;
|
||||||
|
}
|
||||||
|
rt.stream
|
||||||
|
.memcpy_htod(&self.r_f32, &mut self.levels[0].b)
|
||||||
|
.expect("upload r");
|
||||||
|
self.vcycle_on_device();
|
||||||
|
rt.stream
|
||||||
|
.memcpy_dtoh(&self.levels[0].x, &mut self.z_f32)
|
||||||
|
.expect("download z");
|
||||||
|
rt.stream.synchronize().expect("sync");
|
||||||
|
for &idx in &self.fine_cells {
|
||||||
|
z[idx as usize] = self.z_f32[idx as usize] as f64;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,113 @@
|
|||||||
|
//! The f32 hierarchy exported for a device V-cycle (gate 2): per level the
|
||||||
|
//! index lists (cells, colours, parent map, CSR children) and the seven
|
||||||
|
//! coefficient arrays, plus explicit top/bottom neighbour indices so the
|
||||||
|
//! periodic wrap is data, not arithmetic, on the device.
|
||||||
|
|
||||||
|
use super::{Hierarchy3, PoissonProblem3D};
|
||||||
|
use crate::solvers::incompressible::poisson::MultigridParameters;
|
||||||
|
|
||||||
|
/// One exported level. `top`/`bot` hold the neighbour index above/below
|
||||||
|
/// each cell (`u32::MAX` = none); a zero coefficient is never read.
|
||||||
|
pub struct LevelExport3 {
|
||||||
|
pub nx: usize,
|
||||||
|
pub ny: usize,
|
||||||
|
pub nz: usize,
|
||||||
|
pub cells: Vec<u32>,
|
||||||
|
pub red: Vec<u32>,
|
||||||
|
pub black: Vec<u32>,
|
||||||
|
pub top: Vec<u32>,
|
||||||
|
pub bot: Vec<u32>,
|
||||||
|
/// Fine cell → coarse cell (`u32::MAX` without an equation; empty on
|
||||||
|
/// the coarsest level).
|
||||||
|
pub coarse_of: Vec<u32>,
|
||||||
|
/// For the NEXT level's cells in its `cells` order: the fine cells
|
||||||
|
/// restricting into each (CSR).
|
||||||
|
pub children_ptr: Vec<u32>,
|
||||||
|
pub children_idx: Vec<u32>,
|
||||||
|
pub ae: Vec<f32>,
|
||||||
|
pub aw: Vec<f32>,
|
||||||
|
pub an: Vec<f32>,
|
||||||
|
pub as_: Vec<f32>,
|
||||||
|
pub at: Vec<f32>,
|
||||||
|
pub ab: Vec<f32>,
|
||||||
|
pub ap: Vec<f32>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The f32 hierarchy of `problem`, level 0 fine.
|
||||||
|
pub fn export_hierarchy3(
|
||||||
|
problem: &PoissonProblem3D,
|
||||||
|
params: &MultigridParameters,
|
||||||
|
) -> Vec<LevelExport3> {
|
||||||
|
export_levels3(&Hierarchy3::<f32>::build(problem, params))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn export_levels3(hier: &Hierarchy3<f32>) -> Vec<LevelExport3> {
|
||||||
|
let depth = hier.levels.len();
|
||||||
|
let to_u32 = |v: &[usize]| {
|
||||||
|
v.iter()
|
||||||
|
.map(|&i| if i == usize::MAX { u32::MAX } else { i as u32 })
|
||||||
|
.collect::<Vec<u32>>()
|
||||||
|
};
|
||||||
|
(0..depth)
|
||||||
|
.map(|l| {
|
||||||
|
let lv = &hier.levels[l];
|
||||||
|
let (children_ptr, children_idx) = if l + 1 < depth {
|
||||||
|
let coarse = &hier.levels[l + 1];
|
||||||
|
let nc = coarse.problem.nx * coarse.problem.ny * coarse.problem.nz;
|
||||||
|
let mut pos = vec![usize::MAX; nc];
|
||||||
|
for (k, &c) in coarse.cells.iter().enumerate() {
|
||||||
|
pos[c] = k;
|
||||||
|
}
|
||||||
|
let mut lists: Vec<Vec<u32>> = vec![Vec::new(); coarse.cells.len()];
|
||||||
|
for &idx in &lv.cells {
|
||||||
|
let c = lv.coarse_of[idx];
|
||||||
|
if coarse.active[c] {
|
||||||
|
lists[pos[c]].push(idx as u32);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let mut ptr = Vec::with_capacity(lists.len() + 1);
|
||||||
|
let mut flat = Vec::new();
|
||||||
|
ptr.push(0u32);
|
||||||
|
for list in &lists {
|
||||||
|
flat.extend_from_slice(list);
|
||||||
|
ptr.push(flat.len() as u32);
|
||||||
|
}
|
||||||
|
(ptr, flat)
|
||||||
|
} else {
|
||||||
|
(Vec::new(), Vec::new())
|
||||||
|
};
|
||||||
|
LevelExport3 {
|
||||||
|
nx: lv.problem.nx,
|
||||||
|
ny: lv.problem.ny,
|
||||||
|
nz: lv.problem.nz,
|
||||||
|
cells: to_u32(&lv.cells),
|
||||||
|
red: to_u32(&lv.red),
|
||||||
|
black: to_u32(&lv.black),
|
||||||
|
top: to_u32(&lv.top),
|
||||||
|
bot: to_u32(&lv.bot),
|
||||||
|
coarse_of: to_u32(&lv.coarse_of),
|
||||||
|
children_ptr,
|
||||||
|
children_idx,
|
||||||
|
ae: lv.ae.clone(),
|
||||||
|
aw: lv.aw.clone(),
|
||||||
|
an: lv.an.clone(),
|
||||||
|
as_: lv.as_.clone(),
|
||||||
|
at: lv.at.clone(),
|
||||||
|
ab: lv.ab.clone(),
|
||||||
|
ap: lv.ap.clone(),
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// `z = M⁻¹ r` by the host f32 V-cycle: the reference a device V-cycle is
|
||||||
|
/// measured against.
|
||||||
|
pub fn vcycle_f32_reference3(
|
||||||
|
problem: &PoissonProblem3D,
|
||||||
|
params: &MultigridParameters,
|
||||||
|
r: &[f64],
|
||||||
|
z: &mut [f64],
|
||||||
|
) {
|
||||||
|
let mut hier = Hierarchy3::<f32>::build(problem, params);
|
||||||
|
hier.apply_preconditioner(r, z);
|
||||||
|
}
|
||||||
@@ -11,6 +11,11 @@ use crate::solvers::incompressible::poisson::{
|
|||||||
MgPrecision, MgScalar, MgSmoother, MultigridParameters, PoissonSolution,
|
MgPrecision, MgScalar, MgSmoother, MultigridParameters, PoissonSolution,
|
||||||
};
|
};
|
||||||
|
|
||||||
|
#[cfg(feature = "cuda")]
|
||||||
|
pub mod device;
|
||||||
|
pub mod export;
|
||||||
|
pub use export::{LevelExport3, export_hierarchy3, vcycle_f32_reference3};
|
||||||
|
|
||||||
/// Symmetric GS sweeps on the coarsest level (the 2D value).
|
/// Symmetric GS sweeps on the coarsest level (the 2D value).
|
||||||
const COARSEST_SWEEPS: usize = 50;
|
const COARSEST_SWEEPS: usize = 50;
|
||||||
/// The 2D `COARSE_CORRECTION`, proved dimension-independent there.
|
/// The 2D `COARSE_CORRECTION`, proved dimension-independent there.
|
||||||
|
|||||||
@@ -0,0 +1,134 @@
|
|||||||
|
//! 3D Stage 1, gate 2: the device seven-point V-cycle against the host f32
|
||||||
|
//! reference. `device_vcycle_matches_the_host_reference` is the cheap
|
||||||
|
//! correctness pin (a small channel with a z-cylinder hole, periodic z);
|
||||||
|
//! `bench_anchor_size` (ignored) is the 378 × 62 × 62 cylinder channel and
|
||||||
|
//! prints the ms per V-cycle (recorded, not asserted).
|
||||||
|
//!
|
||||||
|
//! `RTX_CUDA_ARCH=sm_120 cargo test --release -p rtx-cfd --features cuda --test three_d_gpu_vcycle_bench -- --nocapture [--ignored]`
|
||||||
|
#![cfg(feature = "cuda")]
|
||||||
|
|
||||||
|
use rtx_cfd::solvers::incompressible::three_d::poisson::device::DeviceVcycle3;
|
||||||
|
use rtx_cfd::solvers::incompressible::three_d::poisson::{
|
||||||
|
PoissonProblem3D, export_hierarchy3, vcycle_f32_reference3,
|
||||||
|
};
|
||||||
|
use rtx_cfd::solvers::incompressible::{MgSmoother, MultigridParameters};
|
||||||
|
use std::time::Instant;
|
||||||
|
|
||||||
|
/// A channel `2.5 × 0.41 × 0.41` (or a slice of it) on cubic cells with a
|
||||||
|
/// z-cylinder hole at (0.2, 0.2) r 0.05, an outlet Dirichlet, seeded rhs.
|
||||||
|
fn channel(nx: usize, ny: usize, nz: usize, periodic_z: bool, seed: u64) -> PoissonProblem3D {
|
||||||
|
let mut p = PoissonProblem3D::new(nx, ny, nz);
|
||||||
|
p.periodic_z = periodic_z;
|
||||||
|
let h = 0.41 / ny as f64;
|
||||||
|
let dt = 3.24e-4;
|
||||||
|
let a = dt * h * h / h; // dt · face area / distance on cubic cells
|
||||||
|
let hole = |i: usize, j: usize| {
|
||||||
|
let (x, y) = ((i as f64 + 0.5) * h, (j as f64 + 0.5) * h);
|
||||||
|
(x - 0.2).powi(2) + (y - 0.2).powi(2) < 0.05 * 0.05
|
||||||
|
};
|
||||||
|
for k in 0..nz {
|
||||||
|
for j in 0..ny {
|
||||||
|
for i in 0..nx {
|
||||||
|
let idx = p.index(k, j, i);
|
||||||
|
if hole(i, j) {
|
||||||
|
p.active[idx] = false;
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
if i + 1 < nx && !hole(i + 1, j) {
|
||||||
|
p.ae[idx] = a;
|
||||||
|
}
|
||||||
|
if i > 0 && !hole(i - 1, j) {
|
||||||
|
p.aw[idx] = a;
|
||||||
|
}
|
||||||
|
if j + 1 < ny && !hole(i, j + 1) {
|
||||||
|
p.an[idx] = a;
|
||||||
|
}
|
||||||
|
if j > 0 && !hole(i, j - 1) {
|
||||||
|
p.as_[idx] = a;
|
||||||
|
}
|
||||||
|
if k + 1 < nz || periodic_z {
|
||||||
|
p.at[idx] = a;
|
||||||
|
}
|
||||||
|
if k > 0 || periodic_z {
|
||||||
|
p.ab[idx] = a;
|
||||||
|
}
|
||||||
|
if i + 1 == nx {
|
||||||
|
p.extra_diag[idx] = 2.0 * a;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
let mut state = seed | 1;
|
||||||
|
for idx in 0..nx * ny * nz {
|
||||||
|
state ^= state << 13;
|
||||||
|
state ^= state >> 7;
|
||||||
|
state ^= state << 17;
|
||||||
|
p.rhs[idx] = if p.active[idx] {
|
||||||
|
1e-3 * ((state >> 11) as f64 / (1u64 << 53) as f64 - 0.5)
|
||||||
|
} else {
|
||||||
|
0.0
|
||||||
|
};
|
||||||
|
}
|
||||||
|
p
|
||||||
|
}
|
||||||
|
|
||||||
|
fn params() -> MultigridParameters {
|
||||||
|
MultigridParameters {
|
||||||
|
smoother: MgSmoother::RedBlack,
|
||||||
|
..MultigridParameters::default()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn compare(p: &PoissonProblem3D, label: &str) -> (f64, f64) {
|
||||||
|
let params = params();
|
||||||
|
let n = p.nx * p.ny * p.nz;
|
||||||
|
let mut z_ref = vec![0.0; n];
|
||||||
|
vcycle_f32_reference3(p, ¶ms, &p.rhs, &mut z_ref);
|
||||||
|
let levels = export_hierarchy3(p, ¶ms);
|
||||||
|
println!(
|
||||||
|
" {label}: {} levels, cells per level {:?}",
|
||||||
|
levels.len(),
|
||||||
|
levels.iter().map(|l| l.cells.len()).collect::<Vec<_>>()
|
||||||
|
);
|
||||||
|
let mut dev = DeviceVcycle3::new(&levels, params.smoother_sweeps.max(1));
|
||||||
|
let mut z = vec![0.0; n];
|
||||||
|
dev.apply(&p.rhs, &mut z);
|
||||||
|
let scale = z_ref.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
|
||||||
|
let worst = z
|
||||||
|
.iter()
|
||||||
|
.zip(&z_ref)
|
||||||
|
.fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
|
||||||
|
println!(
|
||||||
|
" {label}: device vs host f32 V-cycle max |Δz| {worst:.3e} on a scale of {scale:.3e}"
|
||||||
|
);
|
||||||
|
// Timing (one march, transfers included).
|
||||||
|
let reps = 20;
|
||||||
|
dev.apply(&p.rhs, &mut z);
|
||||||
|
let t0 = Instant::now();
|
||||||
|
for _ in 0..reps {
|
||||||
|
dev.apply(&p.rhs, &mut z);
|
||||||
|
}
|
||||||
|
let ms = t0.elapsed().as_secs_f64() * 1e3 / reps as f64;
|
||||||
|
println!(" {label}: {ms:.3} ms per V-cycle including the transfers ({n} cells)");
|
||||||
|
(worst, scale)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn device_vcycle_matches_the_host_reference() {
|
||||||
|
for (nz, periodic) in [(1usize, false), (8, true), (8, false)] {
|
||||||
|
let p = channel(96, 40, nz, periodic, 3);
|
||||||
|
let (worst, scale) = compare(&p, &format!("96×40×{nz} periodic {periodic}"));
|
||||||
|
assert!(
|
||||||
|
worst < 1e-4 * scale,
|
||||||
|
"device V-cycle differs: {worst:.3e} of {scale:.3e}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
#[ignore = "the anchor-size bench (378 × 62 × 62): prints ms per V-cycle"]
|
||||||
|
fn bench_anchor_size() {
|
||||||
|
let p = channel(378, 62, 62, false, 11);
|
||||||
|
let (worst, scale) = compare(&p, "378×62×62 walls");
|
||||||
|
assert!(worst < 1e-4 * scale);
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user