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
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
c860783c81
commit
348f4621d1
@@ -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>,
|
||||
sweeps: usize,
|
||||
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>,
|
||||
z_f32: Vec<f32>,
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
+10
-1
@@ -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<CudaModule>,
|
||||
scat_f64: CudaFunction,
|
||||
@@ -188,6 +191,8 @@ pub struct DeviceCg {
|
||||
/// Its export, kept likewise.
|
||||
fine_export: Option<super::export::LevelExport>,
|
||||
active_host: Vec<bool>,
|
||||
/// R6-3: the device setup's scratch (`RTX_E3_POISSON_DEVICE=1`).
|
||||
pset: Option<Box<setup::PsState>>,
|
||||
max_iterations: usize,
|
||||
scalar_host: Vec<f64>,
|
||||
}
|
||||
@@ -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<f64>, iterations: usize, residual: f64| {
|
||||
if this.singular {
|
||||
let shift = match anchor {
|
||||
|
||||
+770
@@ -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 geom;
|
||||
mod mask;
|
||||
mod poisson_setup;
|
||||
|
||||
use super::{Side, Solver, StepResult};
|
||||
use crate::solvers::incompressible::embedded3::field::Field;
|
||||
|
||||
@@ -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);
|
||||
|
||||
+54
@@ -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,
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user