rtx-cfd embedded3 item 2: e3_mg.cu + poisson::{export, device} (one shared CUDA runtime and module loader for embedded3); gate 2 HELD: device V-cycle = host f32 to 4e-7 relative on 96×40×{1,8} and 378×62×62, 4.97 ms per V-cycle at 1.45 M cells incl. transfers

Co-Authored-By: Claude Fable 5.1 <[email protected]>

crates/specialized/rtx-cfd/src/kernels/cuda/e3_mg.cu

@@ -0,0 +1,116 @@
+/**
+ * embedded3 item 2 (omni-cortex `docs/embedded3_campaign.md`): the
+ * V-cycle's maps for a MASKED, VARIABLE-COEFFICIENT seven-point operator.
+ * Per-cell arrays are [n] (one march); the index lists (cells, colours,
+ * children) drive the maps; `top`/`bot` give the neighbour above/below per
+ * cell (UINT_MAX = none; a zero coefficient is never read), so a periodic
+ * z is data. Restriction, prolongation and zero are the 2D kernels (they
+ * never touch the stencil).
+ */
+#define NONE 0xFFFFFFFFu
+
+__device__ __forceinline__ float nb_sum3(
+    int g, int nx, const float* ae, const float* aw, const float* an, const float* as_,
+    const float* at, const float* ab, const unsigned int* top, const unsigned int* bot,
+    const float* x)
+{
+    float s = 0.0f;
+    float e = ae[g]; if (e != 0.0f) s += e * x[g + 1];
+    float w = aw[g]; if (w != 0.0f) s += w * x[g - 1];
+    float nn = an[g]; if (nn != 0.0f) s += nn * x[g + nx];
+    float ss = as_[g]; if (ss != 0.0f) s += ss * x[g - nx];
+    float t = at[g]; if (t != 0.0f) s += t * x[top[g]];
+    float b = ab[g]; if (b != 0.0f) s += b * x[bot[g]];
+    return s;
+}
+
+extern "C" __global__ void e3_mg_rb_half(
+    int n_col, const unsigned int* __restrict__ col,
+    const float* __restrict__ ae, const float* __restrict__ aw,
+    const float* __restrict__ an, const float* __restrict__ as_,
+    const float* __restrict__ at, const float* __restrict__ ab,
+    const unsigned int* __restrict__ top, const unsigned int* __restrict__ bot,
+    const float* __restrict__ ap, const float* __restrict__ b,
+    float* __restrict__ x, int nx)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_col) return;
+    int g = col[t];
+    float s = nb_sum3(g, nx, ae, aw, an, as_, at, ab, top, bot, x);
+    x[g] = (b[g] + s) / ap[g];
+}
+
+extern "C" __global__ void e3_mg_residual(
+    int n_cells, const unsigned int* __restrict__ cells,
+    const float* __restrict__ ae, const float* __restrict__ aw,
+    const float* __restrict__ an, const float* __restrict__ as_,
+    const float* __restrict__ at, const float* __restrict__ ab,
+    const unsigned int* __restrict__ top, const unsigned int* __restrict__ bot,
+    const float* __restrict__ ap, const float* __restrict__ b,
+    const float* __restrict__ x, float* __restrict__ r, int nx)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int g = cells[t];
+    float s = nb_sum3(g, nx, ae, aw, an, as_, at, ab, top, bot, x);
+    r[g] = b[g] - (ap[g] * x[g] - s);
+}
+
+/* b_c[c] = sum of r_f over the children of coarse cell c (fixed order). */
+extern "C" __global__ void e3_mg_restrict(
+    int n_coarse, const unsigned int* __restrict__ coarse_cells,
+    const unsigned int* __restrict__ children_ptr, const unsigned int* __restrict__ children_idx,
+    const float* __restrict__ r_f, float* __restrict__ b_c)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_coarse) return;
+    float s = 0.0f;
+    for (unsigned int p = children_ptr[t]; p < children_ptr[t + 1]; ++p) s += r_f[children_idx[p]];
+    b_c[coarse_cells[t]] = s;
+}
+
+/* x_f += 2 x_c[coarse_of[idx]] */
+extern "C" __global__ void e3_mg_prolong(
+    int n_cells, const unsigned int* __restrict__ cells, const unsigned int* __restrict__ coarse_of,
+    float* __restrict__ x_f, const float* __restrict__ x_c)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int idx = cells[t];
+    x_f[idx] += 2.0f * x_c[coarse_of[idx]];
+}
+
+extern "C" __global__ void e3_mg_zero(int n_cells, const unsigned int* __restrict__ cells, float* __restrict__ x)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    x[cells[t]] = 0.0f;
+}
+
+/* The coarsest level: one thread, `sweeps` symmetric RED-BLACK sweeps
+ * (red, black, black, red) from zero — the host's ordering. */
+extern "C" __global__ void e3_mg_coarsest(
+    int n_cells, const unsigned int* __restrict__ cells,
+    int n_red, const unsigned int* __restrict__ red,
+    int n_black, const unsigned int* __restrict__ black,
+    const float* __restrict__ ae, const float* __restrict__ aw,
+    const float* __restrict__ an, const float* __restrict__ as_,
+    const float* __restrict__ at, const float* __restrict__ ab,
+    const unsigned int* __restrict__ top, const unsigned int* __restrict__ bot,
+    const float* __restrict__ ap, const float* __restrict__ b,
+    float* __restrict__ x, int nx, int sweeps)
+{
+    if (blockIdx.x * blockDim.x + threadIdx.x != 0) return;
+    for (int t = 0; t < n_cells; ++t) x[cells[t]] = 0.0f;
+    for (int sw = 0; sw < sweeps; ++sw) {
+        for (int half = 0; half < 4; ++half) {
+            const unsigned int* list = (half == 0 || half == 3) ? red : black;
+            int n_list = (half == 0 || half == 3) ? n_red : n_black;
+            for (int t = 0; t < n_list; ++t) {
+                int g = list[t];
+                float s = nb_sum3(g, nx, ae, aw, an, as_, at, ab, top, bot, x);
+                x[g] = (b[g] + s) / ap[g];
+            }
+        }
+    }
+}