embedded3 R6-3: the moving path's Poisson setup on the device (RTX_E3_POISSON_DEVICE=1, default off)

The operator (poisson_operator + merge_small_cells), Level::new's activity,
masked couplings, z links and cell/colour lists, the link CSR, the
components' count and singular count (lock-free union-find), and the f32
fine export with its parent map, built on the device from DeviceCut's
projection tables into the CG's and V-cycle's persistent buffers
(e3_pset.cu, poisson/device_cg/setup.rs). RTX_E3_BAND_CHECK=1 compares every
structure against the host build bit for bit. Fallback to the host setup
(logged once) with gradient weights or a periodic z of <= 2 planes.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
Omar Sobh
2026-09-25 03:25:44 -05:00
co-authored by Claude Opus 5.5
parent c860783c81
commit 348f4621d1
7 changed files with 1344 additions and 2 deletions
@@ -0,0 +1,463 @@
/**
* embedded3 R6-3: the moving path's Poisson setup on the device, behind
* `RTX_E3_POISSON_DEVICE=1` — the host `Solver::poisson_operator` (with
* `merge_small_cells`), `Level::new`'s masking, lists and z links, the
* link CSR, `Components::summary_planes`'s count and singular flags, and
* `export_fine_from`'s f32 casts and parent map, from the projection tables
* DeviceCut already holds (step apertures, open flags, activity, owners,
* the merged cells' CSR). fp64, FMA contraction off; every sum in the host's
* order:
*
* e3_ps_assemble per cell: the raw seven coefficients, the outlet
* Dirichlet part (x1, x0, y1, y0, z1, z0 order), the
* raw activity and the small (merged) flag;
* e3_ps_links per small cell (ascending): its six candidate links
* in the host's face order (e, w, n, s, t, b) — the
* coefficient on s toward nb is zero when nb is a small
* cell processed before s (nb < s);
* e3_ps_extra per master: its slaves' Dirichlet parts added in
* ascending order;
* e3_ps_merge_zero per cell: a small cell's row, Dirichlet part and
* activity zeroed; every coefficient toward a small
* cell zeroed;
* e3_ps_link_counts / e3_ps_link_fill the link CSR from the host-ordered
* per-cell entries;
* e3_ps_diag ap = stencil (+ the link sum); the level activity;
* e3_ps_mask the masked couplings, the z links, the f32 export and
* the parent map;
* e3_ps_flags_* list flags (active, red, black) for compaction;
* e3_ps_uf_* the connected components by a lock-free union-find
* (roots = the smallest index; the count and the
* singular count are order-independent).
*/
typedef unsigned int u32;
typedef unsigned char u8;
#define NONE 0xFFFFFFFFu
struct PsGrid {
int nx, ny, nz, periodic;
int out_x0, out_x1, out_y0, out_y1, out_z0, out_z1;
double ae_int, an_int, at_int, ae_out, an_out, at_out;
};
/* The seven coefficient arrays (raw, then merged, then masked in place). */
struct PsOp {
double *ae, *aw, *an, *as_, *at, *ab, *extra;
};
/* The projection tables the operator reads. */
struct PsIn {
const int* cell_active;
const int *open_u, *open_v, *open_w;
const double *a_u, *a_v, *a_w;
const u32* owner;
const u32* fold_ptr;
const u32* fold_idx;
};
__device__ __forceinline__ long long ps_top(const PsGrid& g, long long idx, int k)
{
long long nxy = (long long) g.nx * g.ny;
if (k + 1 < g.nz) return idx + nxy;
if (g.periodic && g.nz > 1) return idx - (long long) (g.nz - 1) * nxy;
return -1;
}
__device__ __forceinline__ long long ps_bot(const PsGrid& g, long long idx, int k)
{
long long nxy = (long long) g.nx * g.ny;
if (k > 0) return idx - nxy;
if (g.periodic && g.nz > 1) return idx + (long long) (g.nz - 1) * nxy;
return -1;
}
extern "C" __global__ void e3_ps_assemble(PsGrid g, PsIn in, PsOp op, u8* __restrict__ act, u8* __restrict__ small)
{
long long t = (long long) blockIdx.x * blockDim.x + threadIdx.x;
long long nxy = (long long) g.nx * g.ny;
long long n = nxy * g.nz;
if (t >= n) return;
int k = (int) (t / nxy);
int j = (int) ((t % nxy) / g.nx);
int i = (int) (t % g.nx);
double ae = 0.0, aw = 0.0, an = 0.0, as_ = 0.0, at = 0.0, ab = 0.0, extra = 0.0;
if (!in.cell_active[t]) {
op.ae[t] = 0.0; op.aw[t] = 0.0; op.an[t] = 0.0; op.as_[t] = 0.0;
op.at[t] = 0.0; op.ab[t] = 0.0; op.extra[t] = 0.0;
act[t] = 0;
small[t] = 0;
return;
}
long long nxu = g.nx + 1;
long long nyv = g.ny + 1;
if (i + 1 == g.nx) {
if (g.out_x1) extra += g.ae_out;
} else {
long long f = ((long long) k * g.ny + j) * nxu + (i + 1);
if (in.open_u[f]) ae = g.ae_int * in.a_u[f];
}
if (i == 0) {
if (g.out_x0) extra += g.ae_out;
} else {
long long f = ((long long) k * g.ny + j) * nxu + i;
if (in.open_u[f]) aw = g.ae_int * in.a_u[f];
}
if (j + 1 == g.ny) {
if (g.out_y1) extra += g.an_out;
} else {
long long f = ((long long) k * nyv + (j + 1)) * g.nx + i;
if (in.open_v[f]) an = g.an_int * in.a_v[f];
}
if (j == 0) {
if (g.out_y0) extra += g.an_out;
} else {
long long f = ((long long) k * nyv + j) * g.nx + i;
if (in.open_v[f]) as_ = g.an_int * in.a_v[f];
}
if (k + 1 == g.nz && !g.periodic) {
if (g.out_z1) extra += g.at_out;
} else {
long long f = ((long long) ((k + 1) % g.nz) * g.ny + j) * g.nx + i;
if (in.open_w[f]) at = g.at_int * in.a_w[f];
}
if (k == 0 && !g.periodic) {
if (g.out_z0) extra += g.at_out;
} else {
long long f = ((long long) k * g.ny + j) * g.nx + i;
if (in.open_w[f]) ab = g.at_int * in.a_w[f];
}
op.ae[t] = ae; op.aw[t] = aw; op.an[t] = an; op.as_[t] = as_;
op.at[t] = at; op.ab[t] = ab; op.extra[t] = extra;
act[t] = 1;
small[t] = in.owner[t] != (u32) t ? 1 : 0;
}
/* Per small cell (rank r in the ascending list): six link slots. */
extern "C" __global__ void e3_ps_links(
PsGrid g, int n_small, const u32* __restrict__ small_list, const u8* __restrict__ small,
const u32* __restrict__ owner, PsOp op,
u32* __restrict__ la, u32* __restrict__ lb, double* __restrict__ lc, u8* __restrict__ lvalid)
{
int r = blockIdx.x * blockDim.x + threadIdx.x;
if (r >= n_small) return;
long long nxy = (long long) g.nx * g.ny;
long long s = small_list[r];
int k = (int) (s / nxy);
int j = (int) ((s % nxy) / g.nx);
int i = (int) (s % g.nx);
u32 m = owner[s];
long long nb[6];
double c[6];
nb[0] = i + 1 < g.nx ? s + 1 : -1; c[0] = op.ae[s];
nb[1] = i > 0 ? s - 1 : -1; c[1] = op.aw[s];
nb[2] = j + 1 < g.ny ? s + g.nx : -1; c[2] = op.an[s];
nb[3] = j > 0 ? s - g.nx : -1; c[3] = op.as_[s];
nb[4] = ps_top(g, s, k); c[4] = op.at[s];
nb[5] = ps_bot(g, s, k); c[5] = op.ab[s];
for (int f = 0; f < 6; ++f) {
long long slot = (long long) r * 6 + f;
u8 ok = 0;
u32 a = 0, b = 0;
double cc = 0.0;
if (nb[f] >= 0) {
cc = (small[nb[f]] && nb[f] < s) ? 0.0 : c[f];
u32 tt = owner[nb[f]];
if (cc > 0.0 && tt != m) {
ok = 1;
a = m < tt ? m : tt;
b = m < tt ? tt : m;
}
}
la[slot] = a; lb[slot] = b; lc[slot] = cc; lvalid[slot] = ok;
}
}
/* Per master: the slaves' Dirichlet parts in ascending order (before the zeroing). */
extern "C" __global__ void e3_ps_extra(long long n, PsIn in, const u8* __restrict__ small, PsOp op)
{
long long t = (long long) blockIdx.x * blockDim.x + threadIdx.x;
if (t >= n) return;
u32 p0 = in.fold_ptr[t], p1 = in.fold_ptr[t + 1];
if (p0 == p1 || small[t]) return;
double e = op.extra[t];
for (u32 l = p0; l < p1; ++l) {
u32 s = in.fold_idx[l];
if (small[s]) e += op.extra[s];
}
op.extra[t] = e;
}
extern "C" __global__ void e3_ps_merge_zero(PsGrid g, const u8* __restrict__ small, PsOp op, u8* __restrict__ act)
{
long long t = (long long) blockIdx.x * blockDim.x + threadIdx.x;
long long nxy = (long long) g.nx * g.ny;
long long n = nxy * g.nz;
if (t >= n) return;
if (small[t]) {
op.ae[t] = 0.0; op.aw[t] = 0.0; op.an[t] = 0.0; op.as_[t] = 0.0;
op.at[t] = 0.0; op.ab[t] = 0.0; op.extra[t] = 0.0;
act[t] = 0;
return;
}
int k = (int) (t / nxy);
int j = (int) ((t % nxy) / g.nx);
int i = (int) (t % g.nx);
if (i + 1 < g.nx && small[t + 1]) op.ae[t] = 0.0;
if (i > 0 && small[t - 1]) op.aw[t] = 0.0;
if (j + 1 < g.ny && small[t + g.nx]) op.an[t] = 0.0;
if (j > 0 && small[t - g.nx]) op.as_[t] = 0.0;
long long tp = ps_top(g, t, k);
if (tp >= 0 && small[tp]) op.at[t] = 0.0;
long long bt = ps_bot(g, t, k);
if (bt >= 0 && small[bt]) op.ab[t] = 0.0;
}
/* The link CSR's per-cell counts (cells with links, their counts). */
extern "C" __global__ void e3_ps_link_counts(int m, const u32* __restrict__ cell, const u32* __restrict__ count, u32* __restrict__ counts)
{
int t = blockIdx.x * blockDim.x + threadIdx.x;
if (t < m) counts[cell[t]] = count[t];
}
/* ap = stencil (+ the link sum in CSR order); the level activity. */
extern "C" __global__ void e3_ps_diag(
long long n, PsOp op, const u8* __restrict__ act, int has_links,
const u32* __restrict__ link_ptr, const double* __restrict__ link_coef,
double* __restrict__ ap_out, u8* __restrict__ lact)
{
long long t = (long long) blockIdx.x * blockDim.x + threadIdx.x;
if (t >= n) return;
double stencil = op.ae[t] + op.aw[t] + op.an[t] + op.as_[t] + op.at[t] + op.ab[t] + op.extra[t];
double ap = stencil;
if (has_links) {
double ls = 0.0;
for (u32 l = link_ptr[t]; l < link_ptr[t + 1]; ++l) ls += link_coef[l];
ap = stencil + ls;
}
ap_out[t] = ap;
lact[t] = (act[t] && ap > 0.0) ? 1 : 0;
}
/* Level::new's masked couplings and z links (in place), the f32 export and the parent map. */
extern "C" __global__ void e3_ps_mask(
PsGrid g, PsOp op, const u8* __restrict__ lact, const double* __restrict__ ap,
u32* __restrict__ top, u32* __restrict__ bot,
float* __restrict__ fae, float* __restrict__ faw, float* __restrict__ fan, float* __restrict__ fas,
float* __restrict__ fat, float* __restrict__ fab, float* __restrict__ fap,
u32* __restrict__ ftop, u32* __restrict__ fbot, u32* __restrict__ coarse_of)
{
long long t = (long long) blockIdx.x * blockDim.x + threadIdx.x;
long long nxy = (long long) g.nx * g.ny;
long long n = nxy * g.nz;
if (t >= n) return;
int k = (int) (t / nxy);
int j = (int) ((t % nxy) / g.nx);
int i = (int) (t % g.nx);
int a = lact[t];
long long tp = ps_top(g, t, k);
long long bt = ps_bot(g, t, k);
int tp_ok = tp >= 0 && lact[tp];
int bt_ok = bt >= 0 && lact[bt];
double ae = (a && i + 1 < g.nx && lact[t + 1]) ? op.ae[t] : 0.0;
double aw = (a && i > 0 && lact[t - 1]) ? op.aw[t] : 0.0;
double an = (a && j + 1 < g.ny && lact[t + g.nx]) ? op.an[t] : 0.0;
double as_ = (a && j > 0 && lact[t - g.nx]) ? op.as_[t] : 0.0;
double at = (a && tp_ok) ? op.at[t] : 0.0;
double ab = (a && bt_ok) ? op.ab[t] : 0.0;
op.ae[t] = ae; op.aw[t] = aw; op.an[t] = an; op.as_[t] = as_; op.at[t] = at; op.ab[t] = ab;
u32 tv = (a && tp_ok) ? (u32) tp : NONE;
u32 bv = (a && bt_ok) ? (u32) bt : NONE;
top[t] = tv; bot[t] = bv; ftop[t] = tv; fbot[t] = bv;
fae[t] = (float) ae; faw[t] = (float) aw; fan[t] = (float) an; fas[t] = (float) as_;
fat[t] = (float) at; fab[t] = (float) ab; fap[t] = (float) ap[t];
if (!a) {
coarse_of[t] = NONE;
} else {
int nxc = g.nx / 2 > 1 ? g.nx / 2 : 1;
int nyc = g.ny / 2 > 1 ? g.ny / 2 : 1;
int nzc = g.nz / 2 > 1 ? g.nz / 2 : 1;
int ic = i / 2 < nxc - 1 ? i / 2 : nxc - 1;
int jc = j / 2 < nyc - 1 ? j / 2 : nyc - 1;
int kc = k / 2 < nzc - 1 ? k / 2 : nzc - 1;
coarse_of[t] = (u32) (((long long) kc * nyc + jc) * nxc + ic);
}
}
/* List flags: which = 0 active, 1 red (parity 0), 2 black (parity 1). */
extern "C" __global__ void e3_ps_flags(PsGrid g, const u8* __restrict__ lact, int which, u8* __restrict__ flag)
{
long long t = (long long) blockIdx.x * blockDim.x + threadIdx.x;
long long nxy = (long long) g.nx * g.ny;
long long n = nxy * g.nz;
if (t >= n) return;
u8 f = lact[t];
if (which > 0 && f) {
int k = (int) (t / nxy);
int j = (int) ((t % nxy) / g.nx);
int i = (int) (t % g.nx);
int parity = (i + j + k) % 2;
f = (parity == which - 1) ? 1 : 0;
}
flag[t] = f;
}
/* ---- the connected components: lock-free union-find (ECL-CC style) ---- */
__device__ __forceinline__ u32 uf_find(u32* parent, u32 x)
{
volatile u32* p = parent;
u32 cur = x;
u32 next = p[cur];
while (next != cur) {
u32 nn = p[next];
if (nn != next) p[cur] = nn; /* path halving: nn is an ancestor of cur */
cur = next;
next = p[cur];
}
return cur;
}
__device__ __forceinline__ void uf_union(u32* parent, u32 a, u32 b)
{
u32 ra = uf_find(parent, a), rb = uf_find(parent, b);
while (ra != rb) {
u32 hi = ra > rb ? ra : rb;
u32 lo = ra > rb ? rb : ra;
u32 old = atomicCAS(&parent[hi], hi, lo);
if (old == hi) return;
ra = uf_find(parent, old);
rb = uf_find(parent, lo);
}
}
extern "C" __global__ void e3_ps_uf_init(long long n, u32* __restrict__ parent)
{
long long t = (long long) blockIdx.x * blockDim.x + threadIdx.x;
if (t < n) parent[t] = (u32) t;
}
/* Every positive masked coupling from a raw-active cell toward a raw-active neighbour, and its links. */
extern "C" __global__ void e3_ps_uf_hook(
PsGrid g, PsOp op, const u8* __restrict__ act, const u32* __restrict__ link_ptr,
const u32* __restrict__ link_idx, const double* __restrict__ link_coef, u32* parent)
{
long long t = (long long) blockIdx.x * blockDim.x + threadIdx.x;
long long nxy = (long long) g.nx * g.ny;
long long n = nxy * g.nz;
if (t >= n || !act[t]) return;
int k = (int) (t / nxy);
int j = (int) ((t % nxy) / g.nx);
int i = (int) (t % g.nx);
if (i + 1 < g.nx && op.ae[t] > 0.0 && act[t + 1]) uf_union(parent, (u32) t, (u32) (t + 1));
if (i > 0 && op.aw[t] > 0.0 && act[t - 1]) uf_union(parent, (u32) t, (u32) (t - 1));
if (j + 1 < g.ny && op.an[t] > 0.0 && act[t + g.nx]) uf_union(parent, (u32) t, (u32) (t + g.nx));
if (j > 0 && op.as_[t] > 0.0 && act[t - g.nx]) uf_union(parent, (u32) t, (u32) (t - g.nx));
long long tp = ps_top(g, t, k);
if (tp >= 0 && op.at[t] > 0.0 && act[tp]) uf_union(parent, (u32) t, (u32) tp);
long long bt = ps_bot(g, t, k);
if (bt >= 0 && op.ab[t] > 0.0 && act[bt]) uf_union(parent, (u32) t, (u32) bt);
for (u32 l = link_ptr[t]; l < link_ptr[t + 1]; ++l) {
u32 o = link_idx[l];
if (link_coef[l] > 0.0 && act[o]) uf_union(parent, (u32) t, o);
}
}
/* Every raw-active cell with a Dirichlet part marks its root. */
extern "C" __global__ void e3_ps_uf_dirichlet(long long n, const u8* __restrict__ act, const double* __restrict__ extra, u32* parent, u8* __restrict__ dir)
{
long long t = (long long) blockIdx.x * blockDim.x + threadIdx.x;
if (t >= n || !act[t]) return;
if (extra[t] > 0.0) dir[uf_find(parent, (u32) t)] = 1;
}
/* counts[0] = the roots (components), counts[1] = the roots without a Dirichlet part. */
extern "C" __global__ void e3_ps_uf_count(long long n, const u8* __restrict__ act, const u32* __restrict__ parent, const u8* __restrict__ dir, u32* __restrict__ counts)
{
long long t = (long long) blockIdx.x * blockDim.x + threadIdx.x;
if (t >= n || !act[t] || parent[t] != (u32) t) return;
atomicAdd(&counts[0], 1u);
if (!dir[t]) atomicAdd(&counts[1], 1u);
}
/* ---- block scans (the e3_mask.cu scans, own names) ---- */
#define PS_SCAN_BLOCK 1024
__device__ u32 ps_block_excl_scan(u32 v, u32* total)
{
__shared__ u32 warp_sums[32];
int lane = threadIdx.x & 31, wid = threadIdx.x >> 5;
u32 x = v;
for (int o = 1; o < 32; o <<= 1) {
u32 y = __shfl_up_sync(0xffffffffu, x, o);
if (lane >= o) x += y;
}
if (lane == 31) warp_sums[wid] = x;
__syncthreads();
if (wid == 0) {
u32 s = warp_sums[lane];
for (int o = 1; o < 32; o <<= 1) {
u32 y = __shfl_up_sync(0xffffffffu, s, o);
if (lane >= o) s += y;
}
warp_sums[lane] = s;
}
__syncthreads();
u32 incl = x + (wid > 0 ? warp_sums[wid - 1] : 0u);
*total = warp_sums[31];
__syncthreads();
return incl - v;
}
extern "C" __global__ void e3_ps_scan_count_flags(long long n, const u8* __restrict__ flags, u32* __restrict__ sums)
{
long long i = (long long) blockIdx.x * PS_SCAN_BLOCK + threadIdx.x;
u32 total;
ps_block_excl_scan(i < n ? (u32) (flags[i] != 0) : 0u, &total);
if (threadIdx.x == 0) sums[blockIdx.x] = total;
}
extern "C" __global__ void e3_ps_scan_count_vals(long long n, const u32* __restrict__ vals, u32* __restrict__ sums)
{
long long i = (long long) blockIdx.x * PS_SCAN_BLOCK + threadIdx.x;
u32 total;
ps_block_excl_scan(i < n ? vals[i] : 0u, &total);
if (threadIdx.x == 0) sums[blockIdx.x] = total;
}
extern "C" __global__ void e3_ps_scan_top(int nb, u32* __restrict__ sums)
{
__shared__ u32 carry;
if (threadIdx.x == 0) carry = 0u;
__syncthreads();
for (int base = 0; base < nb; base += PS_SCAN_BLOCK) {
int i = base + threadIdx.x;
u32 v = i < nb ? sums[i] : 0u;
u32 total;
u32 ex = ps_block_excl_scan(v, &total);
u32 c = carry;
if (i < nb) sums[i] = c + ex;
__syncthreads();
if (threadIdx.x == 0) carry = c + total;
__syncthreads();
}
if (threadIdx.x == 0) sums[nb] = carry;
}
extern "C" __global__ void e3_ps_scan_compact(long long n, const u8* __restrict__ flags, const u32* __restrict__ offsets, u32* __restrict__ out)
{
long long i = (long long) blockIdx.x * PS_SCAN_BLOCK + threadIdx.x;
u32 v = i < n ? (u32) (flags[i] != 0) : 0u;
u32 total;
u32 ex = ps_block_excl_scan(v, &total);
if (v) out[offsets[blockIdx.x] + ex] = (u32) i;
}
extern "C" __global__ void e3_ps_scan_values(long long n, const u32* __restrict__ vals, const u32* __restrict__ offsets, u32* __restrict__ out)
{
long long i = (long long) blockIdx.x * PS_SCAN_BLOCK + threadIdx.x;
u32 total;
u32 ex = ps_block_excl_scan(i < n ? vals[i] : 0u, &total);
if (i < n) out[i] = offsets[blockIdx.x] + ex;
}
@@ -105,6 +105,8 @@ pub struct DeviceVcycle {
pub(crate) levels: Vec<DevLevel>, pub(crate) levels: Vec<DevLevel>,
sweeps: usize, sweeps: usize,
fine_cells: Vec<u32>, fine_cells: Vec<u32>,
/// R6-3: the fine lists were built on the device (`fine_cells` is read back on demand).
fine_cells_stale: bool,
r_f32: Vec<f32>, r_f32: Vec<f32>,
z_f32: Vec<f32>, z_f32: Vec<f32>,
} }
@@ -162,6 +164,7 @@ impl DeviceVcycle {
levels: dev, levels: dev,
sweeps, sweeps,
fine_cells: levels[0].cells.clone(), fine_cells: levels[0].cells.clone(),
fine_cells_stale: false,
r_f32: vec![0.0; n0], r_f32: vec![0.0; n0],
z_f32: vec![0.0; n0], z_f32: vec![0.0; n0],
} }
@@ -201,6 +204,7 @@ impl DeviceVcycle {
lv.ab = up_f(&level0.ab); lv.ab = up_f(&level0.ab);
lv.ap = up_f(&level0.ap); lv.ap = up_f(&level0.ap);
self.fine_cells = level0.cells.clone(); self.fine_cells = level0.cells.clone();
self.fine_cells_stale = false;
} }
/// P1-5 (c): as [`Self::refresh_fine`] with only the rows of `rows` /// P1-5 (c): as [`Self::refresh_fine`] with only the rows of `rows`
@@ -238,6 +242,7 @@ impl DeviceVcycle {
super::device_cg::scatter_f32(rows, &pick_f(&level0.ab), &mut lv.ab); 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); super::device_cg::scatter_f32(rows, &pick_f(&level0.ap), &mut lv.ap);
self.fine_cells = level0.cells.clone(); self.fine_cells = level0.cells.clone();
self.fine_cells_stale = false;
} }
/// Zero the finest level's right-hand side (a gather writes only the /// Zero the finest level's right-hand side (a gather writes only the
@@ -408,9 +413,25 @@ impl DeviceVcycle {
} }
} }
/// R6-3: the finest level's lists were written on the device.
pub(crate) fn fine_cells_on_device(&mut self) {
self.fine_cells_stale = true;
}
/// `z = M⁻¹ r` on the active cells (upload, V-cycle, download). /// `z = M⁻¹ r` on the active cells (upload, V-cycle, download).
pub fn apply(&mut self, r: &[f64], z: &mut [f64]) { pub fn apply(&mut self, r: &[f64], z: &mut [f64]) {
let rt = runtime(); let rt = runtime();
if self.fine_cells_stale {
let l0 = &self.levels[0];
self.fine_cells = if l0.n_cells == 0 {
Vec::new()
} else {
rt.stream
.memcpy_dtov(&l0.cells.slice(0..l0.n_cells))
.expect("fine cells")
};
self.fine_cells_stale = false;
}
for (dst, &src) in self.r_f32.iter_mut().zip(r) { for (dst, &src) in self.r_f32.iter_mut().zip(r) {
*dst = src as f32; *dst = src as f32;
} }
@@ -13,6 +13,9 @@ use std::sync::{Arc, OnceLock};
const KERNELS: &str = include_str!("../../../../kernels/cuda/e3_cg.cu"); const KERNELS: &str = include_str!("../../../../kernels/cuda/e3_cg.cu");
mod setup;
pub use setup::{OperatorGrid, OperatorInputs, enabled as poisson_device_enabled};
struct CgKernels { struct CgKernels {
_module: Arc<CudaModule>, _module: Arc<CudaModule>,
scat_f64: CudaFunction, scat_f64: CudaFunction,
@@ -188,6 +191,8 @@ pub struct DeviceCg {
/// Its export, kept likewise. /// Its export, kept likewise.
fine_export: Option<super::export::LevelExport>, fine_export: Option<super::export::LevelExport>,
active_host: Vec<bool>, active_host: Vec<bool>,
/// R6-3: the device setup's scratch (`RTX_E3_POISSON_DEVICE=1`).
pset: Option<Box<setup::PsState>>,
max_iterations: usize, max_iterations: usize,
scalar_host: Vec<f64>, scalar_host: Vec<f64>,
} }
@@ -275,6 +280,7 @@ impl DeviceCg {
fine: None, fine: None,
fine_export: None, fine_export: None,
active_host: fine.active.clone(), active_host: fine.active.clone(),
pset: None,
max_iterations: params.max_iterations, max_iterations: params.max_iterations,
scalar_host: vec![0.0], scalar_host: vec![0.0],
} }
@@ -425,6 +431,9 @@ impl DeviceCg {
}); });
self.singular = singular_count > 0; self.singular = singular_count > 0;
self.active_host = fine.active.clone(); self.active_host = fine.active.clone();
if let Some(ps) = self.pset.as_mut() {
ps.live = false;
}
let l_upload = lap.elapsed(); let l_upload = lap.elapsed();
self.key = None; self.key = None;
let l_key = lap.elapsed(); let l_key = lap.elapsed();
@@ -717,7 +726,7 @@ impl DeviceCg {
.expect("e3_cg_copy"); .expect("e3_cg_copy");
} }
self.project_mean(Which::B); self.project_mean(Which::B);
let anchor = anchor.filter(|&a| a < self.n && self.active_host[a]); let anchor = anchor.filter(|&a| a < self.n && self.cell_active(a));
let finish = |this: &mut Self, p: &mut CudaSlice<f64>, iterations: usize, residual: f64| { let finish = |this: &mut Self, p: &mut CudaSlice<f64>, iterations: usize, residual: f64| {
if this.singular { if this.singular {
let shift = match anchor { let shift = match anchor {
@@ -0,0 +1,770 @@
//! R6-3: the moving path's Poisson setup on the device (`e3_pset.cu`),
//! behind `RTX_E3_POISSON_DEVICE=1` (default off). From the projection
//! tables `DeviceCut` holds (step apertures, open flags, activity, owners,
//! the merged cells' CSR) the operator (`Solver::poisson_operator` with
//! `merge_small_cells`), the fine level (`Level::new`: activity, masked
//! couplings, z links, the cell / colour lists), the link CSR, the
//! components' count and singular count (`Components::summary_planes`) and
//! the fine export (`export_fine_from`) are built into the CG's and the
//! V-cycle's persistent buffers; no host level, no host operator.
//! `RTX_E3_BAND_CHECK=1`: every structure against the host build, bit for bit.
use super::super::device::{cfg, load_module, runtime};
use super::super::{Components, Level, Problem};
use super::DeviceCg;
use cudarc::driver::{
CudaFunction, CudaModule, CudaSlice, DevicePtr, DeviceRepr, LaunchConfig, PushKernelArg,
ValidAsZeroBits,
};
use std::sync::{Arc, OnceLock};
use std::time::Instant;
const KERNELS: &str = include_str!("../../../../../kernels/cuda/e3_pset.cu");
const SCAN_BLOCK: usize = 1024;
struct PsKernels {
_module: Arc<CudaModule>,
assemble: CudaFunction,
links: CudaFunction,
extra: CudaFunction,
merge_zero: CudaFunction,
link_counts: CudaFunction,
diag: CudaFunction,
mask: CudaFunction,
flags: CudaFunction,
uf_init: CudaFunction,
uf_hook: CudaFunction,
uf_dirichlet: CudaFunction,
uf_count: CudaFunction,
count_flags: CudaFunction,
count_vals: CudaFunction,
top: CudaFunction,
compact: CudaFunction,
values: CudaFunction,
}
static PS: OnceLock<PsKernels> = OnceLock::new();
fn kernels() -> &'static PsKernels {
PS.get_or_init(|| {
let module = load_module(KERNELS, "e3_pset.cu", true);
let f = |name: &str| module.load_function(name).expect(name);
PsKernels {
assemble: f("e3_ps_assemble"),
links: f("e3_ps_links"),
extra: f("e3_ps_extra"),
merge_zero: f("e3_ps_merge_zero"),
link_counts: f("e3_ps_link_counts"),
diag: f("e3_ps_diag"),
mask: f("e3_ps_mask"),
flags: f("e3_ps_flags"),
uf_init: f("e3_ps_uf_init"),
uf_hook: f("e3_ps_uf_hook"),
uf_dirichlet: f("e3_ps_uf_dirichlet"),
uf_count: f("e3_ps_uf_count"),
count_flags: f("e3_ps_scan_count_flags"),
count_vals: f("e3_ps_scan_count_vals"),
top: f("e3_ps_scan_top"),
compact: f("e3_ps_scan_compact"),
values: f("e3_ps_scan_values"),
_module: module,
}
})
}
/// `RTX_E3_POISSON_DEVICE=1` (default off).
#[must_use]
pub fn enabled() -> bool {
std::env::var("RTX_E3_POISSON_DEVICE").is_ok_and(|v| v == "1")
}
/// `struct PsGrid` in e3_pset.cu.
#[repr(C)]
#[derive(Clone, Copy)]
struct PsGrid {
nx: i32,
ny: i32,
nz: i32,
periodic: i32,
out: [i32; 6],
ae_int: f64,
an_int: f64,
at_int: f64,
ae_out: f64,
an_out: f64,
at_out: f64,
}
unsafe impl DeviceRepr for PsGrid {}
unsafe impl ValidAsZeroBits for PsGrid {}
#[repr(C)]
#[derive(Clone, Copy)]
struct PsOp {
ptrs: [u64; 7],
}
unsafe impl DeviceRepr for PsOp {}
unsafe impl ValidAsZeroBits for PsOp {}
#[repr(C)]
#[derive(Clone, Copy)]
struct PsIn {
ptrs: [u64; 10],
}
unsafe impl DeviceRepr for PsIn {}
unsafe impl ValidAsZeroBits for PsIn {}
/// The grid and the boundary sides of the operator (`poisson_operator`'s inputs besides the mask).
pub struct OperatorGrid {
pub nx: usize,
pub ny: usize,
pub nz: usize,
pub periodic: bool,
/// Pressure outlet on x0, x1, y0, y1, z0, z1.
pub outlet: [bool; 6],
pub dx: f64,
pub dy: f64,
pub dz: f64,
pub dt: f64,
}
/// The projection tables of the step (DeviceCut's projection set).
pub struct OperatorInputs<'a> {
pub cell_active: &'a CudaSlice<i32>,
pub open: [&'a CudaSlice<i32>; 3],
pub a: [&'a CudaSlice<f64>; 3],
pub owner: &'a CudaSlice<u32>,
pub fold_ptr: &'a CudaSlice<u32>,
pub fold_idx: &'a CudaSlice<u32>,
}
/// The device setup's persistent scratch (per operator size).
pub(super) struct PsState {
extra: CudaSlice<f64>,
act: CudaSlice<u8>,
small: CudaSlice<u8>,
pub(super) lact: CudaSlice<u8>,
flags: CudaSlice<u8>,
sums: CudaSlice<u32>,
counts: CudaSlice<u32>,
parent: CudaSlice<u32>,
dir: CudaSlice<u8>,
uf_counts: CudaSlice<u32>,
small_list: CudaSlice<u32>,
/// The CG's operator is the device-built one (the anchor's activity is read from `lact`).
pub(super) live: bool,
}
impl PsState {
fn new(n: usize) -> Self {
let s = &runtime().stream;
let z8 = |m: usize| s.alloc_zeros::<u8>(m.max(1)).expect("alloc");
let z32 = |m: usize| s.alloc_zeros::<u32>(m.max(1)).expect("alloc");
Self {
extra: s.alloc_zeros::<f64>(n).expect("alloc"),
act: z8(n),
small: z8(n),
lact: z8(n),
flags: z8(n),
sums: z32(n.div_ceil(SCAN_BLOCK) + 1),
counts: z32(n),
parent: z32(n),
dir: z8(n),
uf_counts: z32(2),
small_list: z32(1),
live: false,
}
}
}
fn scan_cfg(n: usize) -> LaunchConfig {
LaunchConfig {
grid_dim: (n.div_ceil(SCAN_BLOCK).max(1) as u32, 1, 1),
block_dim: (SCAN_BLOCK as u32, 1, 1),
shared_mem_bytes: 0,
}
}
fn one_block() -> LaunchConfig {
LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (SCAN_BLOCK as u32, 1, 1),
shared_mem_bytes: 0,
}
}
/// Compact the first `n` flags (ascending indices) into `out` (grown to `n` when short); the count.
fn compact(
flags: &CudaSlice<u8>,
sums: &mut CudaSlice<u32>,
n: usize,
out: &mut CudaSlice<u32>,
) -> usize {
let rt = runtime();
let k = kernels();
let nb = n.div_ceil(SCAN_BLOCK).max(1);
let n64 = n as i64;
let nb32 = nb as i32;
unsafe {
rt.stream
.launch_builder(&k.count_flags)
.arg(&n64)
.arg(flags)
.arg(&mut *sums)
.launch(scan_cfg(n))
.expect("e3_ps_scan_count_flags");
rt.stream
.launch_builder(&k.top)
.arg(&nb32)
.arg(&mut *sums)
.launch(one_block())
.expect("e3_ps_scan_top");
}
let total = rt
.stream
.memcpy_dtov(&sums.slice(nb..nb + 1))
.expect("count")[0] as usize;
if out.len() < total.max(1) {
*out = rt.stream.alloc_zeros::<u32>(n.max(1)).expect("alloc");
}
unsafe {
rt.stream
.launch_builder(&k.compact)
.arg(&n64)
.arg(flags)
.arg(&*sums)
.arg(&mut *out)
.launch(scan_cfg(n))
.expect("e3_ps_scan_compact");
}
total
}
impl DeviceCg {
/// Whether the anchor cell is active in the current operator.
pub(super) fn cell_active(&self, a: usize) -> bool {
match &self.pset {
Some(ps) if ps.live => {
let v = runtime()
.stream
.memcpy_dtov(&ps.lact.slice(a..a + 1))
.expect("lact");
v[0] != 0
}
_ => self.active_host[a],
}
}
/// R6-3: the fine operator, level, lists, links, components and export
/// built on the device from the step's projection tables (the multigrid
/// hierarchy's coarser levels kept, as [`Self::refresh`]). `reference`:
/// the host operator of the same step, compared structure by structure
/// (`RTX_E3_BAND_CHECK=1`).
pub fn refresh_device(
&mut self,
grid: &OperatorGrid,
inp: &OperatorInputs<'_>,
reference: Option<Problem>,
t: f64,
) {
let rt = runtime();
let k = kernels();
let s = &rt.stream;
let profile = std::env::var("RTX_E3_MOVING_PROFILE").is_ok();
let lap = Instant::now();
let (nx, ny, nz) = (grid.nx, grid.ny, grid.nz);
let n = nx * ny * nz;
assert_eq!(n, self.n, "refresh_device: grid size");
let (dx, dy, dz, dt) = (grid.dx, grid.dy, grid.dz, grid.dt);
let o = |b: bool| i32::from(b);
let pg = PsGrid {
nx: nx as i32,
ny: ny as i32,
nz: nz as i32,
periodic: o(grid.periodic),
out: grid.outlet.map(o),
ae_int: dt * (dy * dz) / dx,
an_int: dt * (dx * dz) / dy,
at_int: dt * (dx * dy) / dz,
ae_out: dt * (dy * dz) / (0.5 * dx),
an_out: dt * (dx * dz) / (0.5 * dy),
at_out: dt * (dx * dy) / (0.5 * dz),
};
let mut ps = match self.pset.take() {
Some(p) => p,
None => Box::new(PsState::new(n)),
};
let pf = |x: &CudaSlice<f64>| x.device_ptr(s).0;
let pi = |x: &CudaSlice<i32>| x.device_ptr(s).0;
let pu = |x: &CudaSlice<u32>| x.device_ptr(s).0;
let op = PsOp {
ptrs: [
pf(&self.ae),
pf(&self.aw),
pf(&self.an),
pf(&self.as_),
pf(&self.at),
pf(&self.ab),
pf(&ps.extra),
],
};
let pin = PsIn {
ptrs: [
pi(inp.cell_active),
pi(inp.open[0]),
pi(inp.open[1]),
pi(inp.open[2]),
pf(inp.a[0]),
pf(inp.a[1]),
pf(inp.a[2]),
pu(inp.owner),
pu(inp.fold_ptr),
pu(inp.fold_idx),
],
};
let n64 = n as i64;
let sync = |on: bool| {
if on {
s.synchronize().expect("sync");
}
};
// 1. The raw operator; the small cells; their links; the masters' Dirichlet parts; the zeroing.
unsafe {
s.launch_builder(&k.assemble)
.arg(&pg)
.arg(&pin)
.arg(&op)
.arg(&mut ps.act)
.arg(&mut ps.small)
.launch(cfg(n))
.expect("e3_ps_assemble");
}
let n_small = compact(&ps.small, &mut ps.sums, n, &mut ps.small_list);
let mut links: Vec<(u32, u32, f64)> = Vec::new();
if n_small > 0 {
let slots = 6 * n_small;
let mut la = s.alloc_zeros::<u32>(slots).expect("alloc");
let mut lb = s.alloc_zeros::<u32>(slots).expect("alloc");
let mut lc = s.alloc_zeros::<f64>(slots).expect("alloc");
let mut lv = s.alloc_zeros::<u8>(slots).expect("alloc");
let ns = n_small as i32;
unsafe {
s.launch_builder(&k.links)
.arg(&pg)
.arg(&ns)
.arg(&ps.small_list)
.arg(&ps.small)
.arg(inp.owner)
.arg(&op)
.arg(&mut la)
.arg(&mut lb)
.arg(&mut lc)
.arg(&mut lv)
.launch(cfg(n_small))
.expect("e3_ps_links");
}
let (ha, hb, hc, hv) = (
s.memcpy_dtov(&la).expect("la"),
s.memcpy_dtov(&lb).expect("lb"),
s.memcpy_dtov(&lc).expect("lc"),
s.memcpy_dtov(&lv).expect("lv"),
);
links = (0..slots)
.filter(|&q| hv[q] != 0)
.map(|q| (ha[q], hb[q], hc[q]))
.collect();
}
unsafe {
s.launch_builder(&k.extra)
.arg(&n64)
.arg(&pin)
.arg(&ps.small)
.arg(&op)
.launch(cfg(n))
.expect("e3_ps_extra");
s.launch_builder(&k.merge_zero)
.arg(&pg)
.arg(&ps.small)
.arg(&op)
.arg(&mut ps.act)
.launch(cfg(n))
.expect("e3_ps_merge_zero");
}
sync(profile);
let l_assemble = lap.elapsed();
// 2. The link CSR: per cell the links in the host's order (`Problem::link_map`).
let mut entries: Vec<(u32, u32, f64)> = Vec::with_capacity(2 * links.len());
for &(a, b, c) in &links {
entries.push((a, b, c));
entries.push((b, a, c));
}
entries.sort_by_key(|e| e.0); // stable: the link order within a cell
let mut cell_of: Vec<u32> = Vec::new();
let mut count_of: Vec<u32> = Vec::new();
for e in &entries {
if cell_of.last() == Some(&e.0) {
*count_of.last_mut().expect("count") += 1;
} else {
cell_of.push(e.0);
count_of.push(1);
}
}
if self.link_ptr.len() != n + 1 {
self.link_ptr = s.alloc_zeros::<u32>(n + 1).expect("alloc");
}
s.memset_zeros(&mut ps.counts).expect("counts");
if !cell_of.is_empty() {
let m = cell_of.len() as i32;
let d_cell = s.memcpy_stod(&cell_of).expect("cells");
let d_count = s.memcpy_stod(&count_of).expect("counts");
unsafe {
s.launch_builder(&k.link_counts)
.arg(&m)
.arg(&d_cell)
.arg(&d_count)
.arg(&mut ps.counts)
.launch(cfg(cell_of.len()))
.expect("e3_ps_link_counts");
}
}
let nb32 = n.div_ceil(SCAN_BLOCK).max(1) as i32;
unsafe {
s.launch_builder(&k.count_vals)
.arg(&n64)
.arg(&ps.counts)
.arg(&mut ps.sums)
.launch(scan_cfg(n))
.expect("e3_ps_scan_count_vals");
s.launch_builder(&k.top)
.arg(&nb32)
.arg(&mut ps.sums)
.launch(one_block())
.expect("e3_ps_scan_top");
s.launch_builder(&k.values)
.arg(&n64)
.arg(&ps.counts)
.arg(&ps.sums)
.arg(&mut self.link_ptr)
.launch(scan_cfg(n))
.expect("e3_ps_scan_values");
}
s.memcpy_htod(
&[entries.len() as u32][..],
&mut self.link_ptr.slice_mut(n..n + 1),
)
.expect("link_ptr end");
let idx: Vec<u32> = entries.iter().map(|e| e.1).collect();
let coef: Vec<f64> = entries.iter().map(|e| e.2).collect();
self.link_idx = s
.memcpy_stod(if idx.is_empty() { &[0u32][..] } else { &idx })
.expect("link_idx");
self.link_coef = s
.memcpy_stod(if coef.is_empty() {
&[0.0f64][..]
} else {
&coef
})
.expect("link_coef");
sync(profile);
let l_links = lap.elapsed();
// 3. The level: the diagonal, the activity, the masked couplings, the export.
let has_links = i32::from(!links.is_empty());
let fine = &mut self.vcycle.levels[0];
unsafe {
s.launch_builder(&k.diag)
.arg(&n64)
.arg(&op)
.arg(&ps.act)
.arg(&has_links)
.arg(&self.link_ptr)
.arg(&self.link_coef)
.arg(&mut self.ap)
.arg(&mut ps.lact)
.launch(cfg(n))
.expect("e3_ps_diag");
s.launch_builder(&k.mask)
.arg(&pg)
.arg(&op)
.arg(&ps.lact)
.arg(&self.ap)
.arg(&mut self.top)
.arg(&mut self.bot)
.arg(&mut fine.ae)
.arg(&mut fine.aw)
.arg(&mut fine.an)
.arg(&mut fine.as_)
.arg(&mut fine.at)
.arg(&mut fine.ab)
.arg(&mut fine.ap)
.arg(&mut fine.top)
.arg(&mut fine.bot)
.arg(&mut fine.coarse_of)
.launch(cfg(n))
.expect("e3_ps_mask");
}
sync(profile);
let l_level = lap.elapsed();
// 4. The lists (ascending; the colours by (i + j + k) % 2).
let mut counts3 = [0usize; 3];
for which in 0..3i32 {
unsafe {
s.launch_builder(&k.flags)
.arg(&pg)
.arg(&ps.lact)
.arg(&which)
.arg(&mut ps.flags)
.launch(cfg(n))
.expect("e3_ps_flags");
}
let out = match which {
0 => &mut self.cells,
1 => &mut fine.red,
_ => &mut fine.black,
};
counts3[which as usize] = compact(&ps.flags, &mut ps.sums, n, out);
}
if fine.cells.len() < counts3[0].max(1) {
fine.cells = s.alloc_zeros::<u32>(n.max(1)).expect("alloc");
}
if counts3[0] > 0 {
s.memcpy_dtod(
&self.cells.slice(0..counts3[0]),
&mut fine.cells.slice_mut(0..counts3[0]),
)
.expect("cells");
}
fine.n_cells = counts3[0];
fine.n_red = counts3[1];
fine.n_black = counts3[2];
self.vcycle.fine_cells_on_device();
self.n_cells = counts3[0];
self.n_blocks = self.n_cells.div_ceil(256).max(1);
if self.partial.len() < self.n_blocks {
self.partial = s.alloc_zeros::<f64>(self.n_blocks).expect("alloc");
}
sync(profile);
let l_lists = lap.elapsed();
// 5. The components: count and singular count.
unsafe {
s.launch_builder(&k.uf_init)
.arg(&n64)
.arg(&mut ps.parent)
.launch(cfg(n))
.expect("e3_ps_uf_init");
s.launch_builder(&k.uf_hook)
.arg(&pg)
.arg(&op)
.arg(&ps.act)
.arg(&self.link_ptr)
.arg(&self.link_idx)
.arg(&self.link_coef)
.arg(&mut ps.parent)
.launch(cfg(n))
.expect("e3_ps_uf_hook");
}
s.memset_zeros(&mut ps.dir).expect("dir");
s.memset_zeros(&mut ps.uf_counts).expect("uf counts");
unsafe {
s.launch_builder(&k.uf_dirichlet)
.arg(&n64)
.arg(&ps.act)
.arg(&ps.extra)
.arg(&mut ps.parent)
.arg(&mut ps.dir)
.launch(cfg(n))
.expect("e3_ps_uf_dirichlet");
s.launch_builder(&k.uf_count)
.arg(&n64)
.arg(&ps.act)
.arg(&ps.parent)
.arg(&ps.dir)
.arg(&mut ps.uf_counts)
.launch(cfg(n))
.expect("e3_ps_uf_count");
}
let uc = s.memcpy_dtov(&ps.uf_counts).expect("uf counts");
let (n_components, singular_count) = (uc[0] as usize, uc[1] as usize);
let l_components = lap.elapsed();
assert!(
singular_count == 0 || n_components == 1,
"DeviceCg::refresh_device: {n_components} components with {singular_count} singular"
);
self.singular = singular_count > 0;
self.key = None;
self.fine = None;
self.fine_export = None;
ps.live = true;
if profile {
let ms = |d: std::time::Duration| d.as_secs_f64() * 1e3;
eprintln!(
" device poisson laps: total {:.0} ms (assemble + merge {:.0}, links {:.0}, level {:.0}, lists {:.0}, components {:.0}); {} small, {} links, {} cells, {} components",
ms(l_components),
ms(l_assemble),
ms(l_links - l_assemble),
ms(l_level - l_links),
ms(l_lists - l_level),
ms(l_components - l_lists),
n_small,
links.len(),
counts3[0],
n_components
);
}
if let Some(problem) = reference {
self.check_device_setup(problem, &ps, n_components, singular_count, t);
}
self.pset = Some(ps);
}
/// `RTX_E3_BAND_CHECK=1`: the device setup against the host build of the
/// same step (operator → `Level::new` → `summary_planes`, the link CSR of
/// `refresh_with`, `export_fine_from`), bit for bit.
fn check_device_setup(
&self,
problem: Problem,
ps: &PsState,
n_components: usize,
singular_count: usize,
t: f64,
) {
let s = &runtime().stream;
let n = self.n;
let full = Level::<f64>::new(problem);
let (host_nc, host_flags) = Components::summary_planes(&full.problem, &full.cells);
let host_singular = host_flags.iter().filter(|&&f| f).count();
let map = full.problem.link_map();
let mut link_ptr = Vec::with_capacity(n + 1);
let mut link_idx = Vec::new();
let mut link_coef = Vec::new();
link_ptr.push(0u32);
for idx in 0..n {
if let Some(list) = map.get(&idx) {
for &(other, c) in list {
link_idx.push(other as u32);
link_coef.push(c);
}
}
link_ptr.push(link_idx.len() as u32);
}
let export = super::super::export::export_fine_from(&full);
let to_u32 = |v: &[usize]| {
v.iter()
.map(|&i| if i == usize::MAX { u32::MAX } else { i as u32 })
.collect::<Vec<u32>>()
};
let down_f = |x: &CudaSlice<f64>| s.memcpy_dtov(&x.slice(0..n)).expect("down");
let down_f32 = |x: &CudaSlice<f32>| s.memcpy_dtov(&x.slice(0..n)).expect("down");
let down_u = |x: &CudaSlice<u32>, m: usize| -> Vec<u32> {
if m == 0 {
Vec::new()
} else {
s.memcpy_dtov(&x.slice(0..m)).expect("down")
}
};
let down_b = |x: &CudaSlice<u8>| {
s.memcpy_dtov(&x.slice(0..n))
.expect("down")
.into_iter()
.map(|v| v != 0)
.collect::<Vec<bool>>()
};
let same_f = |a: &[f64], b: &[f64]| {
a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
};
let same_f32 = |a: &[f32], b: &[f32]| {
a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
};
let mut differs: Vec<&str> = Vec::new();
let mut want = |name: &'static str, ok: bool| {
if !ok {
differs.push(name);
}
};
want("level active", down_b(&ps.lact) == full.active);
want("problem active", down_b(&ps.act) == full.problem.active);
want(
"extra",
same_f(&down_f(&ps.extra), &full.problem.extra_diag),
);
want(
"cells",
down_u(&self.cells, self.n_cells) == to_u32(&full.cells),
);
want("top", down_u(&self.top, n) == to_u32(&full.top));
want("bot", down_u(&self.bot, n) == to_u32(&full.bot));
for (name, dev, host) in [
("ae", &self.ae, &full.ae),
("aw", &self.aw, &full.aw),
("an", &self.an, &full.an),
("as", &self.as_, &full.as_),
("at", &self.at, &full.at),
("ab", &self.ab, &full.ab),
("ap", &self.ap, &full.ap),
] {
want(name, same_f(&down_f(dev), host));
}
want("link_ptr", down_u(&self.link_ptr, n + 1) == link_ptr);
if !link_idx.is_empty() {
want(
"link_idx",
down_u(&self.link_idx, link_idx.len()) == link_idx,
);
want(
"link_coef",
same_f(
&s.memcpy_dtov(&self.link_coef.slice(0..link_coef.len()))
.expect("down"),
&link_coef,
),
);
} else {
want("links present", self.link_idx.len() == 1);
}
want(
"components",
n_components == host_nc && singular_count == host_singular,
);
let fine = &self.vcycle.levels[0];
want(
"export cells",
down_u(&fine.cells, fine.n_cells) == export.cells,
);
want("export red", down_u(&fine.red, fine.n_red) == export.red);
want(
"export black",
down_u(&fine.black, fine.n_black) == export.black,
);
want("export top", down_u(&fine.top, n) == export.top);
want("export bot", down_u(&fine.bot, n) == export.bot);
want(
"export coarse_of",
down_u(&fine.coarse_of, n) == export.coarse_of,
);
for (name, dev, host) in [
("export ae", &fine.ae, &export.ae),
("export aw", &fine.aw, &export.aw),
("export an", &fine.an, &export.an),
("export as", &fine.as_, &export.as_),
("export at", &fine.at, &export.at),
("export ab", &fine.ab, &export.ab),
("export ap", &fine.ap, &export.ap),
] {
want(name, same_f32(&down_f32(dev), host));
}
if differs.is_empty() {
eprintln!(
" poisson check t {t:.6}: device operator/level/lists/links/components/export IDENTICAL to the host build ({} cells, {} links, {} components, {} singular)",
full.cells.len(),
full.problem.links.len(),
host_nc,
host_singular
);
} else {
eprintln!(
" poisson check t {t:.6}: DIFFERS in {differs:?} (device {n_components}/{singular_count} components, host {host_nc}/{host_singular})"
);
if std::env::var("RTX_E3_GEOM_CHECK_SOFT").is_err() {
panic!("R6-3 poisson check: device setup differs from the host build: {differs:?}");
}
}
}
}
@@ -8,6 +8,7 @@
mod cut; mod cut;
mod geom; mod geom;
mod mask; mod mask;
mod poisson_setup;
use super::{Side, Solver, StepResult}; use super::{Side, Solver, StepResult};
use crate::solvers::incompressible::embedded3::field::Field; use crate::solvers::incompressible::embedded3::field::Field;
@@ -7,7 +7,7 @@
use super::{DeviceStep, E3Params, E3Ptrs, StepResult}; use super::{DeviceStep, E3Params, E3Ptrs, StepResult};
use crate::solvers::incompressible::embedded3::field::Field; use crate::solvers::incompressible::embedded3::field::Field;
use crate::solvers::incompressible::embedded3::poisson::device::{cfg, load_module, runtime}; use crate::solvers::incompressible::embedded3::poisson::device::{cfg, load_module, runtime};
use crate::solvers::incompressible::embedded3::poisson::device_cg::DeviceCg; use crate::solvers::incompressible::embedded3::poisson::device_cg::{DeviceCg, OperatorInputs};
use crate::solvers::incompressible::embedded3::poisson::GuessBasis; use crate::solvers::incompressible::embedded3::poisson::GuessBasis;
use crate::solvers::incompressible::embedded3::step::Solver; use crate::solvers::incompressible::embedded3::step::Solver;
use crate::solvers::incompressible::embedded3::wall::FaceKind; use crate::solvers::incompressible::embedded3::wall::FaceKind;
@@ -738,6 +738,19 @@ impl DeviceCut {
same_u("fold_idx", &self.fold_idx, &full.fold_idx); same_u("fold_idx", &self.fold_idx, &full.fold_idx);
} }
/// R6-3: the projection set the device Poisson setup reads.
fn poisson_inputs(&self) -> OperatorInputs<'_> {
assert_eq!(self.phase, Phase::Projection, "poisson inputs: projection phase");
OperatorInputs {
cell_active: &self.active,
open: [&self.open[0], &self.open[1], &self.open[2]],
a: [&self.a[0], &self.a[1], &self.a[2]],
owner: &self.owner,
fold_ptr: &self.fold_ptr,
fold_idx: &self.fold_idx,
}
}
fn ptrs(&self) -> E3CutPtrs { fn ptrs(&self) -> E3CutPtrs {
let rt = runtime(); let rt = runtime();
let s = &rt.stream; let s = &rt.stream;
@@ -1054,6 +1067,17 @@ impl DeviceStep {
if self.steps_since_hierarchy >= every || self.cg_dt != dt { if self.steps_since_hierarchy >= every || self.cg_dt != dt {
self.cg = None; self.cg = None;
self.steps_since_hierarchy = 0; self.steps_since_hierarchy = 0;
} else if let (Some(cg), Some(dc), Some(grid)) = (
self.cg.as_mut(),
self.cut.as_ref(),
super::poisson_setup::device_grid(&self.solver, g, dt),
) {
// R6-3 (`RTX_E3_POISSON_DEVICE=1`): the operator, level, lists,
// links, components and export on the device from the step's
// projection tables; `RTX_E3_BAND_CHECK=1` hands the host
// operator as the reference.
let reference = check.then(|| self.solver.poisson_operator(g, dt));
cg.refresh_device(&grid, &dc.poisson_inputs(), reference, t_new);
} else if let Some(cg) = self.cg.as_mut() { } else if let Some(cg) = self.cg.as_mut() {
let lap_op = Instant::now(); let lap_op = Instant::now();
let problem = self.solver.poisson_operator(g, dt); let problem = self.solver.poisson_operator(g, dt);
@@ -0,0 +1,54 @@
//! R6-3: when the moving path's Poisson setup runs on the device
//! (`RTX_E3_POISSON_DEVICE=1`, default off) and its grid description.
use crate::solvers::incompressible::embedded3::Grid;
use crate::solvers::incompressible::embedded3::poisson::device_cg::{
OperatorGrid, poisson_device_enabled,
};
use crate::solvers::incompressible::embedded3::step::{Side, Solver};
use std::sync::Once;
static FALLBACK: Once = Once::new();
/// The operator's grid and sides when the device setup applies this step;
/// `None` with the knob off or for a case it does not carry (logged once):
/// the gradient weights (`pressure_centroid`), a periodic z of two planes or
/// fewer (a cell's top and bottom neighbours coincide).
pub(super) fn device_grid(solver: &Solver, g: Grid, dt: f64) -> Option<OperatorGrid> {
if !poisson_device_enabled() {
return None;
}
let b = solver.params.boundaries;
let periodic = b.periodic_z();
let reason = match solver.mask() {
None => Some("no mask"),
Some(m) if m.grad_weights.is_some() => Some("gradient weights"),
Some(_) if periodic && g.nz <= 2 => Some("periodic z with nz <= 2"),
Some(_) => None,
};
if let Some(r) = reason {
FALLBACK.call_once(|| {
eprintln!("R6-3: RTX_E3_POISSON_DEVICE=1 falls back to the host setup ({r})")
});
return None;
}
let out = |s: Side| s == Side::PressureOutlet;
Some(OperatorGrid {
nx: g.nx,
ny: g.ny,
nz: g.nz,
periodic,
outlet: [
out(b.x0),
out(b.x1),
out(b.y0),
out(b.y1),
out(b.z0),
out(b.z1),
],
dx: g.dx,
dy: g.dy,
dz: g.dz,
dt,
})
}