kernels.cuh

#pragma once
// Shared CUDA kernels used by v2..v6. Each translation unit that includes
// this header compiles its own copy (static __global__ has internal linkage),
// so there are no link-time conflicts and unused kernels are dropped by
// ptxas in TUs that do not invoke them.

#include <cstddef>
#include <cstdint>
#include <cufft.h>
#include <cuda_runtime.h>

namespace sbfft {
namespace kern {

constexpr int kPerThread = 4;   // outputs per thread in sparse_row_dft_iter
constexpr int kTx        = 32;  // transpose tile dim
constexpr int kTy        = 8;   // transpose rows-per-block
constexpr int kStatsBlock = 256;

// Twiddle table: W[k] = exp(-2*PI*i * k / N). One thread per k.
static __global__ void build_twiddle_table(int N,
                                           cufftComplex* __restrict__ W) {
    int k = blockIdx.x * blockDim.x + threadIdx.x;
    if (k >= N) return;
    float ang = -6.2831853071795864769f * (float)k / (float)N;
    float c, s;
    __sincosf(ang, &s, &c);
    W[k].x = c;
    W[k].y = s;
}

// Sparse-row direct DFT with iterative twiddle update. Templated on the
// column-index type (int for v3, uint16_t for v4/v5/v6). Each thread owns
// kPerThread consecutive output frequencies; chains kPerThread complex
// multiplies of `step = W[x]` after one initial `W[(v_init*x) mod N]` load.
template <typename ColIdx>
static __global__ void sparse_row_dft_iter(const int*    __restrict__ rowPtr,
                                           const ColIdx* __restrict__ colIdx,
                                           const cufftComplex* __restrict__ W,
                                           int N, int v0, int T,
                                           cufftComplex* __restrict__ B_tile) {
    int row = blockIdx.x;
    int s = rowPtr[row];
    int e = rowPtr[row + 1];
    int k = e - s;
    cufftComplex* outRow = B_tile + (size_t)row * T;

    int my_v_start = threadIdx.x * kPerThread;

    if (k == 0) {
        #pragma unroll
        for (int j = 0; j < kPerThread; ++j) {
            int idx = my_v_start + j;
            if (idx < T) { outRow[idx].x = 0.f; outRow[idx].y = 0.f; }
        }
        return;
    }

    float acc_re[kPerThread];
    float acc_im[kPerThread];
    #pragma unroll
    for (int j = 0; j < kPerThread; ++j) { acc_re[j] = 0.f; acc_im[j] = 0.f; }

    const int v_init = v0 + my_v_start;

    for (int kk = s; kk < e; ++kk) {
        int x = colIdx[kk];
        int idx0 = (int)(((long long)v_init * x) % N);
        cufftComplex twid = W[idx0];
        cufftComplex step = W[x];

        #pragma unroll
        for (int j = 0; j < kPerThread; ++j) {
            acc_re[j] += twid.x;
            acc_im[j] += twid.y;
            float nx = twid.x * step.x - twid.y * step.y;
            float ny = twid.x * step.y + twid.y * step.x;
            twid.x = nx;
            twid.y = ny;
        }
    }

    #pragma unroll
    for (int j = 0; j < kPerThread; ++j) {
        int idx = my_v_start + j;
        if (idx < T) {
            outRow[idx].x = acc_re[j];
            outRow[idx].y = acc_im[j];
        }
    }
}

// 32x32 padded shared-memory tile transpose. src is [rows][cols] row-major;
// dst is [cols][rows] row-major. Pad to TX+1 to defeat bank conflicts.
static __global__ void transpose_kernel(const cufftComplex* __restrict__ src,
                                        cufftComplex* __restrict__ dst,
                                        int rows, int cols) {
    __shared__ cufftComplex tile[kTx][kTx + 1];
    int x = blockIdx.x * kTx + threadIdx.x;
    int y = blockIdx.y * kTx + threadIdx.y;
    #pragma unroll
    for (int dy = 0; dy < kTx; dy += kTy) {
        int yy = y + dy;
        if (x < cols && yy < rows)
            tile[threadIdx.y + dy][threadIdx.x] = src[(size_t)yy * cols + x];
    }
    __syncthreads();
    int x2 = blockIdx.y * kTx + threadIdx.x;
    int y2 = blockIdx.x * kTx + threadIdx.y;
    #pragma unroll
    for (int dy = 0; dy < kTx; dy += kTy) {
        int yy2 = y2 + dy;
        if (x2 < rows && yy2 < cols)
            dst[(size_t)yy2 * rows + x2] = tile[threadIdx.x][threadIdx.y + dy];
    }
}

static inline void launch_transpose(const cufftComplex* src,
                                    cufftComplex* dst,
                                    int rows, int cols,
                                    cudaStream_t stream = 0) {
    dim3 block(kTx, kTy);
    dim3 grid((cols + kTx - 1) / kTx, (rows + kTx - 1) / kTx);
    transpose_kernel<<<grid, block, 0, stream>>>(src, dst, rows, cols);
}

// Per-tile partial reduction for SIGMA|Y|^2 and max|Y| over the valid prefix
// [0, valid_T) of each row of an N-by-T tile. Writes one (sum, max) per
// block; caller reduces on host.
static __global__ void tile_stats(const cufftComplex* __restrict__ tile,
                                  int N, int T, int valid_T,
                                  double* __restrict__ block_sum,
                                  float*  __restrict__ block_max) {
    extern __shared__ unsigned char smem_raw[];
    double* s_sum = reinterpret_cast<double*>(smem_raw);
    float*  s_max = reinterpret_cast<float*>(s_sum + blockDim.x);

    double mySum = 0.0;
    float  myMax = 0.f;

    size_t total = (size_t)N * valid_T;
    size_t i = blockIdx.x * blockDim.x + threadIdx.x;
    size_t stride = (size_t)gridDim.x * blockDim.x;
    while (i < total) {
        size_t row = i / valid_T;
        size_t j   = i - row * valid_T;
        cufftComplex v = tile[row * T + j];
        double m2 = (double)v.x * v.x + (double)v.y * v.y;
        mySum += m2;
        float mag = sqrtf((float)m2);
        if (mag > myMax) myMax = mag;
        i += stride;
    }
    s_sum[threadIdx.x] = mySum;
    s_max[threadIdx.x] = myMax;
    __syncthreads();
    for (int off = blockDim.x / 2; off > 0; off >>= 1) {
        if (threadIdx.x < off) {
            s_sum[threadIdx.x] += s_sum[threadIdx.x + off];
            float other = s_max[threadIdx.x + off];
            if (other > s_max[threadIdx.x]) s_max[threadIdx.x] = other;
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        block_sum[blockIdx.x] = s_sum[0];
        block_max[blockIdx.x] = s_max[0];
    }
}

} // namespace kern
} // namespace sbfft