scan.cu

/**
 * Zero-Loss Parallel Prefix Sum (Scan) Optimization for Turing (sm_75)
 * 
 * DESIGN RATIONALE:
 * 1. Two-Phase Recursive Architecture:
 *    - Processes 2048 elements per block (256 threads * 8 elements/thread).
 *    - Phase 1: Local scan within each 2048-element tile.
 *               Coalesced load -> Shared memory with padding -> Local register scan ->
 *               Warp shuffle reduction -> Uniform offset correction -> Coalesced store.
 *    - Phase 2: Recursively scans block sums. If the block sums array size <= 2048,
 *               it scans in a single block, avoiding launch overhead.
 *    - Phase 3: Uniformly adds block sums back to scanned elements.
 * 
 * 2. Shared Memory Bank Conflict Mitigation:
 *    - Layout: flat array of 2048 elements + 64 padding elements = 2112 floats.
 *    - Padding Formula: physical_index = logical_index + (logical_index >> 5).
 *    - Mathematically eliminates 100% of 32-bank shared memory conflicts on both 
 *      contiguous chunk reads and strided block-write operations.
 * 
 * 3. Register Budget Compliance:
 *    - Enforced with __launch_bounds__(256, 4) to ensure 100% occupancy (1024 threads/SM).
 *    - Caps registers per thread to 64, eliminating local memory spills (0.3 ns vs 400 ns).
 */

#include <iostream>
#include <vector>
#include <cmath>
#include <chrono>
#include <cuda_runtime.h>

#define TILE_SIZE 2048
#define BLOCK_SIZE 256

// Host-side CUDA error-checking macro
#define CUDA_CHECK(call) \
    do { \
        cudaError_t err = call; \
        if (err != cudaSuccess) { \
            std::cerr << "CUDA Error: " << cudaGetErrorString(err) \
                      << " at " << __FILE__ << ":" << __LINE__ << std::endl; \
            exit(EXIT_FAILURE); \
        } \
    } while (0)

#define CUDA_POST_KERNEL_CHECK() \
    do { \
        cudaError_t err = cudaGetLastError(); \
        if (err != cudaSuccess) { \
            std::cerr << "CUDA Kernel Error: " << cudaGetErrorString(err) \
                      << " at " << __FILE__ << ":" << __LINE__ << std::endl; \
            exit(EXIT_FAILURE); \
        } \
    } while (0)

// Shared memory physical index mapping to eliminate bank conflicts
__device__ __forceinline__ size_t get_phys_idx(size_t logical_idx) {
    return logical_idx + (logical_idx >> 5);
}

/**
 * local_scan_kernel
 * Performs block-level inclusive or exclusive scan.
 * Note: __restrict__ is removed to prevent aliasing violations during in-place recursive calls.
 */
template <bool INCLUSIVE>
__global__ void __launch_bounds__(BLOCK_SIZE, 4)
local_scan_kernel(const float* input,
                  float* output,
                  float* d_block_sums,
                  size_t N) {
    // 2048 elements + 64 pads = 2112 floats (~8.25 KB per block)
    __shared__ float s_data[2112];
    __shared__ float s_warp_sums[8];

    size_t block_offset = (size_t)blockIdx.x * TILE_SIZE;
    size_t thread_id = threadIdx.x;

    // 1. Coalesced load from global memory to padded shared memory (64-bit safe indices)
    #pragma unroll
    for (size_t i = 0; i < 8; ++i) {
        size_t logical_idx = thread_id + i * BLOCK_SIZE;
        size_t idx = block_offset + logical_idx;
        float val = (idx < N) ? input[idx] : 0.0f;
        s_data[get_phys_idx(logical_idx)] = val;
    }
    __syncthreads();

    // 2. Load contiguous 8-element chunk to registers
    float local_vals[8];
    size_t chunk_start = thread_id * 8;
    #pragma unroll
    for (size_t i = 0; i < 8; ++i) {
        local_vals[i] = s_data[get_phys_idx(chunk_start + i)];
    }

    // Parallel Prefix Scan Tree in registers (eliminates the 8-addition serial dependency chain)
    float x0 = local_vals[0];
    float x1 = local_vals[1];
    float x2 = local_vals[2];
    float x3 = local_vals[3];
    float x4 = local_vals[4];
    float x5 = local_vals[5];
    float x6 = local_vals[6];
    float x7 = local_vals[7];

    float s01 = x0 + x1;
    float s23 = x2 + x3;
    float s45 = x4 + x5;
    float s67 = x6 + x7;

    float s03 = s01 + s23;
    float s47 = s45 + s67;

    float sum = s03 + s47; // Total sum of the chunk

    local_vals[0] = x0;
    local_vals[1] = s01;
    local_vals[2] = s01 + x2;
    local_vals[3] = s03;
    local_vals[4] = s03 + x4;
    local_vals[5] = s03 + s45;
    local_vals[6] = s03 + s45 + x6;
    local_vals[7] = sum;

    // 3. Cooperative warp-level scan of local sums
    int lane = thread_id & 31;
    int warp_id = thread_id >> 5;
    float warp_sum = sum;

    #pragma unroll
    for (int offset = 1; offset < 32; offset *= 2) {
        float tmp = __shfl_up_sync(0xffffffff, warp_sum, offset);
        if (lane >= offset) {
            warp_sum += tmp;
        }
    }

    // Write warp sums to shared memory
    if (lane == 31) {
        s_warp_sums[warp_id] = warp_sum;
    }
    __syncthreads();

    // Warp 0 scans the warp sums
    if (warp_id == 0) {
        float block_warp_sum = (thread_id < 8) ? s_warp_sums[thread_id] : 0.0f;
        #pragma unroll
        for (int offset = 1; offset < 8; offset *= 2) {
            float tmp = __shfl_up_sync(0xffffffff, block_warp_sum, offset);
            if (thread_id < 8 && lane >= offset) {
                block_warp_sum += tmp;
            }
        }
        if (thread_id < 8) {
            s_warp_sums[thread_id] = block_warp_sum;
        }
    }
    __syncthreads();

    // Compute the cumulative prefix sum offset for this thread
    float prev_warp_sum = (warp_id > 0) ? s_warp_sums[warp_id - 1] : 0.0f;
    float prev_block_sum = prev_warp_sum + (warp_sum - sum);

    // 4. Store scanned values back to padded shared memory
    if (INCLUSIVE) {
        #pragma unroll
        for (size_t i = 0; i < 8; ++i) {
            s_data[get_phys_idx(chunk_start + i)] = local_vals[i] + prev_block_sum;
        }
    } else {
        float prev = prev_block_sum;
        #pragma unroll
        for (size_t i = 0; i < 8; ++i) {
            float curr = local_vals[i] + prev_block_sum;
            s_data[get_phys_idx(chunk_start + i)] = prev;
            prev = curr;
        }
    }
    __syncthreads();

    // Store raw inclusive block sum for hierarchical scan propagation
    if (thread_id == 255 && d_block_sums != nullptr) {
        d_block_sums[blockIdx.x] = s_warp_sums[7];
    }

    // 5. Coalesced write-back to global memory
    #pragma unroll
    for (size_t i = 0; i < 8; ++i) {
        size_t logical_idx = thread_id + i * BLOCK_SIZE;
        size_t idx = block_offset + logical_idx;
        if (idx < N) {
            output[idx] = s_data[get_phys_idx(logical_idx)];
        }
    }
}

/**
 * uniform_add_kernel
 * Uniformly adds the scanned block sums back to local scanned blocks.
 */
__global__ void __launch_bounds__(BLOCK_SIZE, 4)
uniform_add_kernel(float* __restrict__ output,
                   const float* __restrict__ d_block_sums,
                   size_t N) {
    size_t block_id = (size_t)blockIdx.x + 1; // Block 0 needs no offset
    size_t block_offset = block_id * TILE_SIZE;
    size_t thread_id = threadIdx.x;

    float add_val = d_block_sums[block_id];

    #pragma unroll
    for (size_t i = 0; i < 8; ++i) {
        size_t logical_idx = thread_id + i * BLOCK_SIZE;
        size_t idx = block_offset + logical_idx;
        if (idx < N) {
            output[idx] += add_val;
        }
    }
}

/**
 * Hierarchical Recursive GPU Scan Driver
 */
template <bool INCLUSIVE>
void gpu_scan(const float* d_in, float* d_out, size_t N, float** temp_arrays, float* d_block_sums, int temp_level = 0) {
    if (N <= TILE_SIZE) {
        local_scan_kernel<INCLUSIVE><<<1, BLOCK_SIZE>>>(d_in, d_out, nullptr, N);
        CUDA_POST_KERNEL_CHECK();
        return;
    }

    size_t num_blocks = (N + TILE_SIZE - 1) / TILE_SIZE;
    local_scan_kernel<INCLUSIVE><<<num_blocks, BLOCK_SIZE>>>(d_in, d_out, d_block_sums, N);
    CUDA_POST_KERNEL_CHECK();

    if (num_blocks > TILE_SIZE) {
        size_t next_num_blocks = (num_blocks + TILE_SIZE - 1) / TILE_SIZE;
        float* next_block_sums = temp_arrays[temp_level];
        // Hierarchical block-sum scans must execute exclusively
        gpu_scan<false>(d_block_sums, d_block_sums, num_blocks, temp_arrays, next_block_sums, temp_level + 1);
    } else {
        local_scan_kernel<false><<<1, BLOCK_SIZE>>>(d_block_sums, d_block_sums, nullptr, num_blocks);
        CUDA_POST_KERNEL_CHECK();
    }

    uniform_add_kernel<<<num_blocks - 1, BLOCK_SIZE>>>(d_out, d_block_sums, N);
    CUDA_POST_KERNEL_CHECK();
}

/**
 * Static Memory Cache for workspace pointers
 */
struct WorkspaceCache {
    float* block_sums = nullptr;
    size_t block_sums_size = 0;
    std::vector<float*> temp_arrays;
    std::vector<size_t> temp_sizes;

    ~WorkspaceCache() {
        if (block_sums) {
            cudaFree(block_sums);
        }
        for (float* ptr : temp_arrays) {
            if (ptr) cudaFree(ptr);
        }
    }
};

/**
 * Host Interface
 * Reuses device allocations across identical calling dimensions via a local workspace cache.
 */
extern "C" void solution(float* input, float* output, int N) {
    static WorkspaceCache cache;
    size_t num_blocks = ((size_t)N + TILE_SIZE - 1) / TILE_SIZE;

    // Compute required size bounds for temporary workspace arrays
    std::vector<size_t> needed_temp_sizes;
    size_t temp_n = num_blocks;
    while (temp_n > TILE_SIZE) {
        size_t next_n = (temp_n + TILE_SIZE - 1) / TILE_SIZE;
        needed_temp_sizes.push_back(next_n);
        temp_n = next_n;
    }

    // Resolve allocations for the block sums array
    if (num_blocks > 1) {
        if (cache.block_sums == nullptr || num_blocks > cache.block_sums_size) {
            if (cache.block_sums) {
                CUDA_CHECK(cudaFree(cache.block_sums));
            }
            cache.block_sums_size = num_blocks;
            CUDA_CHECK(cudaMalloc(&cache.block_sums, cache.block_sums_size * sizeof(float)));
        }
    }

    // Resolve allocations for hierarchical temporary arrays
    if (needed_temp_sizes.size() > cache.temp_arrays.size()) {
        cache.temp_arrays.resize(needed_temp_sizes.size(), nullptr);
        cache.temp_sizes.resize(needed_temp_sizes.size(), 0);
    }

    for (size_t i = 0; i < needed_temp_sizes.size(); ++i) {
        if (cache.temp_arrays[i] == nullptr || needed_temp_sizes[i] > cache.temp_sizes[i]) {
            if (cache.temp_arrays[i]) {
                CUDA_CHECK(cudaFree(cache.temp_arrays[i]));
            }
            cache.temp_sizes[i] = needed_temp_sizes[i];
            CUDA_CHECK(cudaMalloc(&cache.temp_arrays[i], cache.temp_sizes[i] * sizeof(float)));
        }
    }

    gpu_scan<true>(input, output, N, cache.temp_arrays.data(), cache.block_sums);
    CUDA_CHECK(cudaDeviceSynchronize());
}

/**
 * CPU Reference Implementations for Correctness Validation
 */
void cpu_inclusive_scan(const float* input, float* output, size_t N) {
    if (N == 0) return;
    output[0] = input[0];
    for (size_t i = 1; i < N; ++i) {
        output[i] = output[i - 1] + input[i];
    }
}

void cpu_exclusive_scan(const float* input, float* output, size_t N) {
    if (N == 0) return;
    output[0] = 0.0f;
    for (size_t i = 1; i < N; ++i) {
        output[i] = output[i - 1] + input[i - 1];
    }
}

/**
 * Verification Utility
 */
bool verify_results(const float* host, const float* device, size_t N, float tolerance = 1e-3f) {
    for (size_t i = 0; i < N; ++i) {
        float diff = std::abs(host[i] - device[i]);
        float ref = std::abs(host[i]);
        float rel_err = (ref > 1.0f) ? (diff / ref) : diff;
        if (rel_err > tolerance) {
            std::cout << "Mismatch at index " << i << " | CPU: " << host[i] << " | GPU: " << device[i] << " | Rel Error: " << rel_err << std::endl;
            return false;
        }
    }
    return true;
}

/**
 * Benchmark Harness
 */
void run_benchmark(size_t N) {
    std::cout << "\n==================================================" << std::endl;
    std::cout << "Benchmarking N = " << N << " (" << (static_cast<double>(N) / 1e6) << "M elements)" << std::endl;
    std::cout << "==================================================" << std::endl;

    size_t bytes = N * sizeof(float);
    std::vector<float> h_in(N);
    std::vector<float> h_out_cpu(N);
    std::vector<float> h_out_gpu(N);

    // Initialize input with deterministic values in [-1.0, 1.0] range
    for (size_t i = 0; i < N; ++i) {
        h_in[i] = static_cast<float>(rand()) / static_cast<float>(RAND_MAX) * 2.0f - 1.0f;
    }

    float *d_in, *d_out;
    CUDA_CHECK(cudaMalloc(&d_in, bytes));
    CUDA_CHECK(cudaMalloc(&d_out, bytes));

    CUDA_CHECK(cudaMemcpy(d_in, h_in.data(), bytes, cudaMemcpyHostToDevice));

    // Prepare workspace memory once outside benchmark loops
    size_t num_blocks = (N + TILE_SIZE - 1) / TILE_SIZE;
    std::vector<float*> temp_arrays;
    size_t temp_n = num_blocks;

    while (temp_n > TILE_SIZE) {
        size_t next_n = (temp_n + TILE_SIZE - 1) / TILE_SIZE;
        float* d_temp;
        CUDA_CHECK(cudaMalloc(&d_temp, next_n * sizeof(float)));
        temp_arrays.push_back(d_temp);
        temp_n = next_n;
    }

    float* d_block_sums = nullptr;
    if (num_blocks > 1) {
        CUDA_CHECK(cudaMalloc(&d_block_sums, num_blocks * sizeof(float)));
    }

    // CPU Scan (for validation)
    auto cpu_start = std::chrono::high_resolution_clock::now();
    cpu_inclusive_scan(h_in.data(), h_out_cpu.data(), N);
    auto cpu_end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double, std::milli> cpu_dur = cpu_end - cpu_start;

    // GPU Warmup & Latent Device Synchronize
    gpu_scan<true>(d_in, d_out, N, temp_arrays.data(), d_block_sums);
    CUDA_CHECK(cudaDeviceSynchronize());

    // Verify Inclusive Scan Correctness
    CUDA_CHECK(cudaMemcpy(h_out_gpu.data(), d_out, bytes, cudaMemcpyDeviceToHost));
    bool inclusive_passed = verify_results(h_out_cpu.data(), h_out_gpu.data(), N);

    // Benchmark Inclusive GPU Execution
    cudaEvent_t start, stop;
    CUDA_CHECK(cudaEventCreate(&start));
    CUDA_CHECK(cudaEventCreate(&stop));

    int iterations = 50;
    CUDA_CHECK(cudaEventRecord(start));
    for (int i = 0; i < iterations; ++i) {
        gpu_scan<true>(d_in, d_out, N, temp_arrays.data(), d_block_sums);
    }
    CUDA_CHECK(cudaEventRecord(stop));
    CUDA_CHECK(cudaEventSynchronize(stop));

    float ms = 0;
    CUDA_CHECK(cudaEventElapsedTime(&ms, start, stop));
    float avg_ms = ms / iterations;

    // Calculations based on 2x N floats (1 Read, 1 Write)
    double throughput_gb = (2.0 * bytes) / (avg_ms * 1e-3) / 1e9;
    double throughput_elem = N / (avg_ms * 1e-3) / 1e9;

    std::cout << "INCLUSIVE SCAN RESULTS:" << std::endl;
    std::cout << "  Verification : " << (inclusive_passed ? "PASSED" : "FAILED") << std::endl;
    std::cout << "  CPU Time     : " << cpu_dur.count() << " ms" << std::endl;
    std::cout << "  GPU Time     : " << avg_ms << " ms" << std::endl;
    std::cout << "  Throughput   : " << throughput_gb << " GB/s" << std::endl;
    std::cout << "  Rate         : " << throughput_elem << " Giga-elements/sec" << std::endl;

    // Verify Exclusive Scan Correctness
    cpu_exclusive_scan(h_in.data(), h_out_cpu.data(), N);
    gpu_scan<false>(d_in, d_out, N, temp_arrays.data(), d_block_sums);
    CUDA_CHECK(cudaDeviceSynchronize());
    CUDA_CHECK(cudaMemcpy(h_out_gpu.data(), d_out, bytes, cudaMemcpyDeviceToHost));
    bool exclusive_passed = verify_results(h_out_cpu.data(), h_out_gpu.data(), N);

    // Benchmark Exclusive GPU Execution
    CUDA_CHECK(cudaEventRecord(start));
    for (int i = 0; i < iterations; ++i) {
        gpu_scan<false>(d_in, d_out, N, temp_arrays.data(), d_block_sums);
    }
    CUDA_CHECK(cudaEventRecord(stop));
    CUDA_CHECK(cudaEventSynchronize(stop));

    CUDA_CHECK(cudaEventElapsedTime(&ms, start, stop));
    float avg_ms_ex = ms / iterations;

    double throughput_gb_ex = (2.0 * bytes) / (avg_ms_ex * 1e-3) / 1e9;
    double throughput_elem_ex = N / (avg_ms_ex * 1e-3) / 1e9;

    std::cout << "EXCLUSIVE SCAN RESULTS:" << std::endl;
    std::cout << "  Verification : " << (exclusive_passed ? "PASSED" : "FAILED") << std::endl;
    std::cout << "  GPU Time     : " << avg_ms_ex << " ms" << std::endl;
    std::cout << "  Throughput   : " << throughput_gb_ex << " GB/s" << std::endl;
    std::cout << "  Rate         : " << throughput_elem_ex << " Giga-elements/sec" << std::endl;

    // Clean up temporary workspaces
    CUDA_CHECK(cudaEventDestroy(start));
    CUDA_CHECK(cudaEventDestroy(stop));
    if (d_block_sums) CUDA_CHECK(cudaFree(d_block_sums));
    for (float* ptr : temp_arrays) {
        CUDA_CHECK(cudaFree(ptr));
    }
    CUDA_CHECK(cudaFree(d_in));
    CUDA_CHECK(cudaFree(d_out));
}

int main() {
    // Benchmark requested test sizes: 1M, 10M, 100M elements
    run_benchmark(1000000);
    run_benchmark(10000000);
    run_benchmark(100000000);
    return 0;
}