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

#ifndef FLT_MAX
#define FLT_MAX 3.402823466e+38f
#endif
/**
 * CUDA kernels for reduction operations
 *
 * Implements efficient parallel reduction operations for CFD:
 * - Max reduction (for residual checking)
 * - Sum reduction (for norms and dot products)
 * - Min reduction (for stability checks)
 *
 * Using warp shuffle operations for modern GPUs
 */

// (no host headers: NVRTC compiles these as device code with the built-ins only)

// Warp size constant
#define WARP_SIZE 32

/**
 * Warp-level reduction using shuffle operations
 */
__device__ float warp_reduce_sum(float val) {
    for (int offset = WARP_SIZE/2; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__device__ float warp_reduce_max(float val) {
    for (int offset = WARP_SIZE/2; offset > 0; offset /= 2) {
        val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
    }
    return val;
}

__device__ float warp_reduce_min(float val) {
    for (int offset = WARP_SIZE/2; offset > 0; offset /= 2) {
        val = fminf(val, __shfl_down_sync(0xffffffff, val, offset));
    }
    return val;
}

/**
 * Block-level reduction using shared memory and warp reduction
 */
__device__ float block_reduce_sum(float val) {
    __shared__ float shared[32]; // One value per warp
    int lane = threadIdx.x % WARP_SIZE;
    int wid = threadIdx.x / WARP_SIZE;

    // First reduce within warps
    val = warp_reduce_sum(val);

    // Write warp results to shared memory
    if (lane == 0) shared[wid] = val;
    __syncthreads();

    // Final reduction of warp results
    val = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared[lane] : 0;
    if (wid == 0) val = warp_reduce_sum(val);

    return val;
}

__device__ float block_reduce_max(float val) {
    __shared__ float shared[32];
    int lane = threadIdx.x % WARP_SIZE;
    int wid = threadIdx.x / WARP_SIZE;

    val = warp_reduce_max(val);

    if (lane == 0) shared[wid] = val;
    __syncthreads();

    val = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared[lane] : -FLT_MAX;
    if (wid == 0) val = warp_reduce_max(val);

    return val;
}

/**
 * Sum reduction kernel - reduces array to single value
 * Each block produces a partial sum, requires second pass for final result
 */
extern "C" __global__ void reduce_sum(
    const float* __restrict__ input,
    float* __restrict__ output,
    int n
) {
    float sum = 0;

    // Grid-stride loop for coalesced memory access
    for (int i = blockIdx.x * blockDim.x + threadIdx.x;
         i < n;
         i += blockDim.x * gridDim.x) {
        sum += input[i];
    }

    // Reduce within block
    sum = block_reduce_sum(sum);

    // Write block result
    if (threadIdx.x == 0) {
        output[blockIdx.x] = sum;
    }
}

/**
 * Max reduction kernel - finds maximum value in array
 */
extern "C" __global__ void reduce_max(
    const float* __restrict__ input,
    float* __restrict__ output,
    int n
) {
    float max_val = -FLT_MAX;

    // Grid-stride loop
    for (int i = blockIdx.x * blockDim.x + threadIdx.x;
         i < n;
         i += blockDim.x * gridDim.x) {
        max_val = fmaxf(max_val, input[i]);
    }

    // Reduce within block
    max_val = block_reduce_max(max_val);

    // Write block result
    if (threadIdx.x == 0) {
        output[blockIdx.x] = max_val;
    }
}

/**
 * Absolute max reduction kernel - finds maximum absolute value
 * Used for residual checking in iterative solvers
 */
extern "C" __global__ void reduce_abs_max(
    const float* __restrict__ input,
    float* __restrict__ output,
    int n
) {
    float max_val = 0;

    // Grid-stride loop with absolute value
    for (int i = blockIdx.x * blockDim.x + threadIdx.x;
         i < n;
         i += blockDim.x * gridDim.x) {
        max_val = fmaxf(max_val, fabsf(input[i]));
    }

    // Reduce within block
    max_val = block_reduce_max(max_val);

    // Write block result
    if (threadIdx.x == 0) {
        output[blockIdx.x] = max_val;
    }
}

/**
 * L2 norm squared reduction - computes sum of squares
 */
extern "C" __global__ void reduce_norm2_squared(
    const float* __restrict__ input,
    float* __restrict__ output,
    int n
) {
    float sum = 0;

    // Grid-stride loop with squaring
    for (int i = blockIdx.x * blockDim.x + threadIdx.x;
         i < n;
         i += blockDim.x * gridDim.x) {
        float val = input[i];
        sum += val * val;
    }

    // Reduce within block
    sum = block_reduce_sum(sum);

    // Write block result
    if (threadIdx.x == 0) {
        output[blockIdx.x] = sum;
    }
}

/**
 * Final reduction kernel - reduces partial results from blocks
 * Call with single block after first reduction pass
 */
extern "C" __global__ void reduce_final_sum(
    const float* __restrict__ input,
    float* __restrict__ output,
    int n
) {
    float sum = 0;

    // Single block reduces all partial results
    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        sum += input[i];
    }

    sum = block_reduce_sum(sum);

    if (threadIdx.x == 0) {
        output[0] = sum;
    }
}

extern "C" __global__ void reduce_final_max(
    const float* __restrict__ input,
    float* __restrict__ output,
    int n
) {
    float max_val = -FLT_MAX;

    for (int i = threadIdx.x; i < n; i += blockDim.x) {
        max_val = fmaxf(max_val, input[i]);
    }

    max_val = block_reduce_max(max_val);

    if (threadIdx.x == 0) {
        output[0] = max_val;
    }
}

/**
 * Compute residual and find max for Poisson solver
 * Combines residual computation with reduction for efficiency
 */
extern "C" __global__ void poisson_residual_max(
    const float* __restrict__ phi,
    const float* __restrict__ source,
    float* __restrict__ residual_max,
    float dx2_inv,
    float dy2_inv,
    int nx,
    int ny
) {
    float max_residual = 0;

    // Grid-stride loop over interior points
    for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
         idx < nx * ny;
         idx += blockDim.x * gridDim.x) {

        int i = idx % nx;
        int j = idx / nx;

        // Skip boundary points
        if (i > 0 && i < nx-1 && j > 0 && j < ny-1) {
            // Compute Laplacian
            float laplacian =
                dx2_inv * (phi[idx-1] - 2.0f*phi[idx] + phi[idx+1]) +
                dy2_inv * (phi[idx-nx] - 2.0f*phi[idx] + phi[idx+nx]);

            // Residual = source - laplacian
            float res = fabsf(source[idx] - laplacian);
            max_residual = fmaxf(max_residual, res);
        }
    }

    // Reduce within block
    max_residual = block_reduce_max(max_residual);

    // Write block result
    if (threadIdx.x == 0) {
        residual_max[blockIdx.x] = max_residual;
    }
}