"""Fused CUDA grid-foraging env + 3x MinGRU(h=256) policy rollout.

Structure
---------
* `run()` is a persistent megakernel: one block owns a slice of envs and runs the
  whole horizon without ever returning to the host. Env step, all three MinGRU
  layers, the action head and the argmax all live inside the kernel.
* Layer 0 is algebraically folded. gates0 = w_gru[0] @ (w_enc @ obs + b_enc)
  = (w_gru[0] @ w_enc) @ obs + (w_gru[0] @ b_enc); both factors are constants, so
  the 768x256 GEMM collapses to 768x4. Exact same function, 1.49x less work and
  1.49x less weight traffic (the folded constants are built once per call).
* fp32 FFMA throughout. The check's numeric-stress case allows atol=1e-6 on
  logits/state and the greedy argmax has a measured min top1-top2 gap of 1.1e-5,
  so fp16/bf16/tf32 tensor cores (~1e-5 error) are not usable here.
* The reference advances the LCG for *every* env iff ANY env in the batch hit food
  this step (`if hit.any():`), which is a batch-global condition -- so each
  timestep ends in a grid-wide barrier over a co-resident grid.
"""
from __future__ import annotations

import os

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

BOARD = 11
OBS_DIM = 4
HIDDEN = 256
GRU_LAYERS = 3
NUM_ACTIONS = 4
GRU_OUT = 3 * HIDDEN

_CUDA_SRC = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

#define HID 256
#define NT 4
#define JG 64
#define BOARDI 11

// Block geometry: blockDim = EG*JG threads; thread (eg=tid/64, jg=tid%64) owns
// MT envs x NT hidden units. EC = EG*MT envs per chunk.
// Total accumulator registers per block = 3*EC*HID regardless of how they are
// spread over threads, so raising EG (more threads, same EC) buys occupancy for
// free. The register file (65536/SM) caps 3*EC*256 + operands.
template <int EG, int MT>
struct Cfg {
    static constexpr int NTHREAD = EG * JG;
    static constexpr int EC = EG * MT;
    static constexpr int STRIDE = EC + 4;
    // Per-thread register budget: the file allows 65536/NTHREAD at 1 block/SM,
    // but the hardware caps a thread at 255 regardless.
    static constexpr int RAW = 65536 / NTHREAD;
    static constexpr int REGS = RAW > 240 ? 240 : RAW;
    // Unrolling the k-loop replicates the hv/wv operands on top of acc[3][MT][4].
    static constexpr int KU = (12 * MT + 2 * (MT + 12) + 28 <= REGS) ? 2 : 1;
    static constexpr int SHM = (HID * STRIDE + EC * 4 + EC * 4) * 4;
};

__device__ __forceinline__ float sigmoidf_ref(float x) {
    // matches ATen: 1 / (1 + exp(-x)) in fp32, no fast-math
    return 1.0f / (1.0f + expf(-x));
}

__device__ __forceinline__ unsigned long long lcg_step(unsigned long long r) {
    return (r * 6364136223846793005ULL + 1ULL) & 0x7FFFFFFFFFFFFFFFULL;
}

// ---------------------------------------------------------------- grid barrier
// The reference advances the RNG for every env iff ANY env in the whole batch hit
// food this step, so each timestep needs a batch-global agreement point. The grid
// is sized co-resident (<= SMs * blocks_per_SM), so this cannot deadlock.
__device__ __forceinline__ int grid_sync_any(unsigned int* bar, int* flag,
                                             int blk_any, unsigned int target) {
    __shared__ int s_any;
    if (threadIdx.x == 0 && blk_any) atomicExch(flag, 1);
    __threadfence();
    __syncthreads();
    if (threadIdx.x == 0) {
        atomicAdd(bar, 1u);
        while (*(volatile unsigned int*)bar < target) __nanosleep(64);
        s_any = *(volatile int*)flag;
    }
    __syncthreads();
    return s_any;
}

// ------------------------------------------------------------- MinGRU epilogue
// out = st + sigmoid(zg) * (tanh(zh) - st);  p = sigmoid(zp);  h = p*out + (1-p)*h
// hk holds the incoming h and is overwritten with the outgoing h in place; st is
// read/written in global. Per-env so nothing stacks on acc's live range.
template <int EG, int MT>
__device__ __forceinline__ void mingru_epilogue(float acc[3][MT][NT], float* __restrict__ hk,
                                                float* __restrict__ stp, int e0, int j0) {
    constexpr int STRIDE = Cfg<EG, MT>::STRIDE;
    // Must stay fully unrolled: a runtime `m` would make acc[..][m][..] a dynamic
    // index and demote the whole accumulator array to local memory.
#pragma unroll
    for (int m = 0; m < MT; m++) {
        float4 s4 = *(const float4*)&stp[(long long)(e0 + m) * HID + j0];
        float sv[NT] = {s4.x, s4.y, s4.z, s4.w};
        float o4[NT];
#pragma unroll
        for (int n = 0; n < NT; n++) {
            float hold = hk[(j0 + n) * STRIDE + e0 + m];
            float s = sv[n];
            float o = s + sigmoidf_ref(acc[1][m][n]) * (tanhf(acc[0][m][n]) - s);
            float p = sigmoidf_ref(acc[2][m][n]);
            o4[n] = o;
            hk[(j0 + n) * STRIDE + e0 + m] = p * o + (1.0f - p) * hold;
        }
        *(float4*)&stp[(long long)(e0 + m) * HID + j0] =
            make_float4(o4[0], o4[1], o4[2], o4[3]);
    }
}

// ------------------------------------------------------- one policy step (chunk)
// Mapping (256 threads): eg = tid>>6 owns envs [eg*MT, eg*MT+MT),
//                        jg = tid&63 owns hidden units [jg*4, jg*4+4).
// Every (env, j) pair is owned by exactly one thread, so the MinGRU epilogue
// touches only its own hk slots -> race-free without staging.
// Consumes s_obs[EC][4], advances state in global, leaves h3 in hk[j][e].
template <int EG, int MT>
__device__ __forceinline__ void policy_chunk(
    const float* __restrict__ Afold,  // (4,3,256)
    const float* __restrict__ cfold,  // (3,256)
    const float* __restrict__ Wenc,   // (4,256)
    const float* __restrict__ benc,   // (256)
    const float* __restrict__ Wg,     // (2,256,3,256) layers 1,2 as [k][gate][j]
    float* __restrict__ hk,           // shmem [256][STRIDE]
    const float* __restrict__ s_obs,  // shmem [EC][4]
    float* __restrict__ g_state,      // (3,Sstride,256)
    long long Sstride, long long env_base) {
    constexpr int STRIDE = Cfg<EG, MT>::STRIDE;

    const int tid = threadIdx.x;
    const int e0 = (tid >> 6) * MT;
    const int j0 = (tid & 63) * NT;

    // ---------------- layer 0 (folded)
    // h0 = Wenc@obs + benc goes into hk *before* acc exists, so its operands never
    // stack on top of acc[3][MT][4] (which alone is 192 registers at MT=16). That
    // makes layer 0 structurally identical to layers 1/2: hk holds the incoming h,
    // the epilogue reads it back as `hold` and overwrites it with the outgoing h.
    // Each thread only ever touches its own (j, e) slots, so no barrier is needed.
    {
        float4 bv = *(const float4*)&benc[j0];
        float4 we[4];
#pragma unroll
        for (int o = 0; o < 4; o++) we[o] = *(const float4*)&Wenc[o * HID + j0];
#pragma unroll 2
        for (int m = 0; m < MT; m++) {
            float h0[NT] = {bv.x, bv.y, bv.z, bv.w};
#pragma unroll
            for (int o = 0; o < 4; o++) {
                float ob = s_obs[(e0 + m) * 4 + o];
                h0[0] = fmaf(we[o].x, ob, h0[0]);
                h0[1] = fmaf(we[o].y, ob, h0[1]);
                h0[2] = fmaf(we[o].z, ob, h0[2]);
                h0[3] = fmaf(we[o].w, ob, h0[3]);
            }
#pragma unroll
            for (int n = 0; n < NT; n++) hk[(j0 + n) * STRIDE + e0 + m] = h0[n];
        }
    }

    // acc is the entire register budget at large MT; keep everything else out of
    // its live range.
    float acc[3][MT][NT];
#pragma unroll
    for (int g = 0; g < 3; g++) {
        float4 cv = *(const float4*)&cfold[g * HID + j0];
#pragma unroll
        for (int m = 0; m < MT; m++) {
            acc[g][m][0] = cv.x; acc[g][m][1] = cv.y;
            acc[g][m][2] = cv.z; acc[g][m][3] = cv.w;
        }
    }
#pragma unroll 1
    for (int o = 0; o < 4; o++) {
        float4 av[3];
#pragma unroll
        for (int g = 0; g < 3; g++) av[g] = *(const float4*)&Afold[(o * 3 + g) * HID + j0];
#pragma unroll
        for (int m = 0; m < MT; m++) {
            float ob = s_obs[(e0 + m) * 4 + o];
#pragma unroll
            for (int g = 0; g < 3; g++) {
                const float4 a = av[g];
                acc[g][m][0] = fmaf(a.x, ob, acc[g][m][0]);
                acc[g][m][1] = fmaf(a.y, ob, acc[g][m][1]);
                acc[g][m][2] = fmaf(a.z, ob, acc[g][m][2]);
                acc[g][m][3] = fmaf(a.w, ob, acc[g][m][3]);
            }
        }
    }
    mingru_epilogue<EG, MT>(acc, hk, g_state + env_base * HID, e0, j0);
    __syncthreads();

    // ---------------- layers 1,2 (the two big GEMMs)
#pragma unroll 1
    for (int L = 0; L < 2; L++) {
        const float* __restrict__ Wl = Wg + (long long)L * (HID * 3 * HID);
#pragma unroll
        for (int g = 0; g < 3; g++)
#pragma unroll
            for (int m = 0; m < MT; m++)
#pragma unroll
                for (int n = 0; n < NT; n++) acc[g][m][n] = 0.0f;

#pragma unroll(Cfg<EG, MT>::KU)
        for (int k = 0; k < HID; k++) {
            float hv[MT];
#pragma unroll
            for (int m = 0; m < MT; m += 4) {
                float4 v = *(const float4*)&hk[k * STRIDE + e0 + m];
                hv[m] = v.x; hv[m + 1] = v.y; hv[m + 2] = v.z; hv[m + 3] = v.w;
            }
            float4 wv[3];
#pragma unroll
            for (int g = 0; g < 3; g++)
                wv[g] = __ldg((const float4*)&Wl[(k * 3 + g) * HID + j0]);
#pragma unroll
            for (int m = 0; m < MT; m++) {
#pragma unroll
                for (int g = 0; g < 3; g++) {
                    const float4 w = wv[g];
                    acc[g][m][0] = fmaf(w.x, hv[m], acc[g][m][0]);
                    acc[g][m][1] = fmaf(w.y, hv[m], acc[g][m][1]);
                    acc[g][m][2] = fmaf(w.z, hv[m], acc[g][m][2]);
                    acc[g][m][3] = fmaf(w.w, hv[m], acc[g][m][3]);
                }
            }
        }
        __syncthreads();  // all reads of hk done before we overwrite it
        mingru_epilogue<EG, MT>(acc, hk,
                                g_state + ((long long)(L + 1) * Sstride + env_base) * HID,
                                e0, j0);
        __syncthreads();
    }
}

// ------------------------------------------------------------- the megakernel
template <int EG, int MT>
__global__ __launch_bounds__(EG* JG, 1) void mega_kernel(
    const float* __restrict__ Afold, const float* __restrict__ cfold,
    const float* __restrict__ Wenc, const float* __restrict__ benc,
    const float* __restrict__ Wg, const float* __restrict__ Wa,
    const float* __restrict__ ba,
    int* __restrict__ g_ax, int* __restrict__ g_ay,
    int* __restrict__ g_fx, int* __restrict__ g_fy,
    unsigned long long* __restrict__ g_rng, float* __restrict__ g_rew,
    float* __restrict__ g_state, float* __restrict__ g_logits,
    int* __restrict__ g_hit,
    unsigned int* __restrict__ g_bar, int* __restrict__ g_anyhit,
    long long N, long long Sstride, int H, long long envs_per_block, int nblocks) {
    constexpr int EC = Cfg<EG, MT>::EC;
    constexpr int STRIDE = Cfg<EG, MT>::STRIDE;

    extern __shared__ char smem_raw[];
    float* hk = (float*)smem_raw;
    float* s_obs = hk + (long long)HID * STRIDE;  // [EC][4]
    float* s_lg = s_obs + EC * 4;                 // [EC][4]
    __shared__ int s_blk_any;

    const int tid = threadIdx.x;
    const long long env_base = (long long)blockIdx.x * envs_per_block;
    const long long env_end = min(N, env_base + envs_per_block);
    const int n_chunks = (int)(envs_per_block / EC);

    for (int t = 0; t < H; t++) {
        if (tid == 0) s_blk_any = 0;
        __syncthreads();

        for (int c = 0; c < n_chunks; c++) {
            const long long cb = env_base + (long long)c * EC;
            const int live = (int)min((long long)EC, max(0LL, env_end - cb));

            // obs. torch rewrites `x / 11` (cpu scalar divisor) as `x * (1/11)`.
            if (tid < EC) {
                int e = tid;
                if (e < live) {
                    int ax = g_ax[cb + e], ay = g_ay[cb + e];
                    int fx = g_fx[cb + e], fy = g_fy[cb + e];
                    s_obs[e * 4 + 0] = (float)(fx - ax) * (1.0f / (float)BOARDI);
                    s_obs[e * 4 + 1] = (float)(fy - ay) * (1.0f / (float)BOARDI);
                    s_obs[e * 4 + 2] = (float)ax * (1.0f / (float)(BOARDI - 1));
                    s_obs[e * 4 + 3] = (float)ay * (1.0f / (float)(BOARDI - 1));
                } else {
                    s_obs[e * 4 + 0] = 0.f; s_obs[e * 4 + 1] = 0.f;
                    s_obs[e * 4 + 2] = 0.f; s_obs[e * 4 + 3] = 0.f;
                }
            }
            __syncthreads();

            policy_chunk<EG, MT>(Afold, cfold, Wenc, benc, Wg, hk, s_obs, g_state, Sstride, cb);

            // logits = Wa @ h3 + ba
            if (tid < EC * 4) {
                int e = tid >> 2, a = tid & 3;
                float s = 0.f;
                for (int j = 0; j < HID; j++) s = fmaf(Wa[a * HID + j], hk[j * STRIDE + e], s);
                s_lg[e * 4 + a] = s + ba[a];
            }
            __syncthreads();

            // argmax (ties -> lowest index, like torch) + env move
            if (tid < live) {
                int e = tid;
                float l0 = s_lg[e * 4 + 0], l1 = s_lg[e * 4 + 1];
                float l2 = s_lg[e * 4 + 2], l3 = s_lg[e * 4 + 3];
                int a = 0; float best = l0;
                if (l1 > best) { best = l1; a = 1; }
                if (l2 > best) { best = l2; a = 2; }
                if (l3 > best) { best = l3; a = 3; }
                if (t == H - 1) {
                    g_logits[(cb + e) * 4 + 0] = l0; g_logits[(cb + e) * 4 + 1] = l1;
                    g_logits[(cb + e) * 4 + 2] = l2; g_logits[(cb + e) * 4 + 3] = l3;
                }
                int ax = g_ax[cb + e], ay = g_ay[cb + e];
                int dx = (a == 3) - (a == 2);
                int dy = (a == 1) - (a == 0);
                ax = min(max(ax + dx, 0), BOARDI - 1);
                ay = min(max(ay + dy, 0), BOARDI - 1);
                int hit = (ax == g_fx[cb + e] && ay == g_fy[cb + e]) ? 1 : 0;
                g_ax[cb + e] = ax; g_ay[cb + e] = ay;
                g_hit[cb + e] = hit;
                if (hit) {
                    g_rew[cb + e] += 1.0f;
                    atomicOr(&s_blk_any, 1);
                }
            }
            __syncthreads();
        }

        // batch-global agreement: did ANY env hit this step?
        int any = grid_sync_any(g_bar, &g_anyhit[t], s_blk_any, (unsigned int)(t + 1) * nblocks);

        // rng advances for EVERY env iff any hit; food respawns only for hitters
        if (any) {
            for (long long e = env_base + tid; e < env_end; e += blockDim.x) {
                unsigned long long r = g_rng[e];
                r = lcg_step(r);
                int nfx = (int)(r % (unsigned long long)BOARDI);
                r = lcg_step(r);
                int nfy = (int)(r % (unsigned long long)BOARDI);
                g_rng[e] = r;
                if (g_hit[e]) { g_fx[e] = nfx; g_fy[e] = nfy; }
            }
        }
        __syncthreads();
    }
}

// ---------------------------------------------- single-step policy (for check)
// Shares policy_chunk with the megakernel, so check.py's tight atol=1e-6
// policy_forward test validates the fused path's numerics directly.
template <int EG, int MT>
__global__ __launch_bounds__(EG* JG, 1) void policy_kernel(
    const float* __restrict__ Afold, const float* __restrict__ cfold,
    const float* __restrict__ Wenc, const float* __restrict__ benc,
    const float* __restrict__ Wg, const float* __restrict__ Wa,
    const float* __restrict__ ba, const float* __restrict__ Wv,
    const float* __restrict__ bv,
    const float* __restrict__ obs_in, float* __restrict__ g_state,
    float* __restrict__ g_logits, float* __restrict__ g_value,
    long long N, long long Sstride) {
    constexpr int EC = Cfg<EG, MT>::EC;
    constexpr int STRIDE = Cfg<EG, MT>::STRIDE;

    extern __shared__ char smem_raw[];
    float* hk = (float*)smem_raw;
    float* s_obs = hk + (long long)HID * STRIDE;

    const int tid = threadIdx.x;
    const long long cb = (long long)blockIdx.x * EC;
    const int live = (int)min((long long)EC, max(0LL, N - cb));

    if (tid < EC) {
        int e = tid;
#pragma unroll
        for (int o = 0; o < 4; o++)
            s_obs[e * 4 + o] = (e < live) ? obs_in[(cb + e) * 4 + o] : 0.f;
    }
    __syncthreads();

    policy_chunk<EG, MT>(Afold, cfold, Wenc, benc, Wg, hk, s_obs, g_state, Sstride, cb);

    if (tid < EC * 4) {
        int e = tid >> 2, a = tid & 3;
        if (e < live) {
            float s = 0.f;
            for (int j = 0; j < HID; j++) s = fmaf(Wa[a * HID + j], hk[j * STRIDE + e], s);
            g_logits[(cb + e) * 4 + a] = s + ba[a];
        }
    }
    // hk is read-only here, so the value head can reuse the same threads.
    if (tid < EC) {
        int e = tid;
        if (e < live) {
            float s = 0.f;
            for (int j = 0; j < HID; j++) s = fmaf(Wv[j], hk[j * STRIDE + e], s);
            g_value[cb + e] = s + bv[0];
        }
    }
}

// ------------------------------------------------------------ standalone env_step
__global__ void env_k1(const float* __restrict__ agent, const float* __restrict__ food,
                       const long long* __restrict__ act, float* __restrict__ agent_out,
                       int* __restrict__ hit, float* __restrict__ rew,
                       int* __restrict__ any, long long N) {
    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;
    float ax = agent[i * 2 + 0], ay = agent[i * 2 + 1];
    long long a = act[i];
    float dx = (float)((a == 3) - (a == 2));
    float dy = (float)((a == 1) - (a == 0));
    ax = fminf(fmaxf(ax + dx, 0.f), (float)(BOARDI - 1));
    ay = fminf(fmaxf(ay + dy, 0.f), (float)(BOARDI - 1));
    agent_out[i * 2 + 0] = ax; agent_out[i * 2 + 1] = ay;
    int h = (ax == food[i * 2 + 0] && ay == food[i * 2 + 1]) ? 1 : 0;
    hit[i] = h;
    rew[i] = (float)h;
    if (h) atomicOr(any, 1);
}

__global__ void env_k2(const float* __restrict__ food, const long long* __restrict__ rng_in,
                       const int* __restrict__ hit, float* __restrict__ food_out,
                       long long* __restrict__ rng_out, const int* __restrict__ any,
                       long long N) {
    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;
    unsigned long long r = (unsigned long long)rng_in[i];
    float fx = food[i * 2 + 0], fy = food[i * 2 + 1];
    if (*any) {
        r = lcg_step(r);
        int nfx = (int)(r % (unsigned long long)BOARDI);
        r = lcg_step(r);
        int nfy = (int)(r % (unsigned long long)BOARDI);
        if (hit[i]) { fx = (float)nfx; fy = (float)nfy; }
    }
    food_out[i * 2 + 0] = fx; food_out[i * 2 + 1] = fy;
    rng_out[i] = (long long)r;
}

// ------------------------------------------------------------------- launchers
#define P(t) ((float*)t.data_ptr())

static inline void cuda_ok(const char* what) {
    cudaError_t e = cudaGetLastError();
    TORCH_CHECK(e == cudaSuccess, what, ": ", cudaGetErrorString(e));
}

// (EG, MT) configs. EG*JG threads, EC = EG*MT envs per chunk. Constrained by
// 3*EC*256 accumulator registers + operands <= 65536 per SM.
// Each weight element is fetched once per env-group, so L1 traffic ~ EG while
// L2 traffic (blocks*chunks*1.6MB) is independent of EG -- hence the EG=2 rows.
#define FOR_EACH_CFG(F) \
    F(2, 12) F(2, 16) F(4, 4) F(4, 8) F(4, 12) F(4, 16) F(8, 4) F(8, 8)

#define MEGA_CASE(EG_, MT_)                                                             \
    if (EG == EG_ && MT == MT_) {                                                       \
        constexpr size_t shm = Cfg<EG_, MT_>::SHM;                                      \
        cudaFuncSetAttribute(mega_kernel<EG_, MT_>,                                     \
                             cudaFuncAttributeMaxDynamicSharedMemorySize, (int)shm);    \
        mega_kernel<EG_, MT_><<<nblocks, EG_ * JG, shm>>>(                              \
            P(Afold), P(cfold), P(Wenc), P(benc), P(Wg), P(Wa), P(ba),                  \
            (int*)ax.data_ptr(), (int*)ay.data_ptr(), (int*)fx.data_ptr(),              \
            (int*)fy.data_ptr(), (unsigned long long*)rng.data_ptr(),                   \
            P(rew), P(state), P(logits), (int*)hit.data_ptr(),                          \
            (unsigned int*)bar.data_ptr(), (int*)anyhit.data_ptr(),                     \
            N, Sstride, H, envs_per_block, nblocks);                                    \
        cuda_ok("mega_kernel");                                                         \
        return;                                                                         \
    }

void mega_launch(torch::Tensor Afold, torch::Tensor cfold, torch::Tensor Wenc,
                 torch::Tensor benc, torch::Tensor Wg, torch::Tensor Wa, torch::Tensor ba,
                 torch::Tensor ax, torch::Tensor ay, torch::Tensor fx, torch::Tensor fy,
                 torch::Tensor rng, torch::Tensor rew, torch::Tensor state,
                 torch::Tensor logits, torch::Tensor hit, torch::Tensor bar,
                 torch::Tensor anyhit, long long N, long long Sstride, int H,
                 long long envs_per_block, int nblocks, int EG, int MT) {
    FOR_EACH_CFG(MEGA_CASE)
    TORCH_CHECK(false, "unsupported cfg EG=", EG, " MT=", MT);
}

#define POL_CASE(EG_, MT_)                                                              \
    if (EG == EG_ && MT == MT_) {                                                       \
        constexpr size_t shm = Cfg<EG_, MT_>::SHM;                                      \
        cudaFuncSetAttribute(policy_kernel<EG_, MT_>,                                   \
                             cudaFuncAttributeMaxDynamicSharedMemorySize, (int)shm);    \
        int nb = (int)((N + Cfg<EG_, MT_>::EC - 1) / Cfg<EG_, MT_>::EC);                \
        policy_kernel<EG_, MT_><<<nb, EG_ * JG, shm>>>(                                 \
            P(Afold), P(cfold), P(Wenc), P(benc), P(Wg), P(Wa), P(ba), P(Wv), P(bv),    \
            P(obs), P(state), P(logits), P(value), N, Sstride);                         \
        cuda_ok("policy_kernel");                                                       \
        return;                                                                         \
    }

void policy_launch(torch::Tensor Afold, torch::Tensor cfold, torch::Tensor Wenc,
                   torch::Tensor benc, torch::Tensor Wg, torch::Tensor Wa,
                   torch::Tensor ba, torch::Tensor Wv, torch::Tensor bv,
                   torch::Tensor obs, torch::Tensor state, torch::Tensor logits,
                   torch::Tensor value, long long N, long long Sstride, int EG, int MT) {
    FOR_EACH_CFG(POL_CASE)
    TORCH_CHECK(false, "unsupported cfg EG=", EG, " MT=", MT);
}

void env_step_launch(torch::Tensor agent, torch::Tensor food, torch::Tensor act,
                     torch::Tensor agent_out, torch::Tensor food_out, torch::Tensor rng_in,
                     torch::Tensor rng_out, torch::Tensor rew, torch::Tensor hit,
                     torch::Tensor any, long long N) {
    int th = 256, bl = (int)((N + th - 1) / th);
    env_k1<<<bl, th>>>(P(agent), P(food), (const long long*)act.data_ptr(), P(agent_out),
                       (int*)hit.data_ptr(), P(rew), (int*)any.data_ptr(), N);
    env_k2<<<bl, th>>>(P(food), (const long long*)rng_in.data_ptr(), (const int*)hit.data_ptr(),
                       P(food_out), (long long*)rng_out.data_ptr(), (const int*)any.data_ptr(), N);
    cuda_ok("env_step");
}

#define OCC_CASE(EG_, MT_)                                                              \
    if (EG == EG_ && MT == MT_) {                                                       \
        cudaFuncSetAttribute(mega_kernel<EG_, MT_>,                                     \
                             cudaFuncAttributeMaxDynamicSharedMemorySize,               \
                             (int)Cfg<EG_, MT_>::SHM);                                  \
        cudaOccupancyMaxActiveBlocksPerMultiprocessor(&nb, mega_kernel<EG_, MT_>,       \
                                                      EG_ * JG, Cfg<EG_, MT_>::SHM);    \
        return nb;                                                                      \
    }

int mega_max_blocks_per_sm(int EG, int MT) {
    int nb = 0;
    FOR_EACH_CFG(OCC_CASE)
    return 0;
}
"""

_CPP_SRC = r"""
void mega_launch(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
                 torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
                 torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
                 torch::Tensor, torch::Tensor, torch::Tensor, long long, long long, int,
                 long long, int, int, int);
void policy_launch(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
                   torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
                   torch::Tensor, torch::Tensor, torch::Tensor, long long, long long, int, int);
void env_step_launch(torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
                     torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
                     long long);
int mega_max_blocks_per_sm(int, int);
"""

# (EG, MT) -> EG*64 threads, EC = EG*MT envs per chunk. Must match FOR_EACH_CFG.
_CFGS = ((2, 12), (2, 16), (4, 4), (4, 8), (4, 12), (4, 16), (8, 4), (8, 8))

_MOD = None


def _mod():
    global _MOD
    if _MOD is None:
        os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
        _MOD = load_inline(
            name="grid_mingru_sps",
            cpp_sources=_CPP_SRC,
            cuda_sources=_CUDA_SRC,
            functions=["mega_launch", "policy_launch", "env_step_launch",
                       "mega_max_blocks_per_sm"],
            # No -use_fast_math: expf/tanhf must stay IEEE to match ATen bit-for-bit.
            extra_cuda_cflags=["-O3", "-lineinfo"]
            + (["-Xptxas", "-v"] if os.environ.get("KB_VERBOSE") else []),
            verbose=bool(os.environ.get("KB_VERBOSE")),
        )
    return _MOD


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.w_enc = nn.Parameter(torch.empty(HIDDEN, OBS_DIM))
        self.b_enc = nn.Parameter(torch.zeros(HIDDEN))
        self.w_gru = nn.Parameter(torch.empty(GRU_LAYERS, GRU_OUT, HIDDEN))
        self.w_a = nn.Parameter(torch.empty(NUM_ACTIONS, HIDDEN))
        self.b_a = nn.Parameter(torch.zeros(NUM_ACTIONS))
        self.w_v = nn.Parameter(torch.empty(1, HIDDEN))
        self.b_v = nn.Parameter(torch.zeros(1))
        self.reset_parameters(0)

    def reset_parameters(self, seed: int = 0) -> None:
        g = torch.Generator(device="cpu")
        g.manual_seed(seed)
        for p in self.parameters():
            tmp = torch.empty(p.shape, dtype=p.dtype, device="cpu")
            tmp.normal_(0.0, 0.02, generator=g)
            # p.copy_ (not p.data.copy_) so the version counter bumps and the
            # folded-constant cache in _prep() invalidates.
            with torch.no_grad():
                p.copy_(tmp)

    def forward(self, obs: torch.Tensor, state: torch.Tensor):
        return policy_forward(self, obs, state)


def _prep_key(model: Model):
    # data_ptr catches rebinding, _version catches in-place mutation (both
    # load_state_dict and the numeric-stress rescale go through Tensor.copy_,
    # which bumps the version counter).
    return tuple((p.data_ptr(), p._version) for p in model.parameters())


def _prep(model: Model):
    """Build the folded layer-0 constants (cached; rebuilding ~1.6MB of permuted
    weights every run() call costs several percent at the smallest shape).

    Cached on the model itself, so it cannot outlive the weights it was built from.
    """
    key = _prep_key(model)
    hit = getattr(model, "_kb_prep", None)
    if hit is not None and hit[0] == key:
        return hit[1]
    out = _prep_uncached(model)
    object.__setattr__(model, "_kb_prep", (key, out))
    return out


def _prep_uncached(model: Model):
    dev = torch.device("cuda:0")
    w_enc = model.w_enc.detach().to(dev, torch.float32)      # (256,4)
    b_enc = model.b_enc.detach().to(dev, torch.float32)      # (256,)
    w_gru = model.w_gru.detach().to(dev, torch.float32)      # (3,768,256)

    # gates0 = w_gru[0] @ (w_enc @ obs + b_enc) = (w_gru[0]@w_enc)@obs + w_gru[0]@b_enc
    w0 = w_gru[0].double()
    A = (w0 @ w_enc.double()).float()                        # (768,4)
    c = (w0 @ b_enc.double()).float()                        # (768,)

    return {
        "Afold": A.view(3, HIDDEN, 4).permute(2, 0, 1).contiguous(),      # (4,3,256)
        "cfold": c.view(3, HIDDEN).contiguous(),                          # (3,256)
        "Wenc": w_enc.t().contiguous(),                                   # (4,256)
        "benc": b_enc.contiguous(),                                       # (256,)
        # (2,768,256) -> (L,gate,j,k) -> (L,k,gate,j)
        "Wg": w_gru[1:].reshape(2, 3, HIDDEN, HIDDEN).permute(0, 3, 1, 2).contiguous(),
        "Wa": model.w_a.detach().to(dev, torch.float32).contiguous(),     # (4,256)
        "ba": model.b_a.detach().to(dev, torch.float32).contiguous(),     # (4,)
        "Wv": model.w_v.detach().to(dev, torch.float32).reshape(-1).contiguous(),
        "bv": model.b_v.detach().to(dev, torch.float32).contiguous(),
    }


_SM_COUNT = None
_MAXB = {}


def _cfg_plan(N: int, EG: int, MT: int):
    """Grid must be co-resident (per-step barrier), so blocks <= SMs*blocks_per_SM;
    leftover envs per block are handled as extra chunks."""
    global _SM_COUNT
    if _SM_COUNT is None:
        _SM_COUNT = torch.cuda.get_device_properties(0).multi_processor_count
    if (EG, MT) not in _MAXB:
        _MAXB[(EG, MT)] = _mod().mega_max_blocks_per_sm(EG, MT)
    if _MAXB[(EG, MT)] < 1:
        return None  # not launchable (registers/shared memory)
    # One block per SM. More would be co-resident (the barrier still works) but a
    # doubled-up SM takes twice as long per step and the barrier makes everyone
    # wait for it, so extra blocks only add imbalance.
    cap = _SM_COUNT
    EC = EG * MT
    blocks = min(cap, (N + EC - 1) // EC)
    E0 = (N + blocks - 1) // blocks
    chunks = (E0 + EC - 1) // EC
    E = chunks * EC                       # envs_per_block, a whole number of chunks
    blocks = min(cap, (N + E - 1) // E)
    return blocks, E, chunks


# Roofline constants for the config chooser (measured on RTX PRO 6000 Blackwell).
_FFMA_PER_SM_S = 128 * 2.37e9   # 123.2 FFMA/clk/SM measured, saturated at 8 warps/SM
_WBYTES = 1.60e6                # weights restreamed by every block on every chunk-step
_L2_BW = float(os.environ.get("KB_L2_BW", 4.0e12))


def _cost(H, blocks, E, chunks):
    """Seconds. The two terms pull opposite ways: a bigger chunk (EC) cuts weight
    traffic but coarsens E, and a smaller chunk fills more SMs but restreams the
    weights more often. They overlap, so time is the max, not the sum."""
    t_ffma = H * E * 768.0 * HIDDEN * 2 / _FFMA_PER_SM_S
    t_l2 = H * blocks * chunks * _WBYTES / _L2_BW
    return max(t_ffma, t_l2)


def _plan(N: int, H: int = 32):
    """Pick (EG, MT, blocks, envs_per_block)."""
    env = os.environ.get("KB_CFG")
    if env:
        EG, MT = (int(x) for x in env.split(","))
        blocks, E, _ = _cfg_plan(N, EG, MT)
        return EG, MT, blocks, E
    best = None
    for EG, MT in _CFGS:
        p = _cfg_plan(N, EG, MT)
        if p is None:
            continue
        blocks, E, chunks = p
        # Tie-break toward EG=4: 8 warps/SM is the geometry measured to saturate
        # FFMA (123.2/clk of 128); EG=2 is only 4 warps and EG=8 spills.
        key = (_cost(H, blocks, E, chunks), E, abs(EG - 4))
        if best is None or key < best[0]:
            best = (key, EG, MT, blocks, E)
    _, EG, MT, blocks, E = best
    return EG, MT, blocks, E


def policy_forward(model: Model, obs: torch.Tensor, state: torch.Tensor):
    """obs (N,4), state (N,L,H) -> logits (N,4), new_state (N,L,H), value (N,)."""
    m = _mod()
    dev = obs.device
    W = _prep(model)
    N = obs.shape[0]
    EG, MT = 4, 8
    EC = EG * MT
    Np = ((N + EC - 1) // EC) * EC

    obs_p = torch.zeros(Np, 4, device=dev, dtype=torch.float32)
    obs_p[:N] = obs.to(torch.float32)
    st = torch.zeros(GRU_LAYERS, Np, HIDDEN, device=dev, dtype=torch.float32)
    st[:, :N, :] = state.to(torch.float32).permute(1, 0, 2)
    logits = torch.zeros(Np, NUM_ACTIONS, device=dev, dtype=torch.float32)
    value = torch.zeros(Np, device=dev, dtype=torch.float32)

    m.policy_launch(W["Afold"], W["cfold"], W["Wenc"], W["benc"], W["Wg"], W["Wa"],
                    W["ba"], W["Wv"], W["bv"], obs_p, st, logits, value, N, Np, EG, MT)
    return logits[:N], st[:, :N, :].permute(1, 0, 2).contiguous(), value[:N]


def env_step(agent: torch.Tensor, food: torch.Tensor, actions: torch.Tensor,
             rng_state: torch.Tensor):
    m = _mod()
    dev = agent.device
    N = agent.shape[0]
    agent = agent.to(torch.float32).contiguous()
    food = food.to(torch.float32).contiguous()
    act = actions.to(torch.int64).contiguous()
    rng_in = rng_state.to(torch.int64).contiguous()

    agent_out = torch.empty_like(agent)
    food_out = torch.empty_like(food)
    rng_out = torch.empty_like(rng_in)
    rew = torch.empty(N, device=dev, dtype=torch.float32)
    hit = torch.empty(N, device=dev, dtype=torch.int32)
    any_ = torch.zeros(1, device=dev, dtype=torch.int32)

    m.env_step_launch(agent, food, act, agent_out, food_out, rng_in, rng_out, rew, hit,
                      any_, N)
    return agent_out, food_out, rew, rng_out


def run(num_envs: int, horizon: int, seed: int, model: Model | None = None) -> dict:
    dev = torch.device("cuda:0")
    if model is None:
        model = Model()
    model = model.to(dev).eval()
    m = _mod()
    W = _prep(model)

    EG, MT, blocks, E = _plan(num_envs, horizon)
    Np = blocks * E

    # identical init draw order to the reference
    g = torch.Generator(device="cpu")
    g.manual_seed(seed)
    agent = torch.randint(0, BOARD, (num_envs, 2), generator=g)
    food = torch.randint(0, BOARD, (num_envs, 2), generator=g)

    ax = torch.zeros(Np, device=dev, dtype=torch.int32)
    ay = torch.zeros(Np, device=dev, dtype=torch.int32)
    fx = torch.zeros(Np, device=dev, dtype=torch.int32)
    fy = torch.zeros(Np, device=dev, dtype=torch.int32)
    ax[:num_envs] = agent[:, 0].to(dev, torch.int32)
    ay[:num_envs] = agent[:, 1].to(dev, torch.int32)
    fx[:num_envs] = food[:, 0].to(dev, torch.int32)
    fy[:num_envs] = food[:, 1].to(dev, torch.int32)

    rng = torch.zeros(Np, device=dev, dtype=torch.int64)
    rng[:num_envs] = torch.arange(num_envs, device=dev, dtype=torch.int64) + (seed * 10007)
    rew = torch.zeros(Np, device=dev, dtype=torch.float32)
    state = torch.zeros(GRU_LAYERS, Np, HIDDEN, device=dev, dtype=torch.float32)
    logits = torch.zeros(Np, NUM_ACTIONS, device=dev, dtype=torch.float32)
    hit = torch.zeros(Np, device=dev, dtype=torch.int32)
    bar = torch.zeros(1, device=dev, dtype=torch.int32)
    anyhit = torch.zeros(max(horizon, 1), device=dev, dtype=torch.int32)

    m.mega_launch(W["Afold"], W["cfold"], W["Wenc"], W["benc"], W["Wg"], W["Wa"], W["ba"],
                  ax, ay, fx, fy, rng, rew, state, logits, hit, bar, anyhit,
                  num_envs, Np, horizon, E, blocks, EG, MT)

    return {
        "rewards": rew[:num_envs],
        "positions": torch.stack([ax[:num_envs], ay[:num_envs]], dim=-1).to(torch.int64),
        "last_logits": logits[:num_envs],
        "state": state[:, :num_envs, :].permute(1, 0, 2).contiguous(),
    }


def get_init_inputs():
    return []


def get_inputs():
    return []