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

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