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

The rollout runs as one megastep kernel per env step (env transition + encoder +
3 MinGRU layers + action head, all fused).  Design notes are in _CUDA_SRC below;
the short version:

  * the encoder is algebraically folded into layer 0, turning one of the three
    768x256 GEMMs into a 768x4 one (1.49x fewer FLOPs);
  * the two remaining GEMMs run on fp16 tensor cores using an fp32-accurate
    hi/lo split (3 mma terms), so logits stay within ~3e-8 of the reference;
  * 16 warps x 2 resident blocks per SM lets one block's MinGRU cell overlap the
    other's GEMM, which is what actually keeps the tensor pipe fed.

policy_forward/env_step are separate exact-fp32 kernels (they are correctness
surfaces, not throughput ones).
"""
from __future__ import annotations

import os
from pathlib import Path

import torch
import torch.nn as nn

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

_CUDA_SRC = r"""
// Grid-foraging env + 3x MinGRU(h=256) policy rollout -- fused CUDA megastep kernel.
//
// Key ideas
//  1. Layer-0 fold: gates_0 = W0 @ (Wenc @ obs + benc) = (W0 Wenc) @ obs + W0 benc.
//     Since obs is 4-dim this turns a 768x256 GEMM into a 768x4 one -> 1.49x fewer FLOPs.
//  2. The two remaining 768x256 GEMMs run on fp16 tensor cores with a hi/lo split
//     (3 mma terms) which reproduces fp32 accuracy to ~2e-7 relative.
//  3. Weights are pre-permuted into mma B-fragment order so the inner loop is one
//     perfectly coalesced 128-bit global load per (k-step, n-tile).
//  4. One kernel launch per env step; the previous step's food respawn is folded into
//     the head of the next launch so the global `hit.any()` reduction costs nothing
//     beyond the kernel boundary that is already there.

#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <vector>
#include <cstring>
#if defined(__x86_64__)
#include <immintrin.h>
#endif

#define BOARD 11
#define HID 256
#define NGATE 768
#define NACT 4
#define LCG_A 6364136223846793005ULL
#define LCG_MASK 0x7FFFFFFFFFFFFFFFLL

// hi/lo split scales: keeps the fp16 "lo" limbs inside the normal range.
#define ASCALE 1024.0f
#define BSCALE 1024.0f
#define INVSCALE (1.0f / (ASCALE * BSCALE))

#define RENV 32           // envs per block
#define NWARP 16          // warps per block (2 blocks/SM at 512 threads, 64 regs)
#define NTHREADS (NWARP * 32)
#define HSTRIDE 264      // halves per row of the shared h_hi / h_lo staging arrays
#define GPW 1             // 8-wide hidden groups handled per warp per pass
#define PASSES (32 / (NWARP * GPW))
#define NTPW (3 * GPW)   // n-tiles (of 8 gate columns) handled per warp
#define MTILE (RENV / 16)
#define BPSM 2            // resident blocks per SM

// --------------------------------------------------------------------------------------
// PTX helpers
// --------------------------------------------------------------------------------------
__device__ __forceinline__ void mma_m16n8k16(float *d, const uint32_t *a, const uint32_t *b) {
  asm volatile(
      "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
      "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
      : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
      : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));
}

__device__ __forceinline__ void ldmatrix_x4(uint32_t *r, const __half *p) {
  uint32_t addr = static_cast<uint32_t>(__cvta_generic_to_shared(p));
  asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"
               : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
               : "r"(addr));
}

__device__ __forceinline__ float sigmoidf(float x) { return 1.0f / (1.0f + expf(-x)); }

// Throughput variants for the rollout: 4 instructions per gate instead of ~12/~25.
// The caller folds the accumulator rescale into the exponent constant, so
// sigma(s*z) is  rcp(1 + ex2(z * c))  with c = -s*log2(e).  Absolute error is
// ~5e-8 (sigmoid) / ~1.2e-7 (tanh) -- far below the 1e-3 logit tolerance.
#define LOG2E 1.4426950408889634f
// .ftz is safe here: the exponent argument never lands in the subnormal range
// (|z*c| stays well inside +-150) and rcp's operand is always >= 1.
__device__ __forceinline__ float ex2a(float x) {
  float r;
  asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(x));
  return r;
}
__device__ __forceinline__ float rcpa(float x) {
  float r;
  asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(r) : "f"(x));
  return r;
}
__device__ __forceinline__ float sigc(float z, float c) { return rcpa(1.0f + ex2a(z * c)); }
__device__ __forceinline__ float tanhc(float z, float c) {
  return fmaf(2.0f, sigc(z, 2.0f * c), -1.0f);
}

__device__ __forceinline__ uint32_t pack2(__half a, __half b) {
  __half2 h = __halves2half2(a, b);
  return *reinterpret_cast<uint32_t *>(&h);
}

// --------------------------------------------------------------------------------------
// Weight preprocessing: (768,256) fp32 -> mma B-fragment order, fp16 hi/lo.
//
// Permuted gate column c maps to original row: c = 24*gh + 8*gate + jj  ->  gate*256+gh*8+jj
// so that the three gates of one 8-wide hidden group are three consecutive n-tiles.
// Output index: ((kstep*96 + ntile)*32 + lane), one uint4 = {hi01, hi89, lo01, lo89}.
// --------------------------------------------------------------------------------------
__global__ void permute_weights(const float *__restrict__ W, uint4 *__restrict__ out,
                                int total) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx >= total) return;
  const int layer = idx / (16 * 96 * 32);
  const int sub = idx - layer * (16 * 96 * 32);
  W += (size_t)layer * NGATE * HID;
  int lane = sub & 31;
  int ntile = (sub >> 5) % 96;
  int kstep = sub / (96 * 32);

  int c = ntile * 8 + (lane >> 2);
  int gh = c / 24, rem = c - gh * 24, gate = rem >> 3, jj = rem & 7;
  int grow = gate * HID + gh * 8 + jj;
  int k = kstep * 16 + ((lane & 3) << 1);

  const float *wp = W + (size_t)grow * HID;
  float v[4] = {wp[k] * BSCALE, wp[k + 1] * BSCALE, wp[k + 8] * BSCALE, wp[k + 9] * BSCALE};
  __half hi[4], lo[4];
#pragma unroll
  for (int i = 0; i < 4; i++) {
    hi[i] = __float2half_rn(v[i]);
    lo[i] = __float2half_rn(v[i] - __half2float(hi[i]));
  }
  uint4 o;
  o.x = pack2(hi[0], hi[1]);
  o.y = pack2(hi[2], hi[3]);
  o.z = pack2(lo[0], lo[1]);
  o.w = pack2(lo[2], lo[3]);
  out[idx] = o;
}

// Build the permuted (4,768) / (768,) layer-0 folded constants from natural order.
// Fold the encoder into layer 0's gate weights:  M0 = W0 @ Wenc,  c0 = W0 @ benc.
// fp64 accumulation keeps this below the fp32 round-off of the reference chain.
__global__ __launch_bounds__(256) void fold_l0(const float *__restrict__ wg0,   // (768,256)
                                               const float *__restrict__ wenc,  // (256,4)
                                               const float *__restrict__ benc,  // (256,)
                                               float *__restrict__ M0,          // (768,4)
                                               float *__restrict__ c0) {        // (768,)
  __shared__ double red[8][5];
  const int r = blockIdx.x, k = threadIdx.x, lane = k & 31, w = k >> 5;
  const double wv = (double)wg0[(size_t)r * HID + k];
  double v[5];
#pragma unroll
  for (int o = 0; o < 4; o++) v[o] = wv * (double)wenc[k * 4 + o];
  v[4] = wv * (double)benc[k];
#pragma unroll
  for (int off = 16; off; off >>= 1)
#pragma unroll
    for (int o = 0; o < 5; o++) v[o] += __shfl_down_sync(0xffffffff, v[o], off);
  if (lane == 0)
#pragma unroll
    for (int o = 0; o < 5; o++) red[w][o] = v[o];
  __syncthreads();
  if (k < 5) {
    double s = 0;
#pragma unroll
    for (int i = 0; i < 8; i++) s += red[i][k];
    if (k < 4) M0[r * 4 + k] = (float)s;
    else c0[r] = (float)s;
  }
}

// Layer-0 fragment table.
//   n-tile  0..95 : folded gate weights  gates0 = M0 @ obs + c0   (permuted columns)
//   n-tile 96..127: encoder weights      h0     = Wenc @ obs + benc
// K is padded to 16: rows 0..3 hold the obs weights, row 4 holds the bias (the A
// operand carries a constant 1 there), rows 5..15 are zero.
__global__ void permute_l0(const float *__restrict__ M0,     // (768,4)
                           const float *__restrict__ c0,     // (768,)
                           const float *__restrict__ wenc,   // (256,4)
                           const float *__restrict__ benc,   // (256,)
                           uint4 *__restrict__ out) {        // [128][32]
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx >= 128 * 32) return;
  int lane = idx & 31;
  int ntile = idx >> 5;
  int k0 = (lane & 3) << 1;

  const float *wrow;
  float bias;
  if (ntile < 96) {
    int c = ntile * 8 + (lane >> 2);
    int gh = c / 24, rem = c - gh * 24, gate = rem >> 3, jj = rem & 7;
    int grow = gate * HID + gh * 8 + jj;
    wrow = M0 + (size_t)grow * 4;
    bias = c0[grow];
  } else {
    int j = (ntile - 96) * 8 + (lane >> 2);
    wrow = wenc + (size_t)j * 4;
    bias = benc[j];
  }
  // k values covered by this lane: k0, k0+1, k0+8, k0+9
  float v[4];
  const int ks[4] = {k0, k0 + 1, k0 + 8, k0 + 9};
#pragma unroll
  for (int i = 0; i < 4; i++) {
    int k = ks[i];
    v[i] = (k < 4) ? wrow[k] * BSCALE : ((k == 4) ? bias * BSCALE : 0.0f);
  }
  __half hi[4], lo[4];
#pragma unroll
  for (int i = 0; i < 4; i++) {
    hi[i] = __float2half_rn(v[i]);
    lo[i] = __float2half_rn(v[i] - __half2float(hi[i]));
  }
  uint4 o;
  o.x = pack2(hi[0], hi[1]);
  o.y = pack2(hi[2], hi[3]);
  o.z = pack2(lo[0], lo[1]);
  o.w = pack2(lo[2], lo[3]);
  out[idx] = o;
}

// --------------------------------------------------------------------------------------
// Fused rollout step
// --------------------------------------------------------------------------------------
#define NT0 (NTPW + GPW)  // layer-0 n-tiles per warp: NTPW gate tiles + GPW encoder tiles

struct SmemPtrs {
  __half *hhi;    // [2][RENV*HSTRIDE] ping-pong
  __half *hlo;    // [2][RENV*HSTRIDE]
  float *hf;      // [RENV*HID] fragment-major fp32 h
  float *lred;    // [NACT*NWARP*RENV]
  uint32_t *obsf; // [RENV][4]: packed layer-0 A fragments (hi,lo) x (lane group 0,1)
};

__device__ __forceinline__ SmemPtrs carve(char *base) {
  SmemPtrs s;
  char *p = base;
  s.hhi = (__half *)p;  p += 2 * RENV * HSTRIDE * sizeof(__half);
  s.hlo = (__half *)p;  p += 2 * RENV * HSTRIDE * sizeof(__half);
  s.hf = (float *)p;    p += RENV * HID * sizeof(float);
  s.lred = (float *)p;  p += NACT * NWARP * RENV * sizeof(float);
  s.obsf = (uint32_t *)p; p += RENV * 4 * sizeof(uint32_t);
  return s;
}

static size_t smem_bytes() {
  return 4 * RENV * HSTRIDE * sizeof(__half) + RENV * HID * sizeof(float) +
         NACT * NWARP * RENV * sizeof(float) + RENV * 4 * sizeof(uint32_t);
}

// LCG used by the reference env for food respawns.
__device__ __forceinline__ long long lcg(long long r) {
  return (long long)((unsigned long long)r * LCG_A + 1ULL) & LCG_MASK;
}

__global__ __launch_bounds__(NTHREADS, BPSM) void rollout_step(
    const uint4 *__restrict__ wfrag,   // [2][16][96][32]
    const uint4 *__restrict__ wfrag0,  // [128][32]
    const float *__restrict__ wa_g,    // (4,256)
    const float *__restrict__ ba_g,    // (4,)
    float *__restrict__ state,         // (N,3,256)
    int *__restrict__ agent,           // (N,2)
    int *__restrict__ food,            // (N,2)
    long long *__restrict__ rng,       // (N,)
    float *__restrict__ rewards,       // (N,)
    float *__restrict__ last_logits,   // (N,4)
    int *__restrict__ anyhit,          // (horizon,)
    int N, int step) {
  extern __shared__ char smem_raw[];
  SmemPtrs sm = carve(smem_raw);

  const int tid = threadIdx.x;
  const int warp = tid >> 5;
  const int lane = tid & 31;
  const int e0 = blockIdx.x * RENV;
  const int prev_any = (step > 0) ? anyhit[step - 1] : 0;

  // ---------------------------------------------------------------------------
  // Layer-0 A operand: [obs0 obs1 obs2 obs3 1 0 ... 0] * ASCALE, built straight
  // into mma fragment registers -- lane group (lane&3) selects which two of the
  // 16 padded columns this thread holds, so lanes 2,3 need no env data at all.
  // ---------------------------------------------------------------------------
  if (tid < RENV) {
    const int ge = e0 + tid;
    const int valid = (ge < N);
    int ax = 0, ay = 0, fx = 1, fy = 1;
    if (valid) {
      ax = agent[ge * 2 + 0];
      ay = agent[ge * 2 + 1];
      fx = food[ge * 2 + 0];
      fy = food[ge * 2 + 1];
      if (prev_any && ax == fx && ay == fy) {
        long long r = lcg(rng[ge]);
        fx = (int)(r % BOARD);
        r = lcg(r);
        fy = (int)(r % BOARD);
      }
    }
    const float o[4] = {(float)(fx - ax) / (float)BOARD * ASCALE,
                        (float)(fy - ay) / (float)BOARD * ASCALE,
                        (float)ax / (float)(BOARD - 1) * ASCALE,
                        (float)ay / (float)(BOARD - 1) * ASCALE};
#pragma unroll
    for (int g = 0; g < 2; g++) {
      __half h0 = __float2half_rn(o[2 * g]), h1 = __float2half_rn(o[2 * g + 1]);
      sm.obsf[tid * 4 + g * 2 + 0] = pack2(h0, h1);
      sm.obsf[tid * 4 + g * 2 + 1] = pack2(__float2half_rn(o[2 * g] - __half2float(h0)),
                                           __float2half_rn(o[2 * g + 1] - __half2float(h1)));
    }
  }
  __syncthreads();

  const int cgrp = lane & 3;
  const uint32_t abias = pack2(__float2half_rn(ASCALE), __float2half_rn(0.0f));
  uint32_t a0hi[MTILE][4], a0lo[MTILE][4];
#pragma unroll
  for (int mt = 0; mt < MTILE; mt++) {
#pragma unroll
    for (int half = 0; half < 2; half++) {
      uint32_t hi = 0, lo = 0;
      if (cgrp < 2) {
        const int e = mt * 16 + (lane >> 2) + half * 8;
        hi = sm.obsf[e * 4 + cgrp * 2 + 0];
        lo = sm.obsf[e * 4 + cgrp * 2 + 1];
      } else if (cgrp == 2) {
        hi = abias;  // constant column feeding the folded bias row
      }
      a0hi[mt][half] = hi;
      a0lo[mt][half] = lo;
    }
    a0hi[mt][2] = 0; a0hi[mt][3] = 0;
    a0lo[mt][2] = 0; a0lo[mt][3] = 0;
  }

  // ---------------------------------------------------------------------------
  // Three MinGRU layers.  L=0 is the folded encoder+gate rank-4 product (one
  // k-step, NT0 n-tiles); L=1,2 are the full 768x256 GEMMs (16 k-steps).
  // ---------------------------------------------------------------------------
  const int arow = ((lane >> 3) & 1) * 8 + (lane & 7);
  const int acol = ((lane >> 4) & 1) * 8;
  const float CE = -LOG2E * INVSCALE;

#pragma unroll
  for (int L = 0; L < 3; L++) {
#pragma unroll
   for (int ps = 0; ps < PASSES; ps++) {
    const int gbase = (ps * NWARP + warp) * GPW;
    float acc[NT0][MTILE][4];
#pragma unroll
    for (int n = 0; n < NT0; n++)
#pragma unroll
      for (int m = 0; m < MTILE; m++)
#pragma unroll
        for (int k = 0; k < 4; k++) acc[n][m][k] = 0.0f;

    if (L == 0) {
      const uint4 *wp = wfrag0 + (size_t)(gbase * 3) * 32 + lane;
#pragma unroll
      for (int nt = 0; nt < NTPW; nt++) {
        uint4 b = wp[nt * 32];
        uint32_t bhi[2] = {b.x, b.y}, blo[2] = {b.z, b.w};
#pragma unroll
        for (int mt = 0; mt < MTILE; mt++) {
          mma_m16n8k16(acc[nt][mt], a0hi[mt], bhi);
          mma_m16n8k16(acc[nt][mt], a0lo[mt], bhi);
          mma_m16n8k16(acc[nt][mt], a0hi[mt], blo);
        }
      }
      const uint4 *we = wfrag0 + (size_t)(96 + gbase) * 32 + lane;
#pragma unroll
      for (int gg = 0; gg < GPW; gg++) {
        uint4 b = we[gg * 32];
        uint32_t bhi[2] = {b.x, b.y}, blo[2] = {b.z, b.w};
#pragma unroll
        for (int mt = 0; mt < MTILE; mt++) {
          mma_m16n8k16(acc[NTPW + gg][mt], a0hi[mt], bhi);
          mma_m16n8k16(acc[NTPW + gg][mt], a0lo[mt], bhi);
          mma_m16n8k16(acc[NTPW + gg][mt], a0hi[mt], blo);
        }
      }
    } else {
      const uint4 *wl = wfrag + (size_t)(L - 1) * (16 * 96 * 32);
      const __half *hhi_in = sm.hhi + (size_t)(L - 1) * (RENV * HSTRIDE);
      const __half *hlo_in = sm.hlo + (size_t)(L - 1) * (RENV * HSTRIDE);
      const uint4 *wp0 = wl + (size_t)(gbase * 3) * 32 + lane;
      // software pipeline: the B fragment for k-step ks+1 is issued as soon as
      // the one for ks has been consumed.
      uint4 bcur[NTPW];
#pragma unroll
      for (int nt = 0; nt < NTPW; nt++) bcur[nt] = wp0[nt * 32];
      for (int ks = 0; ks < 16; ks++) {
        uint32_t ahi[MTILE][4], alo[MTILE][4];
#pragma unroll
        for (int mt = 0; mt < MTILE; mt++) {
          ldmatrix_x4(ahi[mt], hhi_in + (mt * 16 + arow) * HSTRIDE + ks * 16 + acol);
          ldmatrix_x4(alo[mt], hlo_in + (mt * 16 + arow) * HSTRIDE + ks * 16 + acol);
        }
        const uint4 *wn = wp0 + (size_t)(ks < 15 ? ks + 1 : 15) * (96 * 32);
        uint4 bv[NTPW];
#pragma unroll
        for (int nt = 0; nt < NTPW; nt++) {
          bv[nt] = bcur[nt];
          bcur[nt] = wn[nt * 32];
        }
        // Issue the three emulation terms as three separate sweeps so that
        // successive mma on the same accumulator are NTPW*MTILE apart.
#pragma unroll
        for (int nt = 0; nt < NTPW; nt++) {
          uint32_t bhi[2] = {bv[nt].x, bv[nt].y};
#pragma unroll
          for (int mt = 0; mt < MTILE; mt++) mma_m16n8k16(acc[nt][mt], ahi[mt], bhi);
        }
#pragma unroll
        for (int nt = 0; nt < NTPW; nt++) {
          uint32_t bhi[2] = {bv[nt].x, bv[nt].y};
#pragma unroll
          for (int mt = 0; mt < MTILE; mt++) mma_m16n8k16(acc[nt][mt], alo[mt], bhi);
        }
#pragma unroll
        for (int nt = 0; nt < NTPW; nt++) {
          uint32_t blo[2] = {bv[nt].z, bv[nt].w};
#pragma unroll
          for (int mt = 0; mt < MTILE; mt++) mma_m16n8k16(acc[nt][mt], ahi[mt], blo);
        }
      }
    }

    // ---- MinGRU cell + highway ----
    __half *hhi_out = sm.hhi + (size_t)(L & 1) * (RENV * HSTRIDE);
    __half *hlo_out = sm.hlo + (size_t)(L & 1) * (RENV * HSTRIDE);
#pragma unroll
    for (int gg = 0; gg < GPW; gg++) {
      const int jc = (gbase + gg) * 8 + ((lane & 3) << 1);
#pragma unroll
      for (int mt = 0; mt < MTILE; mt++) {
#pragma unroll
        for (int half = 0; half < 2; half++) {
          const int e = mt * 16 + (lane >> 2) + half * 8;
          const int ge = e0 + e;
          const int s0 = half * 2;
          const int hidx = ((gbase + gg) * MTILE + mt) * 128 + s0 * 32 + lane;

          float th0 = tanhc(acc[gg * 3 + 0][mt][s0], CE);
          float th1 = tanhc(acc[gg * 3 + 0][mt][s0 + 1], CE);
          float sg0 = sigc(acc[gg * 3 + 1][mt][s0], CE);
          float sg1 = sigc(acc[gg * 3 + 1][mt][s0 + 1], CE);
          float p0 = sigc(acc[gg * 3 + 2][mt][s0], CE);
          float p1 = sigc(acc[gg * 3 + 2][mt][s0 + 1], CE);

          float hp0, hp1;
          if (L == 0) {
            hp0 = acc[NTPW + gg][mt][s0] * INVSCALE;
            hp1 = acc[NTPW + gg][mt][s0 + 1] * INVSCALE;
          } else {
            hp0 = sm.hf[hidx];
            hp1 = sm.hf[hidx + 32];
          }

          float *sp = state + ((size_t)ge * 3 + L) * HID + jc;
          float2 st2 = make_float2(0.f, 0.f);
          if (ge < N) st2 = *reinterpret_cast<float2 *>(sp);

          float o0 = st2.x + sg0 * (th0 - st2.x);
          float o1 = st2.y + sg1 * (th1 - st2.y);
          float hn0 = fmaf(p0, o0 - hp0, hp0);
          float hn1 = fmaf(p1, o1 - hp1, hp1);

          if (ge < N) *reinterpret_cast<float2 *>(sp) = make_float2(o0, o1);
          sm.hf[hidx] = hn0;
          sm.hf[hidx + 32] = hn1;
          if (L < 2) {
            __half2 hi =
                __halves2half2(__float2half_rn(hn0 * ASCALE), __float2half_rn(hn1 * ASCALE));
            __half2 lo = __halves2half2(
                __float2half_rn(hn0 * ASCALE - __half2float(__low2half(hi))),
                __float2half_rn(hn1 * ASCALE - __half2float(__high2half(hi))));
            *reinterpret_cast<__half2 *>(hhi_out + e * HSTRIDE + jc) = hi;
            *reinterpret_cast<__half2 *>(hlo_out + e * HSTRIDE + jc) = lo;
          }
        }
      }
    }
   }
    if (L < 2) __syncthreads();
  }

  // ---- logits = wa h3 + ba ----
  {
    float part[MTILE][2][NACT];
#pragma unroll
    for (int mt = 0; mt < MTILE; mt++)
#pragma unroll
      for (int half = 0; half < 2; half++)
#pragma unroll
        for (int a = 0; a < NACT; a++) part[mt][half][a] = 0.0f;

#pragma unroll
    for (int pg = 0; pg < PASSES * GPW; pg++) {
      const int grp = (pg / GPW) * NWARP * GPW + warp * GPW + (pg % GPW);
      const int jc = grp * 8 + ((lane & 3) << 1);
      float wa0[NACT], wa1[NACT];
#pragma unroll
      for (int a = 0; a < NACT; a++) {
        wa0[a] = wa_g[a * HID + jc];
        wa1[a] = wa_g[a * HID + jc + 1];
      }
#pragma unroll
      for (int mt = 0; mt < MTILE; mt++)
#pragma unroll
        for (int half = 0; half < 2; half++) {
          const int hidx = (grp * MTILE + mt) * 128 + half * 64 + lane;
          float h0 = sm.hf[hidx];
          float h1 = sm.hf[hidx + 32];
#pragma unroll
          for (int a = 0; a < NACT; a++) part[mt][half][a] += wa0[a] * h0 + wa1[a] * h1;
        }
    }
#pragma unroll
    for (int mt = 0; mt < MTILE; mt++)
#pragma unroll
      for (int half = 0; half < 2; half++)
#pragma unroll
        for (int a = 0; a < NACT; a++) {
          float v = part[mt][half][a];
          v += __shfl_xor_sync(0xffffffff, v, 1);
          v += __shfl_xor_sync(0xffffffff, v, 2);
          if ((lane & 3) == 0) {
            const int e = mt * 16 + (lane >> 2) + half * 8;
            sm.lred[(a * NWARP + warp) * RENV + e] = v;
          }
        }
  }
  __syncthreads();

  // ---- action, env transition, reward ----
  if (tid < RENV) {
    const int e = tid;
    const int ge = e0 + e;
    const int valid = (ge < N);
    float lg[NACT];
#pragma unroll
    for (int a = 0; a < NACT; a++) {
      float v = ba_g[a];
      for (int w = 0; w < NWARP; w++) v += sm.lred[(a * NWARP + w) * RENV + e];
      lg[a] = v;
    }
    int act = 0;
    float best = lg[0];
#pragma unroll
    for (int a = 1; a < NACT; a++)
      if (lg[a] > best) { best = lg[a]; act = a; }

    int ax = 0, ay = 0, fx = 1, fy = 1;
    long long r = 0;
    if (valid) {
      ax = agent[ge * 2 + 0];
      ay = agent[ge * 2 + 1];
      fx = food[ge * 2 + 0];
      fy = food[ge * 2 + 1];
      r = rng[ge];
      if (prev_any) {
        long long r1 = lcg(r);
        long long r2 = lcg(r1);
        if (ax == fx && ay == fy) {
          fx = (int)(r1 % BOARD);
          fy = (int)(r2 % BOARD);
        }
        r = r2;
      }
    }
    if (act == 0) ay -= 1;
    else if (act == 1) ay += 1;
    else if (act == 2) ax -= 1;
    else ax += 1;
    ax = min(max(ax, 0), BOARD - 1);
    ay = min(max(ay, 0), BOARD - 1);
    int hit = (valid && ax == fx && ay == fy) ? 1 : 0;

    if (valid) {
      agent[ge * 2 + 0] = ax;
      agent[ge * 2 + 1] = ay;
      food[ge * 2 + 0] = fx;
      food[ge * 2 + 1] = fy;
      rng[ge] = r;
      rewards[ge] += (float)hit;
#pragma unroll
      for (int a = 0; a < NACT; a++) last_logits[ge * NACT + a] = lg[a];
    }
    // Every block races to set the same flag; a cheap pre-read keeps all but the
    // first few from serialising on the L2 atomic (2048 blocks at the largest shape).
    unsigned mask = __ballot_sync(0xffffffff, hit);
    if (lane == 0 && mask && !__ldcv(&anyhit[step])) atomicOr(&anyhit[step], 1);
  }
}

// --------------------------------------------------------------------------------------
// Exact fp32 reference-shaped policy forward (used by the public policy_forward API).
// One block handles PF_E envs; 256 threads = one hidden unit each.
// --------------------------------------------------------------------------------------
#define PF_E 8
__global__ __launch_bounds__(256) void policy_fwd_kernel(
    const float *__restrict__ obs,     // (N,4)
    const float *__restrict__ state,   // (N,3,256)
    const float *__restrict__ wenc,    // (256,4)
    const float *__restrict__ benc,    // (256,)
    const float *__restrict__ wgruT,   // (3,256,768)  transposed: [layer][k][gaterow]
    const float *__restrict__ wa,      // (4,256)
    const float *__restrict__ ba,      // (4,)
    const float *__restrict__ wv,      // (1,256)
    const float *__restrict__ bv,      // (1,)
    float *__restrict__ logits,        // (N,4)
    float *__restrict__ nstate,        // (N,3,256)
    float *__restrict__ value,         // (N,)
    int N) {
  __shared__ float sh[PF_E][HID];
  __shared__ float sred[PF_E][NACT + 1][8];
  const int j = threadIdx.x;
  const int eb = blockIdx.x * PF_E;

  float h[PF_E];
#pragma unroll
  for (int u = 0; u < PF_E; u++) {
    int e = eb + u;
    float acc = benc[j];
    if (e < N) {
#pragma unroll
      for (int i = 0; i < 4; i++) acc += wenc[j * 4 + i] * obs[e * 4 + i];
    }
    h[u] = acc;
    sh[u][j] = acc;
  }
  __syncthreads();

  for (int L = 0; L < 3; L++) {
    float zh[PF_E], zg[PF_E], zp[PF_E];
#pragma unroll
    for (int u = 0; u < PF_E; u++) { zh[u] = 0.f; zg[u] = 0.f; zp[u] = 0.f; }
    const float *wp = wgruT + (size_t)L * HID * NGATE;
    for (int k = 0; k < HID; k++) {
      float a0 = wp[(size_t)k * NGATE + j];
      float a1 = wp[(size_t)k * NGATE + HID + j];
      float a2 = wp[(size_t)k * NGATE + 2 * HID + j];
#pragma unroll
      for (int u = 0; u < PF_E; u++) {
        float hv = sh[u][k];
        zh[u] += a0 * hv;
        zg[u] += a1 * hv;
        zp[u] += a2 * hv;
      }
    }
    __syncthreads();
#pragma unroll
    for (int u = 0; u < PF_E; u++) {
      int e = eb + u;
      float st = (e < N) ? state[((size_t)e * 3 + L) * HID + j] : 0.f;
      float o = st + sigmoidf(zg[u]) * (tanhf(zh[u]) - st);
      float p = sigmoidf(zp[u]);
      h[u] = p * o + (1.0f - p) * h[u];
      sh[u][j] = h[u];
      if (e < N) nstate[((size_t)e * 3 + L) * HID + j] = o;
    }
    __syncthreads();
  }

  // heads
  {
    float acc[PF_E][NACT + 1];
#pragma unroll
    for (int u = 0; u < PF_E; u++)
#pragma unroll
      for (int a = 0; a < NACT + 1; a++) acc[u][a] = 0.f;
#pragma unroll
    for (int u = 0; u < PF_E; u++) {
#pragma unroll
      for (int a = 0; a < NACT; a++) acc[u][a] = wa[a * HID + j] * h[u];
      acc[u][NACT] = wv[j] * h[u];
    }
    const int lane = j & 31, w = j >> 5;
#pragma unroll
    for (int u = 0; u < PF_E; u++)
#pragma unroll
      for (int a = 0; a < NACT + 1; a++) {
        float v = acc[u][a];
#pragma unroll
        for (int off = 16; off; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off);
        if (lane == 0) sred[u][a][w] = v;
      }
    __syncthreads();
    if (j < PF_E * (NACT + 1)) {
      int u = j / (NACT + 1), a = j % (NACT + 1);
      float v = 0.f;
#pragma unroll
      for (int i = 0; i < 8; i++) v += sred[u][a][i];
      int e = eb + u;
      if (e < N) {
        if (a < NACT) logits[e * NACT + a] = v + ba[a];
        else value[e] = v + bv[0];
      }
    }
  }
}

// --------------------------------------------------------------------------------------
// Public env_step (two kernels: move+hit, then respawn)
// --------------------------------------------------------------------------------------
__global__ void env_move(const float *__restrict__ agent, const float *__restrict__ food,
                         const long long *__restrict__ actions, float *__restrict__ nagent,
                         float *__restrict__ reward, int *__restrict__ anyflag, int N) {
  int e = blockIdx.x * blockDim.x + threadIdx.x;
  int hit = 0;
  if (e < N) {
    float ax = agent[e * 2 + 0], ay = agent[e * 2 + 1];
    long long a = actions[e];
    if (a == 0) ay -= 1.f;
    else if (a == 1) ay += 1.f;
    else if (a == 2) ax -= 1.f;
    else if (a == 3) ax += 1.f;
    ax = fminf(fmaxf(ax, 0.f), (float)(BOARD - 1));
    ay = fminf(fmaxf(ay, 0.f), (float)(BOARD - 1));
    nagent[e * 2 + 0] = ax; nagent[e * 2 + 1] = ay;
    hit = (ax == food[e * 2 + 0] && ay == food[e * 2 + 1]) ? 1 : 0;
    reward[e] = (float)hit;
  }
  unsigned m = __ballot_sync(0xffffffff, hit);
  if ((threadIdx.x & 31) == 0 && m) atomicOr(anyflag, 1);
}

__global__ void env_respawn(const float *__restrict__ nagent, float *__restrict__ food,
                            long long *__restrict__ rng, const int *__restrict__ anyflag, int N) {
  int e = blockIdx.x * blockDim.x + threadIdx.x;
  if (e >= N) return;
  if (!*anyflag) return;
  long long r = rng[e];
  r = (long long)((unsigned long long)r * LCG_A + 1ULL) & LCG_MASK;
  float fx = (float)(r % BOARD);
  r = (long long)((unsigned long long)r * LCG_A + 1ULL) & LCG_MASK;
  float fy = (float)(r % BOARD);
  rng[e] = r;
  if (nagent[e * 2 + 0] == food[e * 2 + 0] && nagent[e * 2 + 1] == food[e * 2 + 1]) {
    food[e * 2 + 0] = fx; food[e * 2 + 1] = fy;
  }
}

// --------------------------------------------------------------------------------------
// host entry points
// --------------------------------------------------------------------------------------
torch::Tensor prep_wfrag(torch::Tensor w) {  // w: (L,768,256) contiguous fp32 cuda
  int nl = (int)w.size(0);
  auto out = torch::empty({nl * 16 * 96 * 32 * 4}, w.options().dtype(torch::kInt32));
  int total = nl * 16 * 96 * 32;
  permute_weights<<<(total + 255) / 256, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
      w.data_ptr<float>(), reinterpret_cast<uint4 *>(out.data_ptr<int>()), total);
  return out;
}

torch::Tensor prep_l0(torch::Tensor wg0, torch::Tensor wenc, torch::Tensor benc) {
  auto stream = at::cuda::getCurrentCUDAStream();
  auto M0 = torch::empty({NGATE, 4}, wg0.options());
  auto c0 = torch::empty({NGATE}, wg0.options());
  fold_l0<<<NGATE, 256, 0, stream>>>(wg0.data_ptr<float>(), wenc.data_ptr<float>(),
                                     benc.data_ptr<float>(), M0.data_ptr<float>(),
                                     c0.data_ptr<float>());
  auto out = torch::empty({128 * 32 * 4}, wg0.options().dtype(torch::kInt32));
  permute_l0<<<(128 * 32 + 255) / 256, 256, 0, stream>>>(
      M0.data_ptr<float>(), c0.data_ptr<float>(), wenc.data_ptr<float>(),
      benc.data_ptr<float>(), reinterpret_cast<uint4 *>(out.data_ptr<int>()));
  return out;
}

void rollout(torch::Tensor wfrag, torch::Tensor wfrag0, torch::Tensor wa, torch::Tensor ba,
             torch::Tensor state, torch::Tensor agent, torch::Tensor food, torch::Tensor rng,
             torch::Tensor rewards, torch::Tensor last_logits, torch::Tensor anyhit, int64_t N,
             int64_t horizon) {
  int nblk = (int)((N + RENV - 1) / RENV);
  size_t sb = smem_bytes();
  static bool attr_set = false;
  if (!attr_set) {
    cudaFuncSetAttribute(rollout_step, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)sb);
    cudaFuncSetAttribute(rollout_step, cudaFuncAttributePreferredSharedMemoryCarveout, 100);
    attr_set = true;
  }
  auto stream = at::cuda::getCurrentCUDAStream();
  for (int t = 0; t < horizon; t++) {
    rollout_step<<<nblk, NTHREADS, sb, stream>>>(
        reinterpret_cast<const uint4 *>(wfrag.data_ptr<int>()),
        reinterpret_cast<const uint4 *>(wfrag0.data_ptr<int>()), wa.data_ptr<float>(),
        ba.data_ptr<float>(), state.data_ptr<float>(), agent.data_ptr<int>(),
        food.data_ptr<int>(), reinterpret_cast<long long *>(rng.data_ptr<int64_t>()),
        rewards.data_ptr<float>(), last_logits.data_ptr<float>(), anyhit.data_ptr<int>(),
        (int)N, t);
  }
}

std::vector<torch::Tensor> policy_forward_cuda(torch::Tensor obs, torch::Tensor state,
                                               torch::Tensor wenc, torch::Tensor benc,
                                               torch::Tensor wgruT, torch::Tensor wa,
                                               torch::Tensor ba, torch::Tensor wv,
                                               torch::Tensor bv) {
  int N = obs.size(0);
  auto logits = torch::empty({N, NACT}, obs.options());
  auto nstate = torch::empty({N, 3, HID}, obs.options());
  auto value = torch::empty({N}, obs.options());
  int nblk = (N + PF_E - 1) / PF_E;
  policy_fwd_kernel<<<nblk, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
      obs.data_ptr<float>(), state.data_ptr<float>(), wenc.data_ptr<float>(),
      benc.data_ptr<float>(), wgruT.data_ptr<float>(), wa.data_ptr<float>(), ba.data_ptr<float>(),
      wv.data_ptr<float>(), bv.data_ptr<float>(), logits.data_ptr<float>(),
      nstate.data_ptr<float>(), value.data_ptr<float>(), N);
  return {logits, nstate, value};
}

std::vector<torch::Tensor> env_step_cuda(torch::Tensor agent, torch::Tensor food,
                                         torch::Tensor actions, torch::Tensor rng) {
  int N = agent.size(0);
  auto nagent = torch::empty_like(agent);
  auto nfood = food.clone();
  auto nrng = rng.clone();
  auto reward = torch::empty({N}, agent.options());
  auto flag = torch::zeros({1}, agent.options().dtype(torch::kInt32));
  auto stream = at::cuda::getCurrentCUDAStream();
  int nb = (N + 255) / 256;
  env_move<<<nb, 256, 0, stream>>>(agent.data_ptr<float>(), food.data_ptr<float>(),
                                   reinterpret_cast<const long long *>(actions.data_ptr<int64_t>()), nagent.data_ptr<float>(),
                                   reward.data_ptr<float>(), flag.data_ptr<int>(), N);
  env_respawn<<<nb, 256, 0, stream>>>(nagent.data_ptr<float>(), nfood.data_ptr<float>(),
                                      reinterpret_cast<long long *>(nrng.data_ptr<int64_t>()), flag.data_ptr<int>(), N);
  return {nagent, nfood, reward, nrng};
}


// --------------------------------------------------------------------------------------
// Host-side mt19937 matching at::CPUGeneratorImpl (init_genrand seeding, 32-bit draws).
// torch.randint(0, 11) is `engine() % 11`, so reproducing the raw stream reproduces the
// tensor bit-for-bit; doing it here (vectorised block twist) is ~4x faster than the
// generic per-element ATen path and this sits directly in run()'s critical path.
// --------------------------------------------------------------------------------------
namespace mtgen {
constexpr int N_ = 624, M_ = 397;
constexpr uint32_t MATRIX_A = 0x9908b0dfu, UMASK = 0x80000000u, LMASK = 0x7fffffffu;

static inline uint32_t twist1(uint32_t u, uint32_t v) {
  return (((u & UMASK) | (v & LMASK)) >> 1) ^ ((v & 1u) ? MATRIX_A : 0u);
}
static inline uint32_t temper1(uint32_t y) {
  y ^= (y >> 11);
  y ^= (y << 7) & 0x9d2c5680u;
  y ^= (y << 15) & 0xefc60000u;
  y ^= (y >> 18);
  return y;
}

static void next_state_scalar(uint32_t *st) {
  uint32_t *p = st;
  for (int j = N_ - M_ + 1; --j; p++) *p = p[M_] ^ twist1(p[0], p[1]);
  for (int j = M_; --j; p++) *p = p[M_ - N_] ^ twist1(p[0], p[1]);
  *p = p[M_ - N_] ^ twist1(p[0], st[0]);
}

#if defined(__x86_64__)
__attribute__((target("avx2"))) static inline __m256i twist8(__m256i u, __m256i v) {
  const __m256i um = _mm256_set1_epi32((int)UMASK), lm = _mm256_set1_epi32((int)LMASK);
  __m256i mixed = _mm256_or_si256(_mm256_and_si256(u, um), _mm256_and_si256(v, lm));
  __m256i odd = _mm256_and_si256(v, _mm256_set1_epi32(1));
  __m256i mask = _mm256_sub_epi32(_mm256_setzero_si256(), odd);
  return _mm256_xor_si256(_mm256_srli_epi32(mixed, 1),
                          _mm256_and_si256(mask, _mm256_set1_epi32((int)MATRIX_A)));
}
__attribute__((target("avx2"))) static void next_state_avx2(uint32_t *st) {
  int i = 0;
  for (; i + 8 <= N_ - M_; i += 8) {  // reads are all pre-twist here
    __m256i u = _mm256_loadu_si256((const __m256i *)(st + i));
    __m256i v = _mm256_loadu_si256((const __m256i *)(st + i + 1));
    __m256i w = _mm256_loadu_si256((const __m256i *)(st + i + M_));
    _mm256_storeu_si256((__m256i *)(st + i), _mm256_xor_si256(w, twist8(u, v)));
  }
  for (; i < N_ - M_; i++) st[i] = st[i + M_] ^ twist1(st[i], st[i + 1]);
  for (; i + 8 <= N_ - 1; i += 8) {  // st[i-227] is already updated, distance 227 >= 8
    __m256i u = _mm256_loadu_si256((const __m256i *)(st + i));
    __m256i v = _mm256_loadu_si256((const __m256i *)(st + i + 1));
    __m256i w = _mm256_loadu_si256((const __m256i *)(st + i + M_ - N_));
    _mm256_storeu_si256((__m256i *)(st + i), _mm256_xor_si256(w, twist8(u, v)));
  }
  for (; i < N_ - 1; i++) st[i] = st[i + M_ - N_] ^ twist1(st[i], st[i + 1]);
  st[N_ - 1] = st[M_ - 1] ^ twist1(st[N_ - 1], st[0]);
}
__attribute__((target("avx2"))) static void temper_mod_avx2(const uint32_t *st, int32_t *out,
                                                            int cnt, uint32_t range) {
  int i = 0;
  const __m256i c7 = _mm256_set1_epi32((int)0x9d2c5680u), c15 = _mm256_set1_epi32((int)0xefc60000u);
  for (; i + 8 <= cnt; i += 8) {
    __m256i y = _mm256_loadu_si256((const __m256i *)(st + i));
    y = _mm256_xor_si256(y, _mm256_srli_epi32(y, 11));
    y = _mm256_xor_si256(y, _mm256_and_si256(_mm256_slli_epi32(y, 7), c7));
    y = _mm256_xor_si256(y, _mm256_and_si256(_mm256_slli_epi32(y, 15), c15));
    y = _mm256_xor_si256(y, _mm256_srli_epi32(y, 18));
    // y % 11 via the exact magic multiply  q = (y * 3123612579) >> 35
    const __m256i mg = _mm256_set1_epi64x(3123612579LL);
    __m256i qe = _mm256_srli_epi64(_mm256_mul_epu32(y, mg), 35);
    __m256i qo = _mm256_srli_epi64(_mm256_mul_epu32(_mm256_srli_epi64(y, 32), mg), 35);
    __m256i q = _mm256_blend_epi32(qe, _mm256_slli_epi64(qo, 32), 0xAA);
    __m256i r = _mm256_sub_epi32(y, _mm256_mullo_epi32(q, _mm256_set1_epi32((int)range)));
    _mm256_storeu_si256((__m256i *)(out + i), r);
  }
  for (; i < cnt; i++) out[i] = (int32_t)(temper1(st[i]) % range);
}
static bool have_avx2() {
  static int c = -1;
  if (c < 0) c = __builtin_cpu_supports("avx2") ? 1 : 0;
  return c == 1;
}
#endif

static void randint_fill(int32_t *out, int64_t n, uint64_t seed, uint32_t range) {
  uint32_t st[N_];
  st[0] = (uint32_t)(seed & 0xffffffffu);
  for (int j = 1; j < N_; j++)
    st[j] = 1812433253u * (st[j - 1] ^ (st[j - 1] >> 30)) + (uint32_t)j;
  int64_t done = 0;
  while (done < n) {
#if defined(__x86_64__)
    if (have_avx2()) next_state_avx2(st); else next_state_scalar(st);
#else
    next_state_scalar(st);
#endif
    int cnt = (int)std::min<int64_t>(N_, n - done);
#if defined(__x86_64__)
    if (have_avx2() && range == 11) temper_mod_avx2(st, out + done, cnt, range);
    else
#endif
      for (int i = 0; i < cnt; i++) out[done + i] = (int32_t)(temper1(st[i]) % range);
    done += cnt;
  }
}
}  // namespace mtgen

torch::Tensor randint_cpu(int64_t n, int64_t seed, int64_t range) {
  auto opts = torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU).pinned_memory(true);
  auto t = torch::empty({n}, opts);
  mtgen::randint_fill(t.data_ptr<int32_t>(), n, (uint64_t)seed, (uint32_t)range);
  return t;
}


// One fused setup kernel: splits the uploaded rng draws into agent/food, seeds the
// per-env LCG state and clears the accumulators, so run() issues one launch here
// instead of half a dozen ATen ops.
__global__ void init_run_kernel(const int *__restrict__ draws, int *__restrict__ agent,
                                int *__restrict__ food, long long *__restrict__ rng,
                                float *__restrict__ rewards, float *__restrict__ last_logits,
                                int N, long long rng_base) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i >= N) return;
  agent[i * 2 + 0] = draws[i * 2 + 0];
  agent[i * 2 + 1] = draws[i * 2 + 1];
  food[i * 2 + 0] = draws[(N + i) * 2 + 0];
  food[i * 2 + 1] = draws[(N + i) * 2 + 1];
  rng[i] = rng_base + i;
  rewards[i] = 0.0f;
  last_logits[i * 4 + 0] = 0.0f;
  last_logits[i * 4 + 1] = 0.0f;
  last_logits[i * 4 + 2] = 0.0f;
  last_logits[i * 4 + 3] = 0.0f;
}

std::vector<torch::Tensor> init_run(torch::Tensor draws, int64_t N, int64_t rng_base) {
  auto iopt = draws.options();
  auto fopt = draws.options().dtype(torch::kFloat32);
  auto agent = torch::empty({N, 2}, iopt);
  auto food = torch::empty({N, 2}, iopt);
  auto rng = torch::empty({N}, draws.options().dtype(torch::kInt64));
  auto rewards = torch::empty({N}, fopt);
  auto last_logits = torch::empty({N, NACT}, fopt);
  init_run_kernel<<<(int)((N + 255) / 256), 256, 0, at::cuda::getCurrentCUDAStream()>>>(
      draws.data_ptr<int>(), agent.data_ptr<int>(), food.data_ptr<int>(),
      reinterpret_cast<long long *>(rng.data_ptr<int64_t>()), rewards.data_ptr<float>(),
      last_logits.data_ptr<float>(), (int)N, (long long)rng_base);
  return {agent, food, rng, rewards, last_logits};
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("prep_wfrag", &prep_wfrag);
  m.def("randint_cpu", &randint_cpu);
  m.def("init_run", &init_run);
  m.def("prep_l0", &prep_l0);
  m.def("rollout", &rollout);
  m.def("policy_forward_cuda", &policy_forward_cuda);
  m.def("env_step_cuda", &env_step_cuda);
}
"""


def _build():
    import hashlib
    import shutil
    import sys

    from torch.utils.cpp_extension import _get_build_directory, load

    if shutil.which("ninja") is None:  # venv installs land next to the interpreter
        os.environ["PATH"] = str(Path(sys.executable).parent) + os.pathsep + os.environ["PATH"]

    cc = torch.cuda.get_device_capability(0)
    arch = f"{cc[0]}{cc[1]}"
    tag = hashlib.sha1(_CUDA_SRC.encode()).hexdigest()[:10]
    name = f"grid_mingru_sps_{arch}_{tag}"
    build_dir = Path(_get_build_directory(name, verbose=False))
    build_dir.mkdir(parents=True, exist_ok=True)
    src = build_dir / "kernels.cu"
    if not src.exists() or src.read_text() != _CUDA_SRC:
        src.write_text(_CUDA_SRC)
    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", f"{cc[0]}.{cc[1]}")

    def _try(sfx):
        return load(
            name=name,
            sources=[str(src)],
            extra_cuda_cflags=[
                "-O3",
                f"-gencode=arch=compute_{arch}{sfx},code=sm_{arch}{sfx}",
                "-DNDEBUG",
            ],
            extra_cflags=["-O3"],
            verbose=False,
        )

    try:  # sm_90a/100a/120a expose the arch-specific instruction set
        return _try("a" if cc[0] >= 9 else "")
    except Exception:
        return _try("")


_EXT = None


def _ext():
    global _EXT
    if _EXT is None:
        _EXT = _build()
    return _EXT


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.data.copy_(tmp)

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


# ----------------------------------------------------------------------------------
# public single-step APIs (exact fp32 paths)
# ----------------------------------------------------------------------------------
def policy_forward(model, obs: torch.Tensor, state: torch.Tensor):
    ext = _ext()
    wgruT = model.w_gru.detach().transpose(1, 2).contiguous()
    logits, nstate, value = ext.policy_forward_cuda(
        obs.contiguous().float(),
        state.contiguous().float(),
        model.w_enc.detach().contiguous().float(),
        model.b_enc.detach().contiguous().float(),
        wgruT.float(),
        model.w_a.detach().contiguous().float(),
        model.b_a.detach().contiguous().float(),
        model.w_v.detach().contiguous().float(),
        model.b_v.detach().contiguous().float(),
    )
    return logits, nstate, value


def env_step(agent: torch.Tensor, food: torch.Tensor, actions: torch.Tensor,
             rng_state: torch.Tensor):
    ext = _ext()
    a, f, r, g = ext.env_step_cuda(
        agent.contiguous().float(),
        food.contiguous().float(),
        actions.contiguous().to(torch.int64),
        rng_state.contiguous().to(torch.int64),
    )
    return a.to(agent.dtype), f.to(food.dtype), r, g


# ----------------------------------------------------------------------------------
# rollout
# ----------------------------------------------------------------------------------
def _prepare(model):
    ext = _ext()
    w_gru = model.w_gru.detach()
    wfrag0 = ext.prep_l0(w_gru[0].contiguous(), model.w_enc.detach().contiguous(),
                         model.b_enc.detach().contiguous())
    wfrag = ext.prep_wfrag(w_gru[1:].contiguous())
    return (wfrag, wfrag0, model.w_a.detach().contiguous(), model.b_a.detach().contiguous())


def run(num_envs: int, horizon: int, seed: int, model=None) -> dict:
    ext = _ext()
    device = torch.device("cuda:0")
    if model is None:
        model = Model()
    if model.w_gru.device != device:
        model = model.to(device)

    if num_envs <= 0:
        z = torch.zeros(0, device=device)
        return {"rewards": z, "positions": torch.zeros(0, 2, dtype=torch.int64, device=device),
                "last_logits": torch.zeros(0, NUM_ACTIONS, device=device),
                "state": torch.zeros(0, GRU_LAYERS, HIDDEN, device=device)}

    wfrag, wfrag0, wa, ba = _prepare(model)

    # Reference draws agent then food from a fresh cpu mt19937; randint_cpu reproduces
    # that 32-bit stream exactly (int32/int64 consume identical draws for range < 2^32)
    # into pinned memory, ~6x faster than the generic ATen path.
    draws = ext.randint_cpu(4 * num_envs, seed, BOARD).to(device, non_blocking=True)
    agent, food, rng_state, rewards, last_logits = ext.init_run(draws, num_envs, seed * 10007)
    state = torch.zeros(num_envs, GRU_LAYERS, HIDDEN, device=device)
    anyhit = torch.zeros(max(horizon, 1), device=device, dtype=torch.int32)

    ext.rollout(wfrag, wfrag0, wa, ba, state, agent, food, rng_state,
                rewards, last_logits, anyhit, num_envs, horizon)

    return {
        "rewards": rewards,
        "positions": agent.to(torch.int64),
        "last_logits": last_logits,
        "state": state,
    }


def get_init_inputs():
    return []


def get_inputs():
    return []


_ext()  # compile at import time, not inside a timed run()