diff --git a/crates/specialized/rtx-cfd/src/kernels/cuda/e3_pset.cu b/crates/specialized/rtx-cfd/src/kernels/cuda/e3_pset.cu new file mode 100644 index 0000000..6a5636e --- /dev/null +++ b/crates/specialized/rtx-cfd/src/kernels/cuda/e3_pset.cu @@ -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; +} 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 9eafae2..928d6a4 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 @@ -105,6 +105,8 @@ pub struct DeviceVcycle { pub(crate) levels: Vec, sweeps: usize, fine_cells: Vec, + /// R6-3: the fine lists were built on the device (`fine_cells` is read back on demand). + fine_cells_stale: bool, r_f32: Vec, z_f32: Vec, } @@ -162,6 +164,7 @@ impl DeviceVcycle { levels: dev, sweeps, fine_cells: levels[0].cells.clone(), + fine_cells_stale: false, r_f32: vec![0.0; n0], z_f32: vec![0.0; n0], } @@ -201,6 +204,7 @@ impl DeviceVcycle { lv.ab = up_f(&level0.ab); lv.ap = up_f(&level0.ap); self.fine_cells = level0.cells.clone(); + self.fine_cells_stale = false; } /// 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.ap), &mut lv.ap); self.fine_cells = level0.cells.clone(); + self.fine_cells_stale = false; } /// 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). pub fn apply(&mut self, r: &[f64], z: &mut [f64]) { 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) { *dst = src as f32; } 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 e2377fb..c9705b8 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 @@ -13,6 +13,9 @@ use std::sync::{Arc, OnceLock}; const KERNELS: &str = include_str!("../../../../kernels/cuda/e3_cg.cu"); +mod setup; +pub use setup::{OperatorGrid, OperatorInputs, enabled as poisson_device_enabled}; + struct CgKernels { _module: Arc, scat_f64: CudaFunction, @@ -188,6 +191,8 @@ pub struct DeviceCg { /// Its export, kept likewise. fine_export: Option, active_host: Vec, + /// R6-3: the device setup's scratch (`RTX_E3_POISSON_DEVICE=1`). + pset: Option>, max_iterations: usize, scalar_host: Vec, } @@ -275,6 +280,7 @@ impl DeviceCg { fine: None, fine_export: None, active_host: fine.active.clone(), + pset: None, max_iterations: params.max_iterations, scalar_host: vec![0.0], } @@ -425,6 +431,9 @@ impl DeviceCg { }); self.singular = singular_count > 0; self.active_host = fine.active.clone(); + if let Some(ps) = self.pset.as_mut() { + ps.live = false; + } let l_upload = lap.elapsed(); self.key = None; let l_key = lap.elapsed(); @@ -717,7 +726,7 @@ impl DeviceCg { .expect("e3_cg_copy"); } 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, iterations: usize, residual: f64| { if this.singular { let shift = match anchor { diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device_cg/setup.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device_cg/setup.rs new file mode 100644 index 0000000..9c0b91f --- /dev/null +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/poisson/device_cg/setup.rs @@ -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, + 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 = 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, + pub open: [&'a CudaSlice; 3], + pub a: [&'a CudaSlice; 3], + pub owner: &'a CudaSlice, + pub fold_ptr: &'a CudaSlice, + pub fold_idx: &'a CudaSlice, +} + +/// The device setup's persistent scratch (per operator size). +pub(super) struct PsState { + extra: CudaSlice, + act: CudaSlice, + small: CudaSlice, + pub(super) lact: CudaSlice, + flags: CudaSlice, + sums: CudaSlice, + counts: CudaSlice, + parent: CudaSlice, + dir: CudaSlice, + uf_counts: CudaSlice, + small_list: CudaSlice, + /// 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::(m.max(1)).expect("alloc"); + let z32 = |m: usize| s.alloc_zeros::(m.max(1)).expect("alloc"); + Self { + extra: s.alloc_zeros::(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, + sums: &mut CudaSlice, + n: usize, + out: &mut CudaSlice, +) -> 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::(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, + 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| x.device_ptr(s).0; + let pi = |x: &CudaSlice| x.device_ptr(s).0; + let pu = |x: &CudaSlice| 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::(slots).expect("alloc"); + let mut lb = s.alloc_zeros::(slots).expect("alloc"); + let mut lc = s.alloc_zeros::(slots).expect("alloc"); + let mut lv = s.alloc_zeros::(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 = Vec::new(); + let mut count_of: Vec = 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::(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 = entries.iter().map(|e| e.1).collect(); + let coef: Vec = 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::(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::(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::::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::>() + }; + let down_f = |x: &CudaSlice| s.memcpy_dtov(&x.slice(0..n)).expect("down"); + let down_f32 = |x: &CudaSlice| s.memcpy_dtov(&x.slice(0..n)).expect("down"); + let down_u = |x: &CudaSlice, m: usize| -> Vec { + if m == 0 { + Vec::new() + } else { + s.memcpy_dtov(&x.slice(0..m)).expect("down") + } + }; + let down_b = |x: &CudaSlice| { + s.memcpy_dtov(&x.slice(0..n)) + .expect("down") + .into_iter() + .map(|v| v != 0) + .collect::>() + }; + 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:?}"); + } + } + } +} diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device.rs index c81dc15..52fe7f5 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device.rs @@ -8,6 +8,7 @@ mod cut; mod geom; mod mask; +mod poisson_setup; use super::{Side, Solver, StepResult}; use crate::solvers::incompressible::embedded3::field::Field; diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device/cut.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device/cut.rs index 4d1d0be..764aac0 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device/cut.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device/cut.rs @@ -7,7 +7,7 @@ use super::{DeviceStep, E3Params, E3Ptrs, StepResult}; 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_cg::DeviceCg; +use crate::solvers::incompressible::embedded3::poisson::device_cg::{DeviceCg, OperatorInputs}; use crate::solvers::incompressible::embedded3::poisson::GuessBasis; use crate::solvers::incompressible::embedded3::step::Solver; use crate::solvers::incompressible::embedded3::wall::FaceKind; @@ -738,6 +738,19 @@ impl DeviceCut { 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 { let rt = runtime(); let s = &rt.stream; @@ -1054,6 +1067,17 @@ impl DeviceStep { if self.steps_since_hierarchy >= every || self.cg_dt != dt { self.cg = None; 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() { let lap_op = Instant::now(); let problem = self.solver.poisson_operator(g, dt); diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device/poisson_setup.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device/poisson_setup.rs new file mode 100644 index 0000000..a83d4da --- /dev/null +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/embedded3/step/device/poisson_setup.rs @@ -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 { + 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, + }) +}