v4_bitpacked.cu

// V4 = V3 (iterative twiddle, column-block streaming) +
//      (a) transpose-to-unit-stride layout for cuFFT column FFT
//      (b) uint16 column indices in CSR (N < 65536 holds for all targets).

#include "versions.h"
#include "common/auto_tile.h"
#include "common/cuda_check.h"
#include "common/kernels.cuh"
#include "common/timer.h"

#include <algorithm>
#include <cmath>
#include <cstdio>

namespace sbfft {

namespace {

constexpr int PER_THREAD = kern::kPerThread;
using kern::build_twiddle_table;
using kern::launch_transpose;
using kern::tile_stats;

} // anonymous

VersionResult run_v4_bitpacked(const CsrBinary& csr, const VersionConfig& cfg) {
    VersionResult R;
    R.name = "v4_layout_u16";

    const int N = csr.N;
    if (N >= 65536) {
        std::fprintf(stderr,
                     "v4: N=%d >= 65536, uint16 colIdx not applicable; "
                     "falling back to v3 path is not implemented.\n", N);
        std::exit(1);
    }
    int T = cfg.tile_T > 0 ? cfg.tile_T
                           : pick_tile_T(N, "v4", 1, cfg.verbose);
    const int round = PER_THREAD * 32;
    if (T % round != 0) T = ((T + round - 1) / round) * round;
    const size_t N2 = static_cast<size_t>(N) * N;

    CpuTimer wall;
    wall.tic();

    // Pack colIdx to uint16 on host (cheap; one-time CPU pass).
    std::vector<uint16_t> colIdx_u16(csr.nnz);
    for (int64_t i = 0; i < csr.nnz; ++i)
        colIdx_u16[i] = static_cast<uint16_t>(csr.colIdx[i]);

    int* d_rowPtr = nullptr;
    uint16_t* d_colIdx16 = nullptr;
    CUDA_CHECK(cudaMalloc(&d_rowPtr, (N + 1) * sizeof(int)));
    CUDA_CHECK(cudaMalloc(&d_colIdx16, csr.nnz * sizeof(uint16_t)));

    GpuTimer gt;
    gt.tic();
    CUDA_CHECK(cudaMemcpy(d_rowPtr, csr.rowPtr.data(),
                          (N + 1) * sizeof(int), cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_colIdx16, colIdx_u16.data(),
                          csr.nnz * sizeof(uint16_t), cudaMemcpyHostToDevice));
    R.t_h2d_input_ms = gt.toc_ms();

    // Twiddle table.
    cufftComplex* d_W = nullptr;
    CUDA_CHECK(cudaMalloc(&d_W, N * sizeof(cufftComplex)));

    GpuTimer gt_other;
    double acc_other = 0.0;
    gt_other.tic();
    {
        int block = 256;
        int grid = (N + block - 1) / block;
        build_twiddle_table<<<grid, block>>>(N, d_W);
        CUDA_CHECK(cudaGetLastError());
    }
    acc_other += gt_other.toc_ms();

    // Two slabs.
    cufftComplex* d_tile_a = nullptr;  // row-major [N][T]
    cufftComplex* d_tile_b = nullptr;  // col-major [T][N]
    CUDA_CHECK(cudaMalloc(&d_tile_a, (size_t)N * T * sizeof(cufftComplex)));
    CUDA_CHECK(cudaMalloc(&d_tile_b, (size_t)T * N * sizeof(cufftComplex)));

    cufftComplex* h_full = nullptr;
    double pin_alloc_ms = 0;
    if (cfg.keep_full_output) {
        CpuTimer pa; pa.tic();
        CUDA_CHECK(cudaMallocHost(&h_full, N2 * sizeof(cufftComplex)));
        pin_alloc_ms = pa.toc_ms();
    }

    constexpr int STATS_BLOCK = 256;
    int stats_grid = std::min(2048, (N * T + STATS_BLOCK - 1) / STATS_BLOCK);
    double* d_bsum = nullptr;
    float*  d_bmax = nullptr;
    CUDA_CHECK(cudaMalloc(&d_bsum, stats_grid * sizeof(double)));
    CUDA_CHECK(cudaMalloc(&d_bmax, stats_grid * sizeof(float)));
    std::vector<double> h_bsum(stats_grid);
    std::vector<float>  h_bmax(stats_grid);

    // Unit-stride col FFT plan over d_tile_b.
    //   length=N, batch=T, istride=1, idist=N
    cufftHandle plan;
    gt.tic();
    {
        int n[1] = { N };
        int em[1] = { N };
        CUFFT_CHECK(cufftPlanMany(&plan,
                                  1, n,
                                  em, /*istride=*/1, /*idist=*/N,
                                  em, /*ostride=*/1, /*odist=*/N,
                                  CUFFT_C2C, /*batch=*/T));
    }
    R.t_plan_ms = gt.toc_ms();

    size_t plan_ws = 0;
    cufftGetSize(plan, &plan_ws);

    int n_blocks = (N + T - 1) / T;
    GpuTimer gt_s1, gt_s2, gt_d2h, gt_tr;
    double acc_s1 = 0, acc_s2 = 0, acc_d2h = 0, acc_tr = 0;

    // Pre-plan probe bins: which tile contains each probe?
    // tile = v / T; col_off = v - tile*T; capture from d_tile_a[u*T + col_off].
    R.probe_values.assign(cfg.probe_bins.size(), {0.f, 0.f});
    std::vector<std::vector<int>> probes_per_tile(n_blocks);
    for (size_t i = 0; i < cfg.probe_bins.size(); ++i) {
        int v = cfg.probe_bins[i].second;
        int t = v / T;
        if (t >= 0 && t < n_blocks) probes_per_tile[t].push_back((int)i);
    }

    double sum_sq = 0.0;
    float  max_mag = 0.f;

    const int row_block = T / PER_THREAD;

    for (int b = 0; b < n_blocks; ++b) {
        int v0 = b * T;
        int valid = std::min(T, N - v0);

        // step 1: row DFT into d_tile_a (row-major [N][T])
        gt_s1.tic();
        kern::sparse_row_dft_iter<uint16_t><<<N, row_block>>>(
            d_rowPtr, d_colIdx16, d_W, N, v0, T, d_tile_a);
        CUDA_CHECK(cudaGetLastError());
        acc_s1 += gt_s1.toc_ms();

        // transpose: d_tile_a (N x T) into d_tile_b (T x N)
        gt_tr.tic();
        launch_transpose(d_tile_a, d_tile_b, N, T);
        acc_tr += gt_tr.toc_ms();

        // step 2: cuFFT col on d_tile_b (unit stride)
        gt_s2.tic();
        CUFFT_CHECK(cufftExecC2C(plan, d_tile_b, d_tile_b, CUFFT_FORWARD));
        acc_s2 += gt_s2.toc_ms();

        // transpose back: d_tile_b (T x N) into d_tile_a (N x T)
        gt_tr.tic();
        launch_transpose(d_tile_b, d_tile_a, T, N);
        acc_tr += gt_tr.toc_ms();

        // stats from d_tile_a
        gt_other.tic();
        {
            size_t smem = STATS_BLOCK * (sizeof(double) + sizeof(float));
            tile_stats<<<stats_grid, STATS_BLOCK, smem>>>(
                d_tile_a, N, T, valid, d_bsum, d_bmax);
            CUDA_CHECK(cudaGetLastError());
        }
        acc_other += gt_other.toc_ms();

        CUDA_CHECK(cudaMemcpy(h_bsum.data(), d_bsum,
                              stats_grid * sizeof(double),
                              cudaMemcpyDeviceToHost));
        CUDA_CHECK(cudaMemcpy(h_bmax.data(), d_bmax,
                              stats_grid * sizeof(float),
                              cudaMemcpyDeviceToHost));
        for (int i = 0; i < stats_grid; ++i) {
            sum_sq += h_bsum[i];
            if (h_bmax[i] > max_mag) max_mag = h_bmax[i];
        }

        if (h_full) {
            gt_d2h.tic();
            CUDA_CHECK(cudaMemcpy2D(h_full + v0,
                                    N * sizeof(cufftComplex),
                                    d_tile_a,
                                    T * sizeof(cufftComplex),
                                    valid * sizeof(cufftComplex),
                                    N,
                                    cudaMemcpyDeviceToHost));
            acc_d2h += gt_d2h.toc_ms();
        }

        // Probe-bin capture for this tile.
        for (int pi : probes_per_tile[b]) {
            int u = cfg.probe_bins[pi].first;
            int v = cfg.probe_bins[pi].second;
            int col_off = v - v0;
            CUDA_CHECK(cudaMemcpy(&R.probe_values[pi],
                                  d_tile_a + (size_t)u * T + col_off,
                                  sizeof(cufftComplex),
                                  cudaMemcpyDeviceToHost));
        }
    }

    R.t_step1_ms = acc_s1;
    R.t_step2_ms = acc_s2;
    R.t_other_ms = acc_other + acc_tr;
    R.t_d2h_ms   = acc_d2h;

    R.stats.sum_sq_mag = sum_sq;
    R.stats.max_mag = max_mag;
    R.stats.mean_mag = std::sqrt(sum_sq / static_cast<double>(N2));
    R.parseval_rel_err = parseval_relative_error(R.stats, N, csr.nnz);

    double finalize_ms = 0;
    if (h_full) {
        CpuTimer ft; ft.tic();
        R.full_output.assign(h_full, h_full + N2);
        CUDA_CHECK(cudaFreeHost(h_full));
        finalize_ms = ft.toc_ms();
    }

    R.peak_device_bytes = (N + 1) * sizeof(int)
                        + csr.nnz * sizeof(uint16_t)
                        + 2 * (size_t)N * T * sizeof(cufftComplex)
                        + N * sizeof(cufftComplex)
                        + plan_ws
                        + stats_grid * (sizeof(double) + sizeof(float));

    cufftDestroy(plan);
    CUDA_CHECK(cudaFree(d_bmax));
    CUDA_CHECK(cudaFree(d_bsum));
    CUDA_CHECK(cudaFree(d_tile_b));
    CUDA_CHECK(cudaFree(d_tile_a));
    CUDA_CHECK(cudaFree(d_W));
    CUDA_CHECK(cudaFree(d_colIdx16));
    CUDA_CHECK(cudaFree(d_rowPtr));

    R.t_total_ms = wall.toc_ms();

    if (cfg.verbose) {
        std::printf("  v4 transpose : %8.3f ms (forward + back)\n", acc_tr);
        if (cfg.keep_full_output) {
            std::printf("  v4 pinned alc: %8.3f ms (cudaMallocHost N*N)\n", pin_alloc_ms);
            std::printf("  v4 finalize  : %8.3f ms (vec::assign + cudaFreeHost)\n", finalize_ms);
        }
    }
    return R;
}

} // namespace sbfft