from __future__ import annotations

import os
import sys
from pathlib import Path

# Ensure ninja from the active venv is on PATH for load_inline.
_venv_bin = str(Path(sys.executable).parent)
if _venv_bin not in os.environ.get("PATH", ""):
    os.environ["PATH"] = _venv_bin + os.pathsep + os.environ.get("PATH", "")
os.environ["KBH_GPU_LOCK_HELD"] = "1"

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 <cuda_runtime.h>
#include <cstdint>
#include <cmath>

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

__device__ __forceinline__ float fast_tanh(float x) {
    return tanhf(x);
}

__device__ __forceinline__ uint64_t lcg_step_dev(uint64_t rng) {
    return (rng * 6364136223846793005ULL + 1ULL) & 0x7FFFFFFFFFFFFFFFULL;
}

// -----------------------------------------------------------------------------
// Exact GPU implementation of PyTorch CPU MT19937 torch.randint(0, 11, (N, 2))
// -----------------------------------------------------------------------------
__global__ void mt19937_init_kernel(
    uint64_t seed,
    int num_envs,
    int8_t* __restrict__ out_agent_food
) {
    __shared__ uint32_t mt[2][624];
    int tid = threadIdx.x;
    if (tid == 0) {
        uint32_t s = (uint32_t)(seed & 0xffffffffULL);
        mt[0][0] = s;
        for (int i = 1; i < 624; ++i) {
            s = 1812433253U * (s ^ (s >> 30)) + (uint32_t)i;
            mt[0][i] = s;
        }
    }
    __syncthreads();

    int cur = 0;
    int total = num_envs * 4;
    for (int base = 0; base < total; base += 624) {
        int nxt = cur ^ 1;
        if (tid < 227) {
            uint32_t y0 = (mt[cur][tid] & 0x80000000U) | (mt[cur][tid + 1] & 0x7fffffffU);
            uint32_t v0 = mt[cur][tid + 397] ^ (y0 >> 1) ^ ((y0 & 1U) ? 0x9908b0dfU : 0U);
            mt[nxt][tid] = v0;

            int out0 = base + tid;
            if (out0 < total) {
                uint32_t t0 = v0;
                t0 ^= (t0 >> 11);
                t0 ^= (t0 << 7) & 0x9d2c5680U;
                t0 ^= (t0 << 15) & 0xefc60000U;
                t0 ^= (t0 >> 18);
                out_agent_food[out0] = (int8_t)(t0 % 11U);
            }

            int k1 = tid + 227;
            uint32_t y1 = (mt[cur][k1] & 0x80000000U) | (mt[cur][k1 + 1] & 0x7fffffffU);
            uint32_t v1 = v0 ^ (y1 >> 1) ^ ((y1 & 1U) ? 0x9908b0dfU : 0U);
            mt[nxt][k1] = v1;

            int out1 = base + k1;
            if (out1 < total) {
                uint32_t t1 = v1;
                t1 ^= (t1 >> 11);
                t1 ^= (t1 << 7) & 0x9d2c5680U;
                t1 ^= (t1 << 15) & 0xefc60000U;
                t1 ^= (t1 >> 18);
                out_agent_food[out1] = (int8_t)(t1 % 11U);
            }

            if (tid < 170) {
                int k2 = tid + 454;
                uint32_t m_next;
                if (tid == 169) {
                    uint32_t y_zero = (mt[cur][0] & 0x80000000U) | (mt[cur][1] & 0x7fffffffU);
                    m_next = mt[cur][397] ^ (y_zero >> 1) ^ ((y_zero & 1U) ? 0x9908b0dfU : 0U);
                } else {
                    m_next = mt[cur][k2 + 1];
                }
                uint32_t y2 = (mt[cur][k2] & 0x80000000U) | (m_next & 0x7fffffffU);
                uint32_t v2 = v1 ^ (y2 >> 1) ^ ((y2 & 1U) ? 0x9908b0dfU : 0U);
                mt[nxt][k2] = v2;

                int out2 = base + k2;
                if (out2 < total) {
                    uint32_t t2 = v2;
                    t2 ^= (t2 >> 11);
                    t2 ^= (t2 << 7) & 0x9d2c5680U;
                    t2 ^= (t2 << 15) & 0xefc60000U;
                    t2 ^= (t2 >> 18);
                    out_agent_food[out2] = (int8_t)(t2 % 11U);
                }
            }
        }
        cur = nxt;
        __syncthreads();
    }
}

// -----------------------------------------------------------------------------
// Standalone env_step kernels (exact hit.any() across batch)
// -----------------------------------------------------------------------------
__global__ void env_step_phase1_kernel(
    int N,
    const float* __restrict__ agent_in,
    const float* __restrict__ food_in,
    const int64_t* __restrict__ actions,
    float* __restrict__ agent_out,
    float* __restrict__ reward_out,
    int* __restrict__ any_hit_flag
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;

    float ax = agent_in[i * 2 + 0];
    float ay = agent_in[i * 2 + 1];
    int64_t act = actions[i];
    if (act == 0) ay -= 1.0f;
    else if (act == 1) ay += 1.0f;
    else if (act == 2) ax -= 1.0f;
    else if (act == 3) ax += 1.0f;

    ax = fminf(fmaxf(ax, 0.0f), 10.0f);
    ay = fminf(fmaxf(ay, 0.0f), 10.0f);
    agent_out[i * 2 + 0] = ax;
    agent_out[i * 2 + 1] = ay;

    float fx = food_in[i * 2 + 0];
    float fy = food_in[i * 2 + 1];
    bool hit = (ax == fx) && (ay == fy);
    reward_out[i] = hit ? 1.0f : 0.0f;
    if (hit) {
        *any_hit_flag = 1;
    }
}

__global__ void env_step_phase2_kernel(
    int N,
    const float* __restrict__ food_in,
    const float* __restrict__ reward_out,
    const int64_t* __restrict__ rng_in,
    float* __restrict__ food_out,
    int64_t* __restrict__ rng_out,
    const int* __restrict__ any_hit_flag
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;

    int any_hit = *any_hit_flag;
    uint64_t rng = (uint64_t)rng_in[i];
    float fx = food_in[i * 2 + 0];
    float fy = food_in[i * 2 + 1];

    if (any_hit) {
        rng = lcg_step_dev(rng);
        float nfx = (float)(rng % 11ULL);
        rng = lcg_step_dev(rng);
        float nfy = (float)(rng % 11ULL);
        if (reward_out[i] > 0.5f) {
            fx = nfx;
            fy = nfy;
        }
    }
    food_out[i * 2 + 0] = fx;
    food_out[i * 2 + 1] = fy;
    rng_out[i] = (int64_t)rng;
}

// -----------------------------------------------------------------------------
// Standalone policy_forward kernel (1-pass over k=0..255 with [l][k][gate][j])
// -----------------------------------------------------------------------------
template <int G, int M_REG>
__global__ __launch_bounds__(G * 64, 2)
void policy_forward_kernel(
    int N,
    const float* __restrict__ obs,        // (N, 4)
    const float* __restrict__ state_in,   // (N, 3, 256)
    const float* __restrict__ w_enc,      // (256, 4)
    const float* __restrict__ b_enc,      // (256,)
    const float* __restrict__ w_gru_T,    // (3, 256, 3, 256) [l][k][gate][j]
    const float* __restrict__ w_a,        // (4, 256)
    const float* __restrict__ b_a,        // (4,)
    const float* __restrict__ w_v,        // (256,)
    const float* __restrict__ b_v,        // (1,)
    float* __restrict__ logits_out,       // (N, 4)
    float* __restrict__ state_out,        // (N, 3, 256)
    float* __restrict__ value_out         // (N,)
) {
    constexpr int M_TILE = G * M_REG;
    __shared__ float4 s_h[256][G];
    __shared__ float s_red[2][G][M_REG][5];

    const int col_tid = threadIdx.x;   // 0..63
    const int env_group = threadIdx.y; // 0..G-1
    const int j_base = col_tid * 4;
    const int env_base = blockIdx.x * M_TILE + env_group * M_REG;

    float h_reg[M_REG][4];

    // 1. Encoder
    {
        float4 b_enc_r = __ldg(&reinterpret_cast<const float4*>(b_enc)[col_tid]);
        float b_arr[4] = {b_enc_r.x, b_enc_r.y, b_enc_r.z, b_enc_r.w};
        float4 o_arr[M_REG];
        #pragma unroll
        for (int dm = 0; dm < M_REG; ++dm) {
            int e = env_base + dm;
            o_arr[dm] = (e < N) ? __ldg(&reinterpret_cast<const float4*>(obs)[e]) : make_float4(0.f, 0.f, 0.f, 0.f);
        }
        #pragma unroll
        for (int dj = 0; dj < 4; ++dj) {
            float4 we = __ldg(&reinterpret_cast<const float4*>(w_enc)[j_base + dj]);
            #pragma unroll
            for (int dm = 0; dm < M_REG; ++dm) {
                float acc = b_arr[dj];
                acc = fmaf(o_arr[dm].x, we.x, acc);
                acc = fmaf(o_arr[dm].y, we.y, acc);
                acc = fmaf(o_arr[dm].z, we.z, acc);
                acc = fmaf(o_arr[dm].w, we.w, acc);
                h_reg[dm][dj] = acc;
            }
            s_h[j_base + dj][env_group] = make_float4(h_reg[0][dj], h_reg[1][dj], h_reg[2][dj], h_reg[3][dj]);
        }
    }
    __syncthreads();

    // 2. 3 x MinGRU layers
    #pragma unroll 1
    for (int l = 0; l < 3; ++l) {
        const float4* __restrict__ w_k = reinterpret_cast<const float4*>(w_gru_T) + l * (256 * 3 * 64) + col_tid;

        float acc_zh[M_REG][4] = {0};
        float acc_zg[M_REG][4] = {0};
        float acc_zp[M_REG][4] = {0};

        #pragma unroll 8
        for (int k = 0; k < 256; ++k) {
            float4 h4 = s_h[k][env_group];
            float4 w0 = __ldg(&w_k[0 * 64]);
            float4 w1 = __ldg(&w_k[1 * 64]);
            float4 w2 = __ldg(&w_k[2 * 64]);
            w_k += 3 * 64;

            float h_arr[4] = {h4.x, h4.y, h4.z, h4.w};
            #pragma unroll
            for (int dm = 0; dm < M_REG; ++dm) {
                acc_zh[dm][0] = fmaf(h_arr[dm], w0.x, acc_zh[dm][0]);
                acc_zh[dm][1] = fmaf(h_arr[dm], w0.y, acc_zh[dm][1]);
                acc_zh[dm][2] = fmaf(h_arr[dm], w0.z, acc_zh[dm][2]);
                acc_zh[dm][3] = fmaf(h_arr[dm], w0.w, acc_zh[dm][3]);

                acc_zg[dm][0] = fmaf(h_arr[dm], w1.x, acc_zg[dm][0]);
                acc_zg[dm][1] = fmaf(h_arr[dm], w1.y, acc_zg[dm][1]);
                acc_zg[dm][2] = fmaf(h_arr[dm], w1.z, acc_zg[dm][2]);
                acc_zg[dm][3] = fmaf(h_arr[dm], w1.w, acc_zg[dm][3]);

                acc_zp[dm][0] = fmaf(h_arr[dm], w2.x, acc_zp[dm][0]);
                acc_zp[dm][1] = fmaf(h_arr[dm], w2.y, acc_zp[dm][1]);
                acc_zp[dm][2] = fmaf(h_arr[dm], w2.z, acc_zp[dm][2]);
                acc_zp[dm][3] = fmaf(h_arr[dm], w2.w, acc_zp[dm][3]);
            }
        }

        #pragma unroll
        for (int dm = 0; dm < M_REG; ++dm) {
            int e = env_base + dm;
            float4 st4 = (e < N)
                ? __ldg(&reinterpret_cast<const float4*>(state_in + (e * 3 + l) * 256)[col_tid])
                : make_float4(0.f, 0.f, 0.f, 0.f);
            float st_arr[4] = {st4.x, st4.y, st4.z, st4.w};
            float st_out[4];
            #pragma unroll
            for (int dj = 0; dj < 4; ++dj) {
                float zh_t = fast_tanh(acc_zh[dm][dj]);
                float g = fast_sigmoid(acc_zg[dm][dj]);
                float out = st_arr[dj] + g * (zh_t - st_arr[dj]);
                st_out[dj] = out;
                float p = fast_sigmoid(acc_zp[dm][dj]);
                h_reg[dm][dj] = p * out + (1.0f - p) * h_reg[dm][dj];
            }
            if (e < N) {
                reinterpret_cast<float4*>(state_out + (e * 3 + l) * 256)[col_tid] =
                    make_float4(st_out[0], st_out[1], st_out[2], st_out[3]);
            }
        }

        __syncthreads();
        if (l < 2) {
            #pragma unroll
            for (int dj = 0; dj < 4; ++dj) {
                s_h[j_base + dj][env_group] = make_float4(h_reg[0][dj], h_reg[1][dj], h_reg[2][dj], h_reg[3][dj]);
            }
            __syncthreads();
        }
    }

    // 3. Action head (4) + Value head (1)
    float4 wa_r[4];
    #pragma unroll
    for (int a = 0; a < 4; ++a) {
        wa_r[a] = __ldg(&reinterpret_cast<const float4*>(w_a + a * 256)[col_tid]);
    }
    float4 wv_r = __ldg(&reinterpret_cast<const float4*>(w_v)[col_tid]);

    int warp_id = col_tid >> 5;
    int lane = col_tid & 31;

    #pragma unroll
    for (int dm = 0; dm < M_REG; ++dm) {
        float sum[5];
        #pragma unroll
        for (int a = 0; a < 4; ++a) {
            sum[a] = h_reg[dm][0] * wa_r[a].x + h_reg[dm][1] * wa_r[a].y +
                     h_reg[dm][2] * wa_r[a].z + h_reg[dm][3] * wa_r[a].w;
        }
        sum[4] = h_reg[dm][0] * wv_r.x + h_reg[dm][1] * wv_r.y +
                 h_reg[dm][2] * wv_r.z + h_reg[dm][3] * wv_r.w;

        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            #pragma unroll
            for (int c = 0; c < 5; ++c) {
                sum[c] += __shfl_down_sync(0xffffffff, sum[c], offset);
            }
        }
        if (lane == 0) {
            #pragma unroll
            for (int c = 0; c < 5; ++c) {
                s_red[warp_id][env_group][dm][c] = sum[c];
            }
        }
    }
    __syncthreads();

    if (col_tid < M_REG) {
        int dm = col_tid;
        int e = env_base + dm;
        if (e < N) {
            float4 ba4 = __ldg(&reinterpret_cast<const float4*>(b_a)[0]);
            float l0 = s_red[0][env_group][dm][0] + s_red[1][env_group][dm][0] + ba4.x;
            float l1 = s_red[0][env_group][dm][1] + s_red[1][env_group][dm][1] + ba4.y;
            float l2 = s_red[0][env_group][dm][2] + s_red[1][env_group][dm][2] + ba4.z;
            float l3 = s_red[0][env_group][dm][3] + s_red[1][env_group][dm][3] + ba4.w;
            reinterpret_cast<float4*>(logits_out)[e] = make_float4(l0, l1, l2, l3);
            value_out[e] = s_red[0][env_group][dm][4] + s_red[1][env_group][dm][4] + b_v[0];
        }
    }
}

// -----------------------------------------------------------------------------
// Fused rollout kernel 1 (num_envs <= 4096):
//   G = 2, M_REG = 11, C_REG = 2 (M_TILE = 22 -> 187 CTAs = 99.5% of 1 wave on 188 SMs)
//   - 3 x float4 packed weight loads per kp pair (k = 2*kp, 2*kp+1)
//   - Layers 0 & 1 recurrent state in shared memory s_state[44][256] (44 KB)
//   - Layer 2 recurrent state in registers st2[11][2] (22 regs)
//   - Reuses acc_zh across layers (0 register spills, 188 regs/thread)
// -----------------------------------------------------------------------------
template <int M_REG, int UNROLL_KP>
__global__ __launch_bounds__(256, 1)
void kernel_c2_g2_any(
    int num_envs,
    int horizon,
    int64_t seed,
    const int8_t* __restrict__ init_agent_food,
    const float* __restrict__ w_enc,
    const float* __restrict__ b_enc,
    const float4* __restrict__ w_gru_packed,
    const float* __restrict__ w_a,
    const float* __restrict__ b_a,
    float* __restrict__ rewards_out,
    int64_t* __restrict__ positions_out,
    float* __restrict__ last_logits_out
) {
    constexpr int G = 2;
    constexpr int M_TILE = G * M_REG;
    constexpr int NUM_THREADS = 256;
    constexpr int V_DIM = (M_REG + 3) / 4;
    constexpr int ROWS = 4 * M_REG;

    extern __shared__ float4 dyn_smem[];
    float4 (*s_h)[V_DIM][256] = reinterpret_cast<float4 (*)[V_DIM][256]>(dyn_smem);
    float (*s_state)[NUM_THREADS] = reinterpret_cast<float (*)[NUM_THREADS]>(dyn_smem + G * V_DIM * 256);

    __shared__ float4 s_red[4][G][M_REG];
    float4* s_obs = &s_red[0][0][0];

    const int col_tid = threadIdx.x;
    const int env_group = threadIdx.y;
    const int tid = env_group * 128 + col_tid;
    const int j_base = col_tid * 2;
    const int cta_env_base = blockIdx.x * M_TILE;

    float st2[M_REG][2];
    #pragma unroll
    for (int r = 0; r < ROWS; ++r) {
        s_state[r][tid] = 0.0f;
    }
    #pragma unroll
    for (int m = 0; m < M_REG; ++m) {
        st2[m][0] = 0.0f;
        st2[m][1] = 0.0f;
    }

    int ax = 0, ay = 0, fx = 0, fy = 0;
    uint64_t rng = 0;
    float reward_sum = 0.0f;
    int e_global_tid = cta_env_base + tid;

    if (tid < M_TILE && e_global_tid < num_envs) {
        ax = (int)init_agent_food[e_global_tid * 2 + 0];
        ay = (int)init_agent_food[e_global_tid * 2 + 1];
        fx = (int)init_agent_food[num_envs * 2 + e_global_tid * 2 + 0];
        fy = (int)init_agent_food[num_envs * 2 + e_global_tid * 2 + 1];
        rng = (uint64_t)((int64_t)e_global_tid + seed * 10007LL);
        s_obs[tid] = make_float4(
            (float)(fx - ax) / 11.0f,
            (float)(fy - ay) / 11.0f,
            (float)ax / 10.0f,
            (float)ay / 10.0f
        );
    }
    __syncthreads();

    const int warp_id = col_tid >> 5;
    const int lane = col_tid & 31;

    #pragma unroll 1
    for (int t = 0; t < horizon; ++t) {
        // 1. Encoder
        {
            float2 b_enc_r = __ldg(&reinterpret_cast<const float2*>(b_enc)[col_tid]);
            float b_arr[2] = {b_enc_r.x, b_enc_r.y};
            #pragma unroll
            for (int dj = 0; dj < 2; ++dj) {
                float4 we = __ldg(&reinterpret_cast<const float4*>(w_enc)[j_base + dj]);
                float b = b_arr[dj];
                #pragma unroll
                for (int v = 0; v < V_DIM; ++v) {
                    int m0 = v * 4 + 0;
                    int m1 = v * 4 + 1;
                    int m2 = v * 4 + 2;
                    int m3 = v * 4 + 3;
                    float4 o0 = s_obs[env_group * M_REG + m0];
                    float4 o1 = (m1 < M_REG) ? s_obs[env_group * M_REG + m1] : make_float4(0.f, 0.f, 0.f, 0.f);
                    float4 o2 = (m2 < M_REG) ? s_obs[env_group * M_REG + m2] : make_float4(0.f, 0.f, 0.f, 0.f);
                    float4 o3 = (m3 < M_REG) ? s_obs[env_group * M_REG + m3] : make_float4(0.f, 0.f, 0.f, 0.f);
                    float h0 = fmaf(o0.x, we.x, fmaf(o0.y, we.y, fmaf(o0.z, we.z, fmaf(o0.w, we.w, b))));
                    float h1 = (m1 < M_REG) ? fmaf(o1.x, we.x, fmaf(o1.y, we.y, fmaf(o1.z, we.z, fmaf(o1.w, we.w, b)))) : 0.0f;
                    float h2 = (m2 < M_REG) ? fmaf(o2.x, we.x, fmaf(o2.y, we.y, fmaf(o2.z, we.z, fmaf(o2.w, we.w, b)))) : 0.0f;
                    float h3 = (m3 < M_REG) ? fmaf(o3.x, we.x, fmaf(o3.y, we.y, fmaf(o3.z, we.z, fmaf(o3.w, we.w, b)))) : 0.0f;
                    s_h[env_group][v][dj * 128 + col_tid] = make_float4(h0, h1, h2, h3);
                }
            }
        }
        __syncthreads();

        // 2. 3 x MinGRU layers
        float acc_zh[M_REG][2];
        #pragma unroll 1
        for (int l = 0; l < 3; ++l) {
            const float4* __restrict__ w_k = w_gru_packed + l * (128 * 3 * 128) + col_tid;

            float acc_zg[M_REG][2];
            float acc_zp[M_REG][2];
            #pragma unroll
            for (int dm = 0; dm < M_REG; ++dm) {
                acc_zh[dm][0] = 0.0f; acc_zh[dm][1] = 0.0f;
                acc_zg[dm][0] = 0.0f; acc_zg[dm][1] = 0.0f;
                acc_zp[dm][0] = 0.0f; acc_zp[dm][1] = 0.0f;
            }

            #pragma unroll UNROLL_KP
            for (int kp = 0; kp < 128; ++kp) {
                float4 hg0 = __ldg(&w_k[0 * 128]);
                float4 p01 = __ldg(&w_k[1 * 128]);
                float4 hg1 = __ldg(&w_k[2 * 128]);
                w_k += 3 * 128;

                // Substep 0: k = 2*kp
                {
                    float h_a[4 * V_DIM];
                    #pragma unroll
                    for (int v = 0; v < V_DIM; ++v) {
                        float4 hv = s_h[env_group][v][kp];
                        h_a[v * 4 + 0] = hv.x;
                        h_a[v * 4 + 1] = hv.y;
                        h_a[v * 4 + 2] = hv.z;
                        h_a[v * 4 + 3] = hv.w;
                    }
                    #pragma unroll
                    for (int dm = 0; dm < M_REG; ++dm) {
                        acc_zh[dm][0] = fmaf(h_a[dm], hg0.x, acc_zh[dm][0]);
                        acc_zh[dm][1] = fmaf(h_a[dm], hg0.y, acc_zh[dm][1]);
                        acc_zg[dm][0] = fmaf(h_a[dm], hg0.z, acc_zg[dm][0]);
                        acc_zg[dm][1] = fmaf(h_a[dm], hg0.w, acc_zg[dm][1]);
                        acc_zp[dm][0] = fmaf(h_a[dm], p01.x, acc_zp[dm][0]);
                        acc_zp[dm][1] = fmaf(h_a[dm], p01.y, acc_zp[dm][1]);
                    }
                }
                // Substep 1: k = 2*kp + 1
                {
                    float h_b[4 * V_DIM];
                    #pragma unroll
                    for (int v = 0; v < V_DIM; ++v) {
                        float4 hv = s_h[env_group][v][kp + 128];
                        h_b[v * 4 + 0] = hv.x;
                        h_b[v * 4 + 1] = hv.y;
                        h_b[v * 4 + 2] = hv.z;
                        h_b[v * 4 + 3] = hv.w;
                    }
                    #pragma unroll
                    for (int dm = 0; dm < M_REG; ++dm) {
                        acc_zh[dm][0] = fmaf(h_b[dm], hg1.x, acc_zh[dm][0]);
                        acc_zh[dm][1] = fmaf(h_b[dm], hg1.y, acc_zh[dm][1]);
                        acc_zg[dm][0] = fmaf(h_b[dm], hg1.z, acc_zg[dm][0]);
                        acc_zg[dm][1] = fmaf(h_b[dm], hg1.w, acc_zg[dm][1]);
                        acc_zp[dm][0] = fmaf(h_b[dm], p01.z, acc_zp[dm][0]);
                        acc_zp[dm][1] = fmaf(h_b[dm], p01.w, acc_zp[dm][1]);
                    }
                }
            }

            if (l < 2) {
                int r_base = l * (2 * M_REG);
                #pragma unroll
                for (int dj = 0; dj < 2; ++dj) {
                    float h_old[4 * V_DIM];
                    #pragma unroll
                    for (int v = 0; v < V_DIM; ++v) {
                        float4 hv = s_h[env_group][v][dj * 128 + col_tid];
                        h_old[v * 4 + 0] = hv.x;
                        h_old[v * 4 + 1] = hv.y;
                        h_old[v * 4 + 2] = hv.z;
                        h_old[v * 4 + 3] = hv.w;
                    }
                    #pragma unroll
                    for (int dm = 0; dm < M_REG; ++dm) {
                        int r = r_base + dm * 2 + dj;
                        float st_m = s_state[r][tid];
                        float zh_t = fast_tanh(acc_zh[dm][dj]);
                        float g = fast_sigmoid(acc_zg[dm][dj]);
                        float out = st_m + g * (zh_t - st_m);
                        s_state[r][tid] = out;
                        float p = fast_sigmoid(acc_zp[dm][dj]);
                        acc_zh[dm][dj] = p * out + (1.0f - p) * h_old[dm];
                    }
                }
                __syncthreads();
                #pragma unroll
                for (int dj = 0; dj < 2; ++dj) {
                    #pragma unroll
                    for (int v = 0; v < V_DIM; ++v) {
                        int m0 = v * 4 + 0;
                        int m1 = v * 4 + 1;
                        int m2 = v * 4 + 2;
                        int m3 = v * 4 + 3;
                        s_h[env_group][v][dj * 128 + col_tid] = make_float4(
                            acc_zh[m0][dj],
                            (m1 < M_REG) ? acc_zh[m1][dj] : 0.0f,
                            (m2 < M_REG) ? acc_zh[m2][dj] : 0.0f,
                            (m3 < M_REG) ? acc_zh[m3][dj] : 0.0f
                        );
                    }
                }
                __syncthreads();
            } else {
                #pragma unroll
                for (int dj = 0; dj < 2; ++dj) {
                    float h_old[4 * V_DIM];
                    #pragma unroll
                    for (int v = 0; v < V_DIM; ++v) {
                        float4 hv = s_h[env_group][v][dj * 128 + col_tid];
                        h_old[v * 4 + 0] = hv.x;
                        h_old[v * 4 + 1] = hv.y;
                        h_old[v * 4 + 2] = hv.z;
                        h_old[v * 4 + 3] = hv.w;
                    }
                    #pragma unroll
                    for (int dm = 0; dm < M_REG; ++dm) {
                        float st_m = st2[dm][dj];
                        float zh_t = fast_tanh(acc_zh[dm][dj]);
                        float g = fast_sigmoid(acc_zg[dm][dj]);
                        float out = st_m + g * (zh_t - st_m);
                        st2[dm][dj] = out;
                        float p = fast_sigmoid(acc_zp[dm][dj]);
                        acc_zh[dm][dj] = p * out + (1.0f - p) * h_old[dm];
                    }
                }
            }
        }

        // 3. Action head reduction
        {
            float2 wa_r[4];
            #pragma unroll
            for (int a = 0; a < 4; ++a) {
                wa_r[a] = __ldg(&reinterpret_cast<const float2*>(w_a + a * 256)[col_tid]);
            }
            #pragma unroll
            for (int dm = 0; dm < M_REG; ++dm) {
                float sum[4];
                #pragma unroll
                for (int a = 0; a < 4; ++a) {
                    sum[a] = acc_zh[dm][0] * wa_r[a].x + acc_zh[dm][1] * wa_r[a].y;
                }
                #pragma unroll
                for (int offset = 16; offset > 0; offset >>= 1) {
                    #pragma unroll
                    for (int a = 0; a < 4; ++a) {
                        sum[a] += __shfl_down_sync(0xffffffff, sum[a], offset);
                    }
                }
                if (lane == 0) {
                    s_red[warp_id][env_group][dm] = make_float4(sum[0], sum[1], sum[2], sum[3]);
                }
            }
        }
        __syncthreads();

        // 4. Greedy action + env step
        if (tid < M_TILE && e_global_tid < num_envs) {
            int eg = tid / M_REG;
            int dm = tid % M_REG;
            float4 r0 = s_red[0][eg][dm];
            float4 r1 = s_red[1][eg][dm];
            float4 r2 = s_red[2][eg][dm];
            float4 r3 = s_red[3][eg][dm];
            float4 ba4 = __ldg(&reinterpret_cast<const float4*>(b_a)[0]);
            float l0 = (r0.x + r1.x) + (r2.x + r3.x) + ba4.x;
            float l1 = (r0.y + r1.y) + (r2.y + r3.y) + ba4.y;
            float l2 = (r0.z + r1.z) + (r2.z + r3.z) + ba4.z;
            float l3 = (r0.w + r1.w) + (r2.w + r3.w) + ba4.w;

            if (t == horizon - 1) {
                reinterpret_cast<float4*>(last_logits_out)[e_global_tid] = make_float4(l0, l1, l2, l3);
            }

            int act = 0;
            float best = l0;
            if (l1 > best) { best = l1; act = 1; }
            if (l2 > best) { best = l2; act = 2; }
            if (l3 > best) { best = l3; act = 3; }

            if (act == 0) ay = max(0, ay - 1);
            else if (act == 1) ay = min(10, ay + 1);
            else if (act == 2) ax = max(0, ax - 1);
            else if (act == 3) ax = min(10, ax + 1);

            bool hit = (ax == fx) && (ay == fy);
            if (hit) reward_sum += 1.0f;

            rng = lcg_step_dev(rng);
            int nfx = (int)(rng % 11ULL);
            rng = lcg_step_dev(rng);
            int nfy = (int)(rng % 11ULL);
            if (hit) {
                fx = nfx;
                fy = nfy;
            }

            if (t + 1 < horizon) {
                s_obs[tid] = make_float4(
                    (float)(fx - ax) / 11.0f,
                    (float)(fy - ay) / 11.0f,
                    (float)ax / 10.0f,
                    (float)ay / 10.0f
                );
            }
        }
        if (t + 1 < horizon) {
            __syncthreads();
        }
    }

    if (tid < M_TILE && e_global_tid < num_envs) {
        rewards_out[e_global_tid] = reward_sum;
        positions_out[e_global_tid * 2 + 0] = (int64_t)ax;
        positions_out[e_global_tid * 2 + 1] = (int64_t)ay;
    }
}

// -----------------------------------------------------------------------------
// Fused rollout kernel 2 (num_envs > 4096):
//   G = 2, M_REG = 22, C_REG = 2 (M_TILE = 44 -> 99.1%..99.5% wave efficiency on 188 SMs)
//   - 24 env-layers (48 rows = 48 KB) in shared memory s_state + 48 KB s_h = 96 KB SMEM
//   - Remaining 42 env-layers (84 rows) in coalesced L2-resident buffer l2_state
//   - Zero recurrent state registers live across inner kp loop (0 spills, 226 regs/thread)
// -----------------------------------------------------------------------------
template <int G, int M_REG, int SMEM_ENVS, int UNROLL_KP>
__global__ __launch_bounds__(G * 128, 1)
void kernel_c2_smem_l2_v2(
    int num_envs,
    int horizon,
    int64_t seed,
    const int8_t* __restrict__ init_agent_food,
    const float* __restrict__ w_enc,
    const float* __restrict__ b_enc,
    const float4* __restrict__ w_gru_packed,
    const float* __restrict__ w_a,
    const float* __restrict__ b_a,
    float* __restrict__ l2_state,
    float* __restrict__ rewards_out,
    int64_t* __restrict__ positions_out,
    float* __restrict__ last_logits_out
) {
    constexpr int M_TILE = G * M_REG;
    constexpr int NUM_THREADS = G * 128;
    constexpr int V_DIM = (M_REG + 3) / 4;
    constexpr int ROWS = 2 * SMEM_ENVS;
    constexpr int L2_ROWS = 2 * (3 * M_REG - SMEM_ENVS);

    extern __shared__ float4 dyn_smem[];
    float4 (*s_h)[V_DIM][256] = reinterpret_cast<float4 (*)[V_DIM][256]>(dyn_smem);
    float (*s_state)[NUM_THREADS] = reinterpret_cast<float (*)[NUM_THREADS]>(dyn_smem + G * V_DIM * 256);

    __shared__ float4 s_red[4][G][M_REG];
    float4* s_obs = &s_red[0][0][0];

    const int col_tid = threadIdx.x;
    const int env_group = threadIdx.y;
    const int tid = env_group * 128 + col_tid;
    const int j_base = col_tid * 2;
    const int cta_env_base = blockIdx.x * M_TILE;

    #pragma unroll
    for (int r = 0; r < ROWS; ++r) {
        s_state[r][tid] = 0.0f;
    }
    int l2_base = (blockIdx.x * L2_ROWS) * NUM_THREADS + tid;
    #pragma unroll
    for (int r_l2 = 0; r_l2 < L2_ROWS; ++r_l2) {
        l2_state[l2_base + r_l2 * NUM_THREADS] = 0.0f;
    }

    int ax = 0, ay = 0, fx = 0, fy = 0;
    uint64_t rng = 0;
    float reward_sum = 0.0f;
    int e_global_tid = cta_env_base + tid;

    if (tid < M_TILE && e_global_tid < num_envs) {
        ax = (int)init_agent_food[e_global_tid * 2 + 0];
        ay = (int)init_agent_food[e_global_tid * 2 + 1];
        fx = (int)init_agent_food[num_envs * 2 + e_global_tid * 2 + 0];
        fy = (int)init_agent_food[num_envs * 2 + e_global_tid * 2 + 1];
        rng = (uint64_t)((int64_t)e_global_tid + seed * 10007LL);
        s_obs[tid] = make_float4(
            (float)(fx - ax) / 11.0f,
            (float)(fy - ay) / 11.0f,
            (float)ax / 10.0f,
            (float)ay / 10.0f
        );
    }
    __syncthreads();

    const int warp_id = col_tid >> 5;
    const int lane = col_tid & 31;

    #pragma unroll 1
    for (int t = 0; t < horizon; ++t) {
        // 1. Encoder
        {
            float2 b_enc_r = __ldg(&reinterpret_cast<const float2*>(b_enc)[col_tid]);
            float b_arr[2] = {b_enc_r.x, b_enc_r.y};
            #pragma unroll
            for (int dj = 0; dj < 2; ++dj) {
                float4 we = __ldg(&reinterpret_cast<const float4*>(w_enc)[j_base + dj]);
                float b = b_arr[dj];
                #pragma unroll
                for (int v = 0; v < V_DIM; ++v) {
                    int m0 = v * 4 + 0;
                    int m1 = v * 4 + 1;
                    int m2 = v * 4 + 2;
                    int m3 = v * 4 + 3;
                    float4 o0 = s_obs[env_group * M_REG + m0];
                    float4 o1 = (m1 < M_REG) ? s_obs[env_group * M_REG + m1] : make_float4(0.f, 0.f, 0.f, 0.f);
                    float4 o2 = (m2 < M_REG) ? s_obs[env_group * M_REG + m2] : make_float4(0.f, 0.f, 0.f, 0.f);
                    float4 o3 = (m3 < M_REG) ? s_obs[env_group * M_REG + m3] : make_float4(0.f, 0.f, 0.f, 0.f);
                    float h0 = fmaf(o0.x, we.x, fmaf(o0.y, we.y, fmaf(o0.z, we.z, fmaf(o0.w, we.w, b))));
                    float h1 = (m1 < M_REG) ? fmaf(o1.x, we.x, fmaf(o1.y, we.y, fmaf(o1.z, we.z, fmaf(o1.w, we.w, b)))) : 0.0f;
                    float h2 = (m2 < M_REG) ? fmaf(o2.x, we.x, fmaf(o2.y, we.y, fmaf(o2.z, we.z, fmaf(o2.w, we.w, b)))) : 0.0f;
                    float h3 = (m3 < M_REG) ? fmaf(o3.x, we.x, fmaf(o3.y, we.y, fmaf(o3.z, we.z, fmaf(o3.w, we.w, b)))) : 0.0f;
                    s_h[env_group][v][dj * 128 + col_tid] = make_float4(h0, h1, h2, h3);
                }
            }
        }
        __syncthreads();

        // 2. 3 x MinGRU layers
        float acc_zh[M_REG][2];
        #pragma unroll 1
        for (int l = 0; l < 3; ++l) {
            const float4* __restrict__ w_k = w_gru_packed + l * (128 * 3 * 128) + col_tid;

            float acc_zg[M_REG][2];
            float acc_zp[M_REG][2];
            #pragma unroll
            for (int dm = 0; dm < M_REG; ++dm) {
                acc_zh[dm][0] = 0.0f; acc_zh[dm][1] = 0.0f;
                acc_zg[dm][0] = 0.0f; acc_zg[dm][1] = 0.0f;
                acc_zp[dm][0] = 0.0f; acc_zp[dm][1] = 0.0f;
            }

            #pragma unroll UNROLL_KP
            for (int kp = 0; kp < 128; ++kp) {
                float4 hg0 = __ldg(&w_k[0 * 128]);
                float4 p01 = __ldg(&w_k[1 * 128]);
                float4 hg1 = __ldg(&w_k[2 * 128]);
                w_k += 3 * 128;

                // Substep 0: k = 2*kp
                #pragma unroll
                for (int v = 0; v < V_DIM; ++v) {
                    float4 hv = s_h[env_group][v][kp];
                    float h_sub[4] = {hv.x, hv.y, hv.z, hv.w};
                    #pragma unroll
                    for (int sub = 0; sub < 4; ++sub) {
                        int dm = v * 4 + sub;
                        if (dm < M_REG) {
                            acc_zh[dm][0] = fmaf(h_sub[sub], hg0.x, acc_zh[dm][0]);
                            acc_zh[dm][1] = fmaf(h_sub[sub], hg0.y, acc_zh[dm][1]);
                            acc_zg[dm][0] = fmaf(h_sub[sub], hg0.z, acc_zg[dm][0]);
                            acc_zg[dm][1] = fmaf(h_sub[sub], hg0.w, acc_zg[dm][1]);
                            acc_zp[dm][0] = fmaf(h_sub[sub], p01.x, acc_zp[dm][0]);
                            acc_zp[dm][1] = fmaf(h_sub[sub], p01.y, acc_zp[dm][1]);
                        }
                    }
                }
                // Substep 1: k = 2*kp + 1
                #pragma unroll
                for (int v = 0; v < V_DIM; ++v) {
                    float4 hv = s_h[env_group][v][kp + 128];
                    float h_sub[4] = {hv.x, hv.y, hv.z, hv.w};
                    #pragma unroll
                    for (int sub = 0; sub < 4; ++sub) {
                        int dm = v * 4 + sub;
                        if (dm < M_REG) {
                            acc_zh[dm][0] = fmaf(h_sub[sub], hg1.x, acc_zh[dm][0]);
                            acc_zh[dm][1] = fmaf(h_sub[sub], hg1.y, acc_zh[dm][1]);
                            acc_zg[dm][0] = fmaf(h_sub[sub], hg1.z, acc_zg[dm][0]);
                            acc_zg[dm][1] = fmaf(h_sub[sub], hg1.w, acc_zg[dm][1]);
                            acc_zp[dm][0] = fmaf(h_sub[sub], p01.z, acc_zp[dm][0]);
                            acc_zp[dm][1] = fmaf(h_sub[sub], p01.w, acc_zp[dm][1]);
                        }
                    }
                }
            }

            // Epilogue: process 1 (dj, v) at a time -> only 4 floats of h_old live at once
            #pragma unroll
            for (int dj = 0; dj < 2; ++dj) {
                #pragma unroll
                for (int v = 0; v < V_DIM; ++v) {
                    float4 hv = s_h[env_group][v][dj * 128 + col_tid];
                    float h_sub[4] = {hv.x, hv.y, hv.z, hv.w};
                    #pragma unroll
                    for (int sub = 0; sub < 4; ++sub) {
                        int dm = v * 4 + sub;
                        if (dm < M_REG) {
                            int r = (l * M_REG + dm) * 2 + dj;
                            float st_m = (r < ROWS)
                                ? s_state[r][tid]
                                : l2_state[l2_base + (r - ROWS) * NUM_THREADS];
                            float zh_t = fast_tanh(acc_zh[dm][dj]);
                            float g = fast_sigmoid(acc_zg[dm][dj]);
                            float out = st_m + g * (zh_t - st_m);
                            if (r < ROWS) {
                                s_state[r][tid] = out;
                            } else {
                                l2_state[l2_base + (r - ROWS) * NUM_THREADS] = out;
                            }
                            float p = fast_sigmoid(acc_zp[dm][dj]);
                            acc_zh[dm][dj] = p * out + (1.0f - p) * h_sub[sub];
                        }
                    }
                }
            }

            if (l < 2) {
                __syncthreads();
                #pragma unroll
                for (int dj = 0; dj < 2; ++dj) {
                    #pragma unroll
                    for (int v = 0; v < V_DIM; ++v) {
                        int m0 = v * 4 + 0;
                        int m1 = v * 4 + 1;
                        int m2 = v * 4 + 2;
                        int m3 = v * 4 + 3;
                        s_h[env_group][v][dj * 128 + col_tid] = make_float4(
                            acc_zh[m0][dj],
                            (m1 < M_REG) ? acc_zh[m1][dj] : 0.0f,
                            (m2 < M_REG) ? acc_zh[m2][dj] : 0.0f,
                            (m3 < M_REG) ? acc_zh[m3][dj] : 0.0f
                        );
                    }
                }
                __syncthreads();
            }
        }

        // 3. Action head reduction
        {
            float2 wa_r[4];
            #pragma unroll
            for (int a = 0; a < 4; ++a) {
                wa_r[a] = __ldg(&reinterpret_cast<const float2*>(w_a + a * 256)[col_tid]);
            }
            #pragma unroll
            for (int dm = 0; dm < M_REG; ++dm) {
                float sum[4];
                #pragma unroll
                for (int a = 0; a < 4; ++a) {
                    sum[a] = acc_zh[dm][0] * wa_r[a].x + acc_zh[dm][1] * wa_r[a].y;
                }
                #pragma unroll
                for (int offset = 16; offset > 0; offset >>= 1) {
                    #pragma unroll
                    for (int a = 0; a < 4; ++a) {
                        sum[a] += __shfl_down_sync(0xffffffff, sum[a], offset);
                    }
                }
                if (lane == 0) {
                    s_red[warp_id][env_group][dm] = make_float4(sum[0], sum[1], sum[2], sum[3]);
                }
            }
        }
        __syncthreads();

        // 4. Greedy action + env step
        if (tid < M_TILE && e_global_tid < num_envs) {
            int eg = tid / M_REG;
            int dm = tid % M_REG;
            float4 r0 = s_red[0][eg][dm];
            float4 r1 = s_red[1][eg][dm];
            float4 r2 = s_red[2][eg][dm];
            float4 r3 = s_red[3][eg][dm];
            float4 ba4 = __ldg(&reinterpret_cast<const float4*>(b_a)[0]);
            float l0 = (r0.x + r1.x) + (r2.x + r3.x) + ba4.x;
            float l1 = (r0.y + r1.y) + (r2.y + r3.y) + ba4.y;
            float l2 = (r0.z + r1.z) + (r2.z + r3.z) + ba4.z;
            float l3 = (r0.w + r1.w) + (r2.w + r3.w) + ba4.w;

            if (t == horizon - 1) {
                reinterpret_cast<float4*>(last_logits_out)[e_global_tid] = make_float4(l0, l1, l2, l3);
            }

            int act = 0;
            float best = l0;
            if (l1 > best) { best = l1; act = 1; }
            if (l2 > best) { best = l2; act = 2; }
            if (l3 > best) { best = l3; act = 3; }

            if (act == 0) ay = max(0, ay - 1);
            else if (act == 1) ay = min(10, ay + 1);
            else if (act == 2) ax = max(0, ax - 1);
            else if (act == 3) ax = min(10, ax + 1);

            bool hit = (ax == fx) && (ay == fy);
            if (hit) reward_sum += 1.0f;

            rng = lcg_step_dev(rng);
            int nfx = (int)(rng % 11ULL);
            rng = lcg_step_dev(rng);
            int nfy = (int)(rng % 11ULL);
            if (hit) {
                fx = nfx;
                fy = nfy;
            }

            if (t + 1 < horizon) {
                s_obs[tid] = make_float4(
                    (float)(fx - ax) / 11.0f,
                    (float)(fy - ay) / 11.0f,
                    (float)ax / 10.0f,
                    (float)ay / 10.0f
                );
            }
        }
        if (t + 1 < horizon) {
            __syncthreads();
        }
    }

    if (tid < M_TILE && e_global_tid < num_envs) {
        rewards_out[e_global_tid] = reward_sum;
        positions_out[e_global_tid * 2 + 0] = (int64_t)ax;
        positions_out[e_global_tid * 2 + 1] = (int64_t)ay;
    }
}

extern "C" {

void c_launch_mt19937_init(int64_t seed, int num_envs, int8_t* out_agent_food, cudaStream_t stream) {
    mt19937_init_kernel<<<1, 256, 0, stream>>>((uint64_t)seed, num_envs, out_agent_food);
}

void c_launch_env_step(
    int N,
    const float* agent_in,
    const float* food_in,
    const int64_t* actions,
    const int64_t* rng_in,
    float* agent_out,
    float* food_out,
    float* reward_out,
    int64_t* rng_out,
    int* any_hit_flag,
    cudaStream_t stream
) {
    cudaMemsetAsync(any_hit_flag, 0, sizeof(int), stream);
    int threads = 256;
    int blocks = (N + threads - 1) / threads;
    env_step_phase1_kernel<<<blocks, threads, 0, stream>>>(
        N, agent_in, food_in, actions, agent_out, reward_out, any_hit_flag
    );
    env_step_phase2_kernel<<<blocks, threads, 0, stream>>>(
        N, food_in, reward_out, rng_in, food_out, rng_out, any_hit_flag
    );
}

void c_launch_policy_forward(
    int N,
    const float* obs,
    const float* state_in,
    const float* w_enc,
    const float* b_enc,
    const float* w_gru_T,
    const float* w_a,
    const float* b_a,
    const float* w_v,
    const float* b_v,
    float* logits_out,
    float* state_out,
    float* value_out,
    cudaStream_t stream
) {
    constexpr int G = 4;
    constexpr int M_REG = 4;
    constexpr int M_TILE = G * M_REG;
    dim3 block(64, G);
    dim3 grid((N + M_TILE - 1) / M_TILE);
    policy_forward_kernel<G, M_REG><<<grid, block, 0, stream>>>(
        N, obs, state_in, w_enc, b_enc, w_gru_T, w_a, b_a, w_v, b_v,
        logits_out, state_out, value_out
    );
}

void c_launch_fused_rollout(
    int num_envs,
    int horizon,
    int64_t seed,
    int8_t* init_agent_food,
    const float* w_enc,
    const float* b_enc,
    const float* w_gru_packed,
    const float* w_a,
    const float* b_a,
    float* l2_state,
    float* rewards_out,
    int64_t* positions_out,
    float* last_logits_out,
    cudaStream_t stream
) {
    mt19937_init_kernel<<<1, 256, 0, stream>>>((uint64_t)seed, num_envs, init_agent_food);
    const float4* wp = reinterpret_cast<const float4*>(w_gru_packed);

    if (num_envs <= 4096) {
        constexpr int G = 2;
        constexpr int M_REG = 11;
        constexpr int UNROLL_KP = 32;
        constexpr int M_TILE = G * M_REG;
        constexpr int V_DIM = (M_REG + 3) / 4;
        constexpr int ROWS = 4 * M_REG;
        constexpr int smem = (G * V_DIM * 256) * sizeof(float4) + ROWS * 256 * sizeof(float);
        static bool attr_set_m11 = false;
        if (!attr_set_m11) {
            cudaFuncSetAttribute(
                kernel_c2_g2_any<M_REG, UNROLL_KP>,
                cudaFuncAttributeMaxDynamicSharedMemorySize,
                smem
            );
            attr_set_m11 = true;
        }
        dim3 block(128, G);
        dim3 grid((num_envs + M_TILE - 1) / M_TILE);
        kernel_c2_g2_any<M_REG, UNROLL_KP><<<grid, block, smem, stream>>>(
            num_envs, horizon, seed, init_agent_food,
            w_enc, b_enc, wp, w_a, b_a,
            rewards_out, positions_out, last_logits_out
        );
    } else {
        constexpr int G = 2;
        constexpr int M_REG = 22;
        constexpr int SMEM_ENVS = 24;
        constexpr int UNROLL_KP = 16;
        constexpr int M_TILE = G * M_REG;
        constexpr int NUM_THREADS = G * 128;
        constexpr int V_DIM = (M_REG + 3) / 4;
        constexpr int ROWS = 2 * SMEM_ENVS;
        constexpr int smem = (G * V_DIM * 256) * sizeof(float4) + ROWS * NUM_THREADS * sizeof(float);
        static bool attr_set_m22 = false;
        if (!attr_set_m22) {
            cudaFuncSetAttribute(
                kernel_c2_smem_l2_v2<G, M_REG, SMEM_ENVS, UNROLL_KP>,
                cudaFuncAttributeMaxDynamicSharedMemorySize,
                smem
            );
            attr_set_m22 = true;
        }
        dim3 block(128, G);
        dim3 grid((num_envs + M_TILE - 1) / M_TILE);
        kernel_c2_smem_l2_v2<G, M_REG, SMEM_ENVS, UNROLL_KP><<<grid, block, smem, stream>>>(
            num_envs, horizon, seed, init_agent_food,
            w_enc, b_enc, wp, w_a, b_a, l2_state,
            rewards_out, positions_out, last_logits_out
        );
    }
}

} // extern "C"
"""

_CPP_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cstdint>
#include <vector>

extern "C" {
void c_launch_mt19937_init(int64_t seed, int num_envs, int8_t* out_agent_food, cudaStream_t stream);
void c_launch_env_step(
    int N,
    const float* agent_in,
    const float* food_in,
    const int64_t* actions,
    const int64_t* rng_in,
    float* agent_out,
    float* food_out,
    float* reward_out,
    int64_t* rng_out,
    int* any_hit_flag,
    cudaStream_t stream
);
void c_launch_policy_forward(
    int N,
    const float* obs,
    const float* state_in,
    const float* w_enc,
    const float* b_enc,
    const float* w_gru_T,
    const float* w_a,
    const float* b_a,
    const float* w_v,
    const float* b_v,
    float* logits_out,
    float* state_out,
    float* value_out,
    cudaStream_t stream
);
void c_launch_fused_rollout(
    int num_envs,
    int horizon,
    int64_t seed,
    int8_t* init_agent_food,
    const float* w_enc,
    const float* b_enc,
    const float* w_gru_packed,
    const float* w_a,
    const float* b_a,
    float* l2_state,
    float* rewards_out,
    int64_t* positions_out,
    float* last_logits_out,
    cudaStream_t stream
);
}

void launch_mt19937_init(int64_t seed, int num_envs, torch::Tensor out_agent_food) {
    c_launch_mt19937_init(seed, num_envs, out_agent_food.data_ptr<int8_t>(), at::cuda::getCurrentCUDAStream());
}

std::vector<torch::Tensor> launch_env_step(
    torch::Tensor agent,
    torch::Tensor food,
    torch::Tensor actions,
    torch::Tensor rng_state,
    torch::Tensor any_hit_flag
) {
    int N = agent.size(0);
    auto agent_out = torch::empty_like(agent);
    auto food_out = torch::empty_like(food);
    auto reward_out = torch::empty({N}, agent.options());
    auto rng_out = torch::empty_like(rng_state);
    c_launch_env_step(
        N,
        agent.data_ptr<float>(),
        food.data_ptr<float>(),
        actions.data_ptr<int64_t>(),
        rng_state.data_ptr<int64_t>(),
        agent_out.data_ptr<float>(),
        food_out.data_ptr<float>(),
        reward_out.data_ptr<float>(),
        rng_out.data_ptr<int64_t>(),
        any_hit_flag.data_ptr<int>(),
        at::cuda::getCurrentCUDAStream()
    );
    return {agent_out, food_out, reward_out, rng_out};
}

std::vector<torch::Tensor> launch_policy_forward(
    torch::Tensor obs,
    torch::Tensor state,
    torch::Tensor w_enc,
    torch::Tensor b_enc,
    torch::Tensor w_gru_T,
    torch::Tensor w_a,
    torch::Tensor b_a,
    torch::Tensor w_v,
    torch::Tensor b_v
) {
    int N = obs.size(0);
    auto logits_out = torch::empty({N, 4}, obs.options());
    auto state_out = torch::empty_like(state);
    auto value_out = torch::empty({N}, obs.options());
    c_launch_policy_forward(
        N,
        obs.data_ptr<float>(),
        state.data_ptr<float>(),
        w_enc.data_ptr<float>(),
        b_enc.data_ptr<float>(),
        w_gru_T.data_ptr<float>(),
        w_a.data_ptr<float>(),
        b_a.data_ptr<float>(),
        w_v.data_ptr<float>(),
        b_v.data_ptr<float>(),
        logits_out.data_ptr<float>(),
        state_out.data_ptr<float>(),
        value_out.data_ptr<float>(),
        at::cuda::getCurrentCUDAStream()
    );
    return {logits_out, state_out, value_out};
}

void launch_fused_rollout(
    int num_envs,
    int horizon,
    int64_t seed,
    torch::Tensor init_agent_food,
    torch::Tensor w_enc,
    torch::Tensor b_enc,
    torch::Tensor w_gru_packed,
    torch::Tensor w_a,
    torch::Tensor b_a,
    torch::Tensor l2_state,
    torch::Tensor rewards_out,
    torch::Tensor positions_out,
    torch::Tensor last_logits_out
) {
    c_launch_fused_rollout(
        num_envs,
        horizon,
        seed,
        init_agent_food.data_ptr<int8_t>(),
        w_enc.data_ptr<float>(),
        b_enc.data_ptr<float>(),
        w_gru_packed.data_ptr<float>(),
        w_a.data_ptr<float>(),
        b_a.data_ptr<float>(),
        l2_state.data_ptr<float>(),
        rewards_out.data_ptr<float>(),
        positions_out.data_ptr<int64_t>(),
        last_logits_out.data_ptr<float>(),
        at::cuda::getCurrentCUDAStream()
    );
}
"""

_ext = load_inline(
    name="grid_mingru_cuda_ext_v9",
    cpp_sources=_CPP_SRC,
    cuda_sources=_CUDA_SRC,
    functions=[
        "launch_mt19937_init",
        "launch_env_step",
        "launch_policy_forward",
        "launch_fused_rollout",
    ],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)


def _pack_w_gru(w: torch.Tensor) -> torch.Tensor:
    wt = w.detach().view(3, 3, 128, 2, 128, 2)
    wp = wt.permute(0, 4, 5, 1, 2, 3).contiguous()
    hg0 = wp[:, :, 0, 0:2, :, :].permute(0, 1, 3, 2, 4).reshape(3, 128, 128, 4)
    p01 = wp[:, :, :, 2, :, :].permute(0, 1, 3, 2, 4).reshape(3, 128, 128, 4)
    hg1 = wp[:, :, 1, 0:2, :, :].permute(0, 1, 3, 2, 4).reshape(3, 128, 128, 4)
    return torch.stack([hg0, p01, hg1], dim=2).contiguous()


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._cached_w_gru_T: torch.Tensor | None = None
        self._cached_w_gru_packed: torch.Tensor | None = None
        self._cached_w_gru_ver: int = -1
        self._cached_w_gru_ptr: int = 0
        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)
        self._cached_w_gru_T = None
        self._cached_w_gru_packed = None

    def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):
        self._cached_w_gru_T = None
        self._cached_w_gru_packed = None
        return super().load_state_dict(state_dict, strict=strict, assign=assign)

    def _ensure_gru_cache(self) -> None:
        w = self.w_gru
        if (
            self._cached_w_gru_T is None
            or self._cached_w_gru_packed is None
            or self._cached_w_gru_ver != w._version
            or self._cached_w_gru_ptr != w.data_ptr()
            or self._cached_w_gru_T.device != w.device
        ):
            wd = w.detach()
            self._cached_w_gru_T = wd.view(3, 3, 256, 256).permute(0, 3, 1, 2).contiguous()
            self._cached_w_gru_packed = _pack_w_gru(wd)
            self._cached_w_gru_ver = w._version
            self._cached_w_gru_ptr = w.data_ptr()

    def get_w_gru_T(self) -> torch.Tensor:
        self._ensure_gru_cache()
        return self._cached_w_gru_T

    def get_w_gru_packed(self) -> torch.Tensor:
        self._ensure_gru_cache()
        return self._cached_w_gru_packed

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


_any_hit_flag_buf: torch.Tensor | None = None
_init_af_buf: torch.Tensor | None = None
_l2_state_buf: torch.Tensor | None = None
_default_model: Model | None = None


def _get_any_hit_flag(device: torch.device) -> torch.Tensor:
    global _any_hit_flag_buf
    if _any_hit_flag_buf is None or _any_hit_flag_buf.device != device:
        _any_hit_flag_buf = torch.zeros(1, dtype=torch.int32, device=device)
    return _any_hit_flag_buf


def _get_init_af_buf(num_envs: int, device: torch.device) -> torch.Tensor:
    global _init_af_buf
    req = max(num_envs, 65536) * 4
    if _init_af_buf is None or _init_af_buf.numel() < req or _init_af_buf.device != device:
        _init_af_buf = torch.empty(req, dtype=torch.int8, device=device)
    return _init_af_buf


def _get_l2_state_buf(num_envs: int, device: torch.device) -> torch.Tensor:
    global _l2_state_buf
    max_ctas = (max(num_envs, 65536) + 43) // 44
    req = max_ctas * 84 * 256
    if _l2_state_buf is None or _l2_state_buf.numel() < req or _l2_state_buf.device != device:
        _l2_state_buf = torch.empty(req, dtype=torch.float32, device=device)
    return _l2_state_buf


def policy_forward(model: Model, obs: torch.Tensor, state: torch.Tensor):
    obs_c = obs.contiguous()
    state_c = state.contiguous()
    if hasattr(model, "get_w_gru_T"):
        w_gru_T = model.get_w_gru_T()
    else:
        w_gru_T = model.w_gru.detach().view(3, 3, 256, 256).permute(0, 3, 1, 2).contiguous()
    logits, new_state, value = _ext.launch_policy_forward(
        obs_c,
        state_c,
        model.w_enc.contiguous(),
        model.b_enc.contiguous(),
        w_gru_T,
        model.w_a.contiguous(),
        model.b_a.contiguous(),
        model.w_v.contiguous(),
        model.b_v.contiguous(),
    )
    return logits, new_state, value


def env_step(
    agent: torch.Tensor,
    food: torch.Tensor,
    actions: torch.Tensor,
    rng_state: torch.Tensor,
):
    flag = _get_any_hit_flag(agent.device)
    agent_out, food_out, reward_out, rng_out = _ext.launch_env_step(
        agent.contiguous(),
        food.contiguous(),
        actions.contiguous(),
        rng_state.contiguous(),
        flag,
    )
    return agent_out, food_out, reward_out, rng_out


def run(num_envs: int, horizon: int, seed: int, model: Model | None = None) -> dict:
    global _default_model
    device = torch.device("cuda:0")
    if model is None:
        if _default_model is None:
            _default_model = Model().to(device).eval()
        model = _default_model
    elif next(model.parameters()).device != device:
        model = model.to(device)

    if num_envs < 512:
        af_buf = _get_init_af_buf(num_envs, device)
        _ext.launch_mt19937_init(seed, num_envs, af_buf)
        agent = af_buf[: num_envs * 2].view(num_envs, 2).float()
        food = af_buf[num_envs * 2 : num_envs * 4].view(num_envs, 2).float()
        rng_state = torch.arange(num_envs, device=device, dtype=torch.int64) + (seed * 10007)
        state = torch.zeros(num_envs, GRU_LAYERS, HIDDEN, device=device, dtype=torch.float32)
        rewards = torch.zeros(num_envs, device=device, dtype=torch.float32)
        last_logits = torch.zeros(num_envs, NUM_ACTIONS, device=device, dtype=torch.float32)
        for _t in range(horizon):
            obs = torch.stack(
                [
                    (food[:, 0] - agent[:, 0]) / BOARD,
                    (food[:, 1] - agent[:, 1]) / BOARD,
                    agent[:, 0] / (BOARD - 1),
                    agent[:, 1] / (BOARD - 1),
                ],
                dim=-1,
            )
            logits, state, _ = policy_forward(model, obs, state)
            last_logits = logits
            actions = torch.argmax(logits, dim=-1)
            agent, food, r, rng_state = env_step(agent, food, actions, rng_state)
            rewards = rewards + r
        return {
            "rewards": rewards,
            "positions": agent.round().long(),
            "last_logits": last_logits,
        }

    if hasattr(model, "get_w_gru_packed"):
        w_gru_packed = model.get_w_gru_packed()
    else:
        w_gru_packed = _pack_w_gru(model.w_gru)

    af_buf = _get_init_af_buf(num_envs, device)
    l2_buf = _get_l2_state_buf(num_envs, device)
    rewards = torch.empty(num_envs, dtype=torch.float32, device=device)
    positions = torch.empty((num_envs, 2), dtype=torch.int64, device=device)
    last_logits = torch.empty((num_envs, NUM_ACTIONS), dtype=torch.float32, device=device)

    _ext.launch_fused_rollout(
        num_envs,
        horizon,
        seed,
        af_buf,
        model.w_enc,
        model.b_enc,
        w_gru_packed,
        model.w_a,
        model.b_a,
        l2_buf,
        rewards,
        positions,
        last_logits,
    )

    return {
        "rewards": rewards,
        "positions": positions,
        "last_logits": last_logits,
    }