rtx-cfd 3D Stage 1 item 3: cg3.cu + DeviceCg3 — the f64 CG entirely on the device (run_pcg's control flow: true-residual stop and resync, breakdown guard, mean projection, anchor shift; fixed-order block reductions, no atomics; the device V-cycle as preconditioner; DevicePcgCache3); gate 3 HELD: iterations = host, |Δp| ≤ 2e-14 relative, run-to-run bit-identical, singular box anchor exact / mean 2e-15

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

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

@@ -0,0 +1,181 @@
+/**
+ * 3D Stage 1, gate 3: the f64 conjugate gradient's maps on the device over
+ * the active cells (index list), with FIXED-ORDER reductions: every block
+ * reduces its 256 lanes by the same shared-memory tree and writes one
+ * partial; `cg3_reduce` sums the partials in index order on one thread.
+ * Run-to-run bit-identical by construction (no atomics).
+ */
+#define NONE 0xFFFFFFFFu
+
+__device__ __forceinline__ double nb_sum3d(
+    int g, int nx, const double* ae, const double* aw, const double* an, const double* as_,
+    const double* at, const double* ab, const unsigned int* top, const unsigned int* bot,
+    const double* x)
+{
+    double s = 0.0;
+    double e = ae[g]; if (e != 0.0) s += e * x[g + 1];
+    double w = aw[g]; if (w != 0.0) s += w * x[g - 1];
+    double nn = an[g]; if (nn != 0.0) s += nn * x[g + nx];
+    double ss = as_[g]; if (ss != 0.0) s += ss * x[g - nx];
+    double t = at[g]; if (t != 0.0) s += t * x[top[g]];
+    double b = ab[g]; if (b != 0.0) s += b * x[bot[g]];
+    return s;
+}
+
+__device__ __forceinline__ double block_reduce(double v)
+{
+    __shared__ double sh[256];
+    int t = threadIdx.x;
+    sh[t] = v;
+    __syncthreads();
+    for (int s = 128; s > 0; s >>= 1) {
+        if (t < s) sh[t] += sh[t + s];
+        __syncthreads();
+    }
+    return sh[0];
+}
+
+/* q = A d on the cells. */
+extern "C" __global__ void cg3_spmv(
+    int n_cells, const unsigned int* __restrict__ cells,
+    const double* __restrict__ ae, const double* __restrict__ aw,
+    const double* __restrict__ an, const double* __restrict__ as_,
+    const double* __restrict__ at, const double* __restrict__ ab,
+    const unsigned int* __restrict__ top, const unsigned int* __restrict__ bot,
+    const double* __restrict__ ap, const double* __restrict__ d, double* __restrict__ q, int nx)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int g = cells[t];
+    q[g] = ap[g] * d[g] - nb_sum3d(g, nx, ae, aw, an, as_, at, ab, top, bot, d);
+}
+
+/* r = b − A p on the cells; partial[block] = Σ |r| over the block's cells. */
+extern "C" __global__ void cg3_residual(
+    int n_cells, const unsigned int* __restrict__ cells,
+    const double* __restrict__ ae, const double* __restrict__ aw,
+    const double* __restrict__ an, const double* __restrict__ as_,
+    const double* __restrict__ at, const double* __restrict__ ab,
+    const unsigned int* __restrict__ top, const unsigned int* __restrict__ bot,
+    const double* __restrict__ ap, const double* __restrict__ b,
+    const double* __restrict__ p, double* __restrict__ r, double* __restrict__ partial, int nx)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    double v = 0.0;
+    if (t < n_cells) {
+        int g = cells[t];
+        v = b[g] - (ap[g] * p[g] - nb_sum3d(g, nx, ae, aw, an, as_, at, ab, top, bot, p));
+        r[g] = v;
+        v = fabs(v);
+    }
+    double s = block_reduce(v);
+    if (threadIdx.x == 0) partial[blockIdx.x] = s;
+}
+
+/* partial[block] = Σ a·b over the block's cells (b = a for a norm). */
+extern "C" __global__ void cg3_dot_partial(
+    int n_cells, const unsigned int* __restrict__ cells,
+    const double* __restrict__ a, const double* __restrict__ b, double* __restrict__ partial)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    double v = 0.0;
+    if (t < n_cells) { int g = cells[t]; v = a[g] * b[g]; }
+    double s = block_reduce(v);
+    if (threadIdx.x == 0) partial[blockIdx.x] = s;
+}
+
+/* partial[block] = Σ |a| over the block's cells. */
+extern "C" __global__ void cg3_l1_partial(
+    int n_cells, const unsigned int* __restrict__ cells,
+    const double* __restrict__ a, double* __restrict__ partial)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    double v = 0.0;
+    if (t < n_cells) v = fabs(a[cells[t]]);
+    double s = block_reduce(v);
+    if (threadIdx.x == 0) partial[blockIdx.x] = s;
+}
+
+/* partial[block] = Σ a over the block's cells. */
+extern "C" __global__ void cg3_sum_partial(
+    int n_cells, const unsigned int* __restrict__ cells,
+    const double* __restrict__ a, double* __restrict__ partial)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    double v = 0.0;
+    if (t < n_cells) v = a[cells[t]];
+    double s = block_reduce(v);
+    if (threadIdx.x == 0) partial[blockIdx.x] = s;
+}
+
+/* out[0] = Σ partial[0..n) in index order, one thread. */
+extern "C" __global__ void cg3_reduce(int n, const double* __restrict__ partial, double* __restrict__ out)
+{
+    if (blockIdx.x * blockDim.x + threadIdx.x != 0) return;
+    double s = 0.0;
+    for (int i = 0; i < n; ++i) s += partial[i];
+    out[0] = s;
+}
+
+/* p += alpha d; r -= alpha q. */
+extern "C" __global__ void cg3_axpy2(
+    int n_cells, const unsigned int* __restrict__ cells, double alpha,
+    const double* __restrict__ d, const double* __restrict__ q,
+    double* __restrict__ p, double* __restrict__ r)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int g = cells[t];
+    p[g] += alpha * d[g];
+    r[g] -= alpha * q[g];
+}
+
+/* d = z + beta d. */
+extern "C" __global__ void cg3_xpay(
+    int n_cells, const unsigned int* __restrict__ cells, double beta,
+    const double* __restrict__ z, double* __restrict__ d)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int g = cells[t];
+    d[g] = z[g] + beta * d[g];
+}
+
+/* dst = src (copy over the cells). */
+extern "C" __global__ void cg3_copy(
+    int n_cells, const unsigned int* __restrict__ cells,
+    const double* __restrict__ src, double* __restrict__ dst)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int g = cells[t];
+    dst[g] = src[g];
+}
+
+/* v -= s over the cells. */
+extern "C" __global__ void cg3_shift(
+    int n_cells, const unsigned int* __restrict__ cells, double s, double* __restrict__ v)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    v[cells[t]] -= s;
+}
+
+/* f32 b0 = (float) r over the cells; f64 z = (double) x0 over the cells. */
+extern "C" __global__ void cg3_gather_f32(
+    int n_cells, const unsigned int* __restrict__ cells, const double* __restrict__ r, float* __restrict__ b0)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int g = cells[t];
+    b0[g] = (float) r[g];
+}
+
+extern "C" __global__ void cg3_scatter_f64(
+    int n_cells, const unsigned int* __restrict__ cells, const float* __restrict__ x0, double* __restrict__ z)
+{
+    int t = blockIdx.x * blockDim.x + threadIdx.x;
+    if (t >= n_cells) return;
+    int g = cells[t];
+    z[g] = (double) x0[g];
+}