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

/**
 * 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];
}