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

/**
 * CUDA kernels for Poisson equation solving
 *
 * Implements iterative solvers for the Poisson equation:
 * ∇²φ = f
 *
 * Includes Jacobi, Gauss-Seidel, and SOR methods for 2D and 3D domains.
 * These are essential for pressure correction in CFD solvers.
 */

extern "C" __global__ void poisson_jacobi_2d(
    float* __restrict__ phi_new,
    const float* __restrict__ phi,
    const float* __restrict__ source,
    float factor,           // 1 / (2 * (1/dx² + 1/dy²))
    float dx2_inv,         // 1/dx²
    float dy2_inv,         // 1/dy²
    int nx,
    int ny
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;

    if (i >= nx || j >= ny) return;

    int idx = j * nx + i;

    // Boundary cells keep the Dirichlet values they hold.
    if (i == 0 || i == nx - 1 || j == 0 || j == ny - 1) {
        phi_new[idx] = phi[idx];
        return;
    }

    int idx_ip1 = j * nx + (i + 1);
    int idx_im1 = j * nx + (i - 1);
    int idx_jp1 = (j + 1) * nx + i;
    int idx_jm1 = (j - 1) * nx + i;

    // Jacobi iteration: φ^{k+1}_ij = factor * (source_ij + (1/dx²)(φ_{i+1,j} + φ_{i-1,j}) + (1/dy²)(φ_{i,j+1} + φ_{i,j-1}))
    phi_new[idx] = factor * (
        -source[idx] +
        dx2_inv * (phi[idx_ip1] + phi[idx_im1]) +
        dy2_inv * (phi[idx_jp1] + phi[idx_jm1])
    );
}

extern "C" __global__ void poisson_gauss_seidel_2d(
    float* __restrict__ phi,
    const float* __restrict__ source,
    float factor,
    float dx2_inv,
    float dy2_inv,
    int nx,
    int ny,
    int red_black_flag    // 0 for red, 1 for black
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;

    if (i >= nx || j >= ny) return;

    // Red-black ordering for parallelization
    if (((i + j) % 2) != red_black_flag) return;

    int idx = j * nx + i;

    // Boundary conditions
    if (i == 0 || i == nx - 1 || j == 0 || j == ny - 1) {
        return;
    }

    int idx_ip1 = j * nx + (i + 1);
    int idx_im1 = j * nx + (i - 1);
    int idx_jp1 = (j + 1) * nx + i;
    int idx_jm1 = (j - 1) * nx + i;

    // Gauss-Seidel update (in-place)
    phi[idx] = factor * (
        -source[idx] +
        dx2_inv * (phi[idx_ip1] + phi[idx_im1]) +
        dy2_inv * (phi[idx_jp1] + phi[idx_jm1])
    );
}

extern "C" __global__ void poisson_sor_2d(
    float* __restrict__ phi,
    const float* __restrict__ source,
    float factor,
    float dx2_inv,
    float dy2_inv,
    float omega,          // SOR relaxation parameter (1 < ω < 2)
    int nx,
    int ny,
    int red_black_flag
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;

    if (i >= nx || j >= ny) return;

    if (((i + j) % 2) != red_black_flag) return;

    int idx = j * nx + i;

    // Boundary conditions
    if (i == 0 || i == nx - 1 || j == 0 || j == ny - 1) {
        return;
    }

    int idx_ip1 = j * nx + (i + 1);
    int idx_im1 = j * nx + (i - 1);
    int idx_jp1 = (j + 1) * nx + i;
    int idx_jm1 = (j - 1) * nx + i;

    float phi_old = phi[idx];

    // Compute Gauss-Seidel update
    float phi_gs = factor * (
        -source[idx] +
        dx2_inv * (phi[idx_ip1] + phi[idx_im1]) +
        dy2_inv * (phi[idx_jp1] + phi[idx_jm1])
    );

    // SOR update: φ^{new} = (1-ω)φ^{old} + ω*φ^{GS}
    phi[idx] = (1.0f - omega) * phi_old + omega * phi_gs;
}

extern "C" __global__ void poisson_residual_2d(
    float* __restrict__ residual,
    const float* __restrict__ phi,
    const float* __restrict__ source,
    float dx2_inv,
    float dy2_inv,
    int nx,
    int ny
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;

    if (i >= nx || j >= ny) return;

    int idx = j * nx + i;

    // Boundary points have zero residual
    if (i == 0 || i == nx - 1 || j == 0 || j == ny - 1) {
        residual[idx] = 0.0f;
        return;
    }

    int idx_ip1 = j * nx + (i + 1);
    int idx_im1 = j * nx + (i - 1);
    int idx_jp1 = (j + 1) * nx + i;
    int idx_jm1 = (j - 1) * nx + i;

    // Compute residual: r = f - L*φ where L is the Laplacian operator
    float laplacian = dx2_inv * (phi[idx_ip1] - 2.0f * phi[idx] + phi[idx_im1]) +
                     dy2_inv * (phi[idx_jp1] - 2.0f * phi[idx] + phi[idx_jm1]);

    residual[idx] = source[idx] - laplacian;
}

// 3D Poisson solver kernels
extern "C" __global__ void poisson_jacobi_3d(
    float* __restrict__ phi_new,
    const float* __restrict__ phi,
    const float* __restrict__ source,
    float factor,           // 1 / (2 * (1/dx² + 1/dy² + 1/dz²))
    float dx2_inv,
    float dy2_inv,
    float dz2_inv,
    int nx,
    int ny,
    int nz
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;
    int k = blockIdx.z * blockDim.z + threadIdx.z;

    if (i >= nx || j >= ny || k >= nz) return;

    int idx = k * nx * ny + j * nx + i;

    // Boundary conditions
    if (i == 0 || i == nx - 1 || j == 0 || j == ny - 1 || k == 0 || k == nz - 1) {
        phi_new[idx] = 0.0f;
        return;
    }

    int idx_ip1 = k * nx * ny + j * nx + (i + 1);
    int idx_im1 = k * nx * ny + j * nx + (i - 1);
    int idx_jp1 = k * nx * ny + (j + 1) * nx + i;
    int idx_jm1 = k * nx * ny + (j - 1) * nx + i;
    int idx_kp1 = (k + 1) * nx * ny + j * nx + i;
    int idx_km1 = (k - 1) * nx * ny + j * nx + i;

    // 3D Jacobi iteration
    phi_new[idx] = factor * (
        -source[idx] +
        dx2_inv * (phi[idx_ip1] + phi[idx_im1]) +
        dy2_inv * (phi[idx_jp1] + phi[idx_jm1]) +
        dz2_inv * (phi[idx_kp1] + phi[idx_km1])
    );
}

// Multigrid helper kernels
extern "C" __global__ void restrict_2d(
    float* __restrict__ coarse,
    const float* __restrict__ fine,
    int nx_fine,
    int ny_fine,
    int nx_coarse,
    int ny_coarse
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;

    if (i >= nx_coarse || j >= ny_coarse) return;

    int idx_coarse = j * nx_coarse + i;

    // Full weighting restriction (9-point stencil in 2D)
    int i_fine = 2 * i;
    int j_fine = 2 * j;

    if (i_fine >= nx_fine - 1 || j_fine >= ny_fine - 1) {
        coarse[idx_coarse] = 0.0f;
        return;
    }

    float sum = 0.0f;

    // Center point (weight 4)
    sum += 4.0f * fine[j_fine * nx_fine + i_fine];

    // Edge neighbors (weight 2 each)
    if (i_fine > 0) sum += 2.0f * fine[j_fine * nx_fine + (i_fine - 1)];
    if (i_fine < nx_fine - 1) sum += 2.0f * fine[j_fine * nx_fine + (i_fine + 1)];
    if (j_fine > 0) sum += 2.0f * fine[(j_fine - 1) * nx_fine + i_fine];
    if (j_fine < ny_fine - 1) sum += 2.0f * fine[(j_fine + 1) * nx_fine + i_fine];

    // Corner neighbors (weight 1 each)
    if (i_fine > 0 && j_fine > 0)
        sum += fine[(j_fine - 1) * nx_fine + (i_fine - 1)];
    if (i_fine < nx_fine - 1 && j_fine > 0)
        sum += fine[(j_fine - 1) * nx_fine + (i_fine + 1)];
    if (i_fine > 0 && j_fine < ny_fine - 1)
        sum += fine[(j_fine + 1) * nx_fine + (i_fine - 1)];
    if (i_fine < nx_fine - 1 && j_fine < ny_fine - 1)
        sum += fine[(j_fine + 1) * nx_fine + (i_fine + 1)];

    coarse[idx_coarse] = sum / 16.0f;  // Normalize by total weight
}

extern "C" __global__ void prolongate_2d(
    float* __restrict__ fine,
    const float* __restrict__ coarse,
    int nx_fine,
    int ny_fine,
    int nx_coarse,
    int ny_coarse
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;

    if (i >= nx_fine || j >= ny_fine) return;

    int idx_fine = j * nx_fine + i;

    // Bilinear interpolation from coarse grid
    float i_coarse = i * 0.5f;
    float j_coarse = j * 0.5f;

    int i0 = (int)i_coarse;
    int j0 = (int)j_coarse;
    int i1 = min(i0 + 1, nx_coarse - 1);
    int j1 = min(j0 + 1, ny_coarse - 1);

    float alpha = i_coarse - i0;
    float beta = j_coarse - j0;

    float val00 = coarse[j0 * nx_coarse + i0];
    float val10 = coarse[j0 * nx_coarse + i1];
    float val01 = coarse[j1 * nx_coarse + i0];
    float val11 = coarse[j1 * nx_coarse + i1];

    float val0 = (1.0f - alpha) * val00 + alpha * val10;
    float val1 = (1.0f - alpha) * val01 + alpha * val11;

    fine[idx_fine] += (1.0f - beta) * val0 + beta * val1;
}

// Specialized kernel for pressure Poisson equation with Neumann boundaries
extern "C" __global__ void pressure_poisson_2d(
    float* __restrict__ pressure_new,
    const float* __restrict__ pressure,
    const float* __restrict__ divergence,
    float dx2_inv,
    float dy2_inv,
    float dt,
    float rho,
    int nx,
    int ny
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;

    if (i >= nx || j >= ny) return;

    int idx = j * nx + i;
    float factor = 1.0f / (2.0f * (dx2_inv + dy2_inv));

    // Neumann boundary conditions (∂p/∂n = 0)
    if (i == 0 || i == nx - 1 || j == 0 || j == ny - 1) {
        // For Neumann BCs, use one-sided differences
        int idx_neighbor;
        if (i == 0) idx_neighbor = j * nx + 1;
        else if (i == nx - 1) idx_neighbor = j * nx + (nx - 2);
        else if (j == 0) idx_neighbor = 1 * nx + i;
        else idx_neighbor = (ny - 2) * nx + i;

        pressure_new[idx] = pressure[idx_neighbor];
        return;
    }

    int idx_ip1 = j * nx + (i + 1);
    int idx_im1 = j * nx + (i - 1);
    int idx_jp1 = (j + 1) * nx + i;
    int idx_jm1 = (j - 1) * nx + i;

    // Source term from velocity divergence (for pressure correction)
    float source = -rho * divergence[idx] / dt;

    pressure_new[idx] = factor * (
        source +
        dx2_inv * (pressure[idx_ip1] + pressure[idx_im1]) +
        dy2_inv * (pressure[idx_jp1] + pressure[idx_jm1])
    );
}