// Dropout kernels (Phase T18).
//
// A counter-based (stateless) RNG: the keep/drop decision for element `i` is a
// pure function of (seed, i) — no global RNG state is advanced. This is what
// makes dropout compatible with activation recomputation (T13): when a
// checkpointed block re-runs its forward in backward, the SAME seed regenerates
// the SAME mask, so the recomputed activations / grads stay bit-identical to the
// forward (no mask drift).
//
// Inverted dropout: at training time kept elements are scaled by 1/(1-p) so the
// expectation E[out] == x (no inference-time rescale needed; eval is identity,
// handled in Rust by simply not calling dropout).
//
//   key  = seed ^ (i * GOLDEN)
//   h    = splitmix64(key)             // a few rounds of xorshift/multiply
//   u    = (h >> 40) / 2^24  in [0,1)  // 24-bit uniform
//   keep = u >= p                      // Bernoulli(keep = 1-p)
//   out  = keep ? x * scale : 0        // scale = 1/(1-p)
//   mask = keep ? scale : 0            // cached for backward (dx = d * mask)
//
// fp32 + bf16 variants: bf16 loads/stores half-size activations but the uniform
// `u` is always computed in fp32, so the mask distribution is identical across
// dtypes (drop decisions don't depend on bf16 rounding). The mask buffer is fp32
// in both cases (it stores `scale` or 0 — exactly representable, tiny relative to
// the activation, reused only elementwise in backward).

#include <cuda_bf16.h>
#include <stdint.h>

extern "C" {

// splitmix64: cheap, well-mixed counter hash. Maps a 64-bit counter to a 64-bit
// pseudo-random output; we only need the high bits for a uniform.
__device__ __forceinline__ uint64_t splitmix64(uint64_t x) {
    x += 0x9E3779B97F4A7C15ULL;
    x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ULL;
    x = (x ^ (x >> 27)) * 0x94D049BB133111EBULL;
    return x ^ (x >> 31);
}

// Uniform [0,1) for element i under `seed`, computed in fp32 (dtype-independent).
__device__ __forceinline__ float dropout_uniform(uint64_t seed, int i) {
    uint64_t key = seed ^ ((uint64_t)i * 0x9E3779B97F4A7C15ULL);
    uint64_t h = splitmix64(key);
    // Top 24 bits → [0,1) with 2^-24 resolution.
    return (float)(h >> 40) * (1.0f / 16777216.0f);  // 1/2^24
}

// fp32 forward: out[i] = keep ? x[i]*scale : 0 ; mask[i] = keep ? scale : 0.
__global__ void dropout_fwd_f32_k(const float* x, float* out, float* mask,
                                  float p, float scale, uint64_t seed, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float keep = (dropout_uniform(seed, i) >= p) ? scale : 0.0f;
        mask[i] = keep;
        out[i] = x[i] * keep;
    }
}
void launch_dropout_fwd_f32(const float* x, float* out, float* mask, float p,
                            float scale, uint64_t seed, int n, void* s) {
    int blk = 256, grid = (n + blk - 1) / blk;
    dropout_fwd_f32_k<<<grid, blk, 0, (cudaStream_t)s>>>(x, out, mask, p, scale,
                                                         seed, n);
}

// Backward applies the SAME cached mask elementwise: dx[i] = d[i] * mask[i].
__global__ void dropout_bwd_f32_k(const float* d, const float* mask, float* dx,
                                  int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) dx[i] = d[i] * mask[i];
}
void launch_dropout_bwd_f32(const float* d, const float* mask, float* dx, int n,
                            void* s) {
    int blk = 256, grid = (n + blk - 1) / blk;
    dropout_bwd_f32_k<<<grid, blk, 0, (cudaStream_t)s>>>(d, mask, dx, n);
}

// bf16 forward: activation is bf16; mask is fp32 (stores `scale` or 0). Uniform
// is fp32, so the mask matches the fp32 path bit-for-bit (same drop decisions).
__global__ void dropout_fwd_bf16_k(const __nv_bfloat16* x, __nv_bfloat16* out,
                                   float* mask, float p, float scale,
                                   uint64_t seed, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float keep = (dropout_uniform(seed, i) >= p) ? scale : 0.0f;
        mask[i] = keep;
        out[i] = __float2bfloat16(__bfloat162float(x[i]) * keep);
    }
}
void launch_dropout_fwd_bf16(const void* x, void* out, float* mask, float p,
                             float scale, uint64_t seed, int n, void* s) {
    int blk = 256, grid = (n + blk - 1) / blk;
    dropout_fwd_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
        (const __nv_bfloat16*)x, (__nv_bfloat16*)out, mask, p, scale, seed, n);
}

__global__ void dropout_bwd_bf16_k(const __nv_bfloat16* d, const float* mask,
                                   __nv_bfloat16* dx, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) dx[i] = __float2bfloat16(__bfloat162float(d[i]) * mask[i]);
}
void launch_dropout_bwd_bf16(const void* d, const float* mask, void* dx, int n,
                             void* s) {
    int blk = 256, grid = (n + blk - 1) / blk;
    dropout_bwd_bf16_k<<<grid, blk, 0, (cudaStream_t)s>>>(
        (const __nv_bfloat16*)d, mask, (__nv_bfloat16*)dx, n);
}

}  // extern "C"