[BUG]: DeviceReduce::ReduceByKey is not run-to-run deterministic for fp64 sums over multi-tile runs #9995

// Reproducer: cub::DeviceReduce::ReduceByKey is not run-to-run
// deterministic for float64 addition over a key run spanning many tiles,
// despite the documented guarantee ("Provides 'run-to-run' determinism for
// pseudo-associative reduction (e.g., addition of floating point types) on
// the same GPU device" -- cub/device/device_reduce.cuh, ReduceByKey doc
// block; present in v3.1.4 and on main).
//
// Observed (NVIDIA T600 Laptop GPU, sm_75, CUDA 13.1, CUB 3.1.4 / 300104):
// experiment A returns 5 distinct bit patterns for the long run's aggregate
// across 200 identical invocations, spread over the last ~3 ulps.
//
// Protocol:
//   - keys/values are uploaded to the device ONCE and never touched again
//     (verified by checksumming the input buffers before and after);
//   - temp storage is allocated ONCE (constant size, printed);
//   - the SAME ReduceByKey call runs REPS times on the default stream; the
//     first (long) run's aggregate is compared bit-for-bit across reps.
//
// Controls (same code path, only the data changes):
//   A) fp64 values, one 50k-item key run  -> flips if the internal
//      combination order varies between invocations;
//   B) same keys, values = small integers stored as fp64. Integer sums
//      below 2^53 are EXACT in any association order, so if A's flips came
//      from memory corruption or a harness bug, B would flip too; if they
//      come from reordered floating-point combination, B is bit-stable;
//   C) fp64 values but every run fits one tile (64 items) -> bit-stable,
//      isolating the nondeterminism to the cross-tile combination.
//
// Build (no dependencies beyond the CUDA toolkit's bundled CCCL):
//   nvcc -O2 -arch=native \
//        -I$CUDA_HOME/targets/x86_64-linux/include/cccl \
//        cub_reducebykey_determinism.cu -o cub_reducebykey_determinism
//   ./cub_reducebykey_determinism

#include <cub/device/device_reduce.cuh>
#include <cub/version.cuh>

#include <cstdint>
#include <cstdio>
#include <cstring>
#include <map>
#include <random>
#include <vector>

#define CK(x)                                                            \
    if ((x) != cudaSuccess) {                                            \
        printf("cuda error %s @%d\n", cudaGetErrorString(x), __LINE__);  \
        return 1;                                                        \
    }

struct Sum {
    __host__ __device__ double operator()(double a, double b) const {
        return a + b;
    }
};

static uint64_t fnv1a(const void* p, size_t n) {
    const auto* b = static_cast<const unsigned char*>(p);
    uint64_t h = 1469598103934665603ULL;
    for (size_t i = 0; i < n; ++i) h = (h ^ b[i]) * 1099511628211ULL;
    return h;
}

static int experiment(const char* name, const std::vector<uint64_t>& keys,
                      const std::vector<double>& vals, int reps) {
    const size_t n = keys.size();
    uint64_t *d_keys, *d_ukeys;
    double *d_vals, *d_aggs;
    size_t* d_nruns;
    CK(cudaMalloc(&d_keys, n * 8));
    CK(cudaMalloc(&d_ukeys, n * 8));
    CK(cudaMalloc(&d_vals, n * 8));
    CK(cudaMalloc(&d_aggs, n * 8));
    CK(cudaMalloc(&d_nruns, 8));
    CK(cudaMemcpy(d_keys, keys.data(), n * 8, cudaMemcpyHostToDevice));
    CK(cudaMemcpy(d_vals, vals.data(), n * 8, cudaMemcpyHostToDevice));
    const uint64_t in_hash =
        fnv1a(keys.data(), n * 8) ^ fnv1a(vals.data(), n * 8);

    size_t temp_bytes = 0;
    cub::DeviceReduce::ReduceByKey(nullptr, temp_bytes, d_keys, d_ukeys,
                                   d_vals, d_aggs, d_nruns, Sum{}, n);
    void* d_temp;
    CK(cudaMalloc(&d_temp, temp_bytes));

    std::map<uint64_t, int> observed;  // aggregate[0] bit pattern -> count
    size_t nruns_first = 0;
    for (int rep = 0; rep < reps; ++rep) {
        CK(cudaMemset(d_aggs, 0xEE, n * 8));  // poison: must be overwritten
        cub::DeviceReduce::ReduceByKey(d_temp, temp_bytes, d_keys, d_ukeys,
                                       d_vals, d_aggs, d_nruns, Sum{}, n);
        uint64_t bits = 0;
        size_t nruns = 0;
        CK(cudaMemcpy(&bits, d_aggs, 8, cudaMemcpyDeviceToHost));
        CK(cudaMemcpy(&nruns, d_nruns, 8, cudaMemcpyDeviceToHost));
        if (rep == 0) nruns_first = nruns;
        if (nruns != nruns_first) printf("  !! run count changed\n");
        ++observed[bits];
    }

    // Re-download the input to prove it was never modified.
    std::vector<uint64_t> keys2(n);
    std::vector<double> vals2(n);
    CK(cudaMemcpy(keys2.data(), d_keys, n * 8, cudaMemcpyDeviceToHost));
    CK(cudaMemcpy(vals2.data(), d_vals, n * 8, cudaMemcpyDeviceToHost));
    const uint64_t out_hash =
        fnv1a(keys2.data(), n * 8) ^ fnv1a(vals2.data(), n * 8);

    // Host reference: strict left-to-right sum of run 0.
    double seq = 0;
    for (size_t i = 0; i < n && keys[i] == keys[0]; ++i) seq += vals[i];

    printf("%s: n=%zu runs=%zu temp=%zu B, input %s, %d reps -> %zu "
           "distinct bit patterns for aggregate[0]\n",
           name, n, nruns_first, temp_bytes,
           in_hash == out_hash ? "UNMODIFIED" : "MODIFIED(!)", reps,
           observed.size());
    for (const auto& [bits, count] : observed) {
        double v;
        memcpy(&v, &bits, 8);
        printf("    %.17g  (0x%016llx)  x%d%s\n", v,
               static_cast<unsigned long long>(bits), count,
               memcmp(&v, &seq, 8) == 0 ? "  == sequential host sum" : "");
    }
    cudaFree(d_keys);
    cudaFree(d_ukeys);
    cudaFree(d_vals);
    cudaFree(d_aggs);
    cudaFree(d_nruns);
    cudaFree(d_temp);
    return 0;
}

int main() {
    int dev = 0;
    cudaDeviceProp prop{};
    CK(cudaGetDeviceProperties(&prop, dev));
    int drv = 0, rt = 0;
    cudaDriverGetVersion(&drv);
    cudaRuntimeGetVersion(&rt);
    printf("GPU: %s (sm_%d%d), driver %d, runtime %d, CUB %d\n\n", prop.name,
           prop.major, prop.minor, drv, rt, CUB_VERSION);

    const size_t n = 65536;
    std::mt19937_64 rng(42);
    std::uniform_real_distribution<double> dist(-1.0, 1.0);

    // A: one 50k-item run of key 0 (fp64 values), then short runs.
    std::vector<uint64_t> keys(n);
    std::vector<double> vals_fp(n), vals_int(n);
    for (size_t i = 0; i < n; ++i) {
        keys[i] = (i < 50000) ? 0 : 1 + (i - 50000) / 64;
        vals_fp[i] = (i % 7 == 0) ? dist(rng) * 1e-2 : 0.0;
        vals_int[i] = static_cast<double>(static_cast<int>(rng() % 1000));
    }
    if (experiment("A) fp64, 50k-item run   ", keys, vals_fp, 200)) return 1;

    // B: identical keys, integer-valued doubles (any order sums exactly).
    if (experiment("B) int-as-fp64, same run", keys, vals_int, 200)) return 1;

    // C: fp64 values, all runs short (64 items, single-tile).
    std::vector<uint64_t> keys_short(n);
    for (size_t i = 0; i < n; ++i) keys_short[i] = i / 64;
    if (experiment("C) fp64, 64-item runs   ", keys_short, vals_fp, 200))
        return 1;
    return 0;
}