Grid + MinGRU SPS Muse Spark 1.3

"""Vectorized grid-foraging + 3xMinGRU(h=256) rollout, fused CUDA path.

run() executes the horizon as: per step, cuBLAS fp32 GEMMs (mm with
out=, pre-transposed weights) for the 3 MinGRU gate projections,
interleaved with three custom CUDA kernels:

  k_respawn_obs : LCG food-respawn from the previous step + obs build +
                  encoder Linear(4->256) fused in (warp per env)
  k_math        : MinGRU gate math (sigmoid/tanh/highway) + state update
                  per layer (x3)
  k_act_env     : action head Linear(256->4) fused in + argmax + env
                  move/clamp + reward + any-hit reduction (warp per env)

The whole H-step sequence is captured once in a CUDA graph (per
(num_envs, horizon, weights) key, with persistent buffers) and replayed,
so per-step launch overhead collapses to a single graph launch.

Math matches the reference exactly (fp32 assoc-order differences only).
Env integers (positions/food/rng/reward) are bitwise exact.
"""

import os

import torch
import torch.nn as nn
import torch.nn.functional as F

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


def _mingru_g(x: torch.Tensor) -> torch.Tensor:
    return torch.tanh(x)


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)


def policy_forward(model: Model, obs: torch.Tensor, state: torch.Tensor):
    """obs (N,4), state (N,L,H) -> logits (N,4), new_state (N,L,H), value (N,)."""
    h = F.linear(obs, model.w_enc, model.b_enc)
    new_states = []
    for layer in range(GRU_LAYERS):
        st = state[:, layer, :]
        gates = F.linear(h, model.w_gru[layer])
        zh, zg, zp = gates.split(HIDDEN, dim=-1)
        out = st + torch.sigmoid(zg) * (_mingru_g(zh) - st)
        p = torch.sigmoid(zp)
        h = p * out + (1.0 - p) * h
        new_states.append(out)
    new_state = torch.stack(new_states, dim=1)
    logits = F.linear(h, model.w_a, model.b_a)
    value = F.linear(h, model.w_v, model.b_v).squeeze(-1)
    return logits, new_state, value


def obs_from_state(agent: torch.Tensor, food: torch.Tensor) -> torch.Tensor:
    return torch.stack(
        [
            (food[:, 0] - agent[:, 0]) / BOARD,
            (food[:, 1] - agent[:, 1]) / BOARD,
            agent[:, 0] / (BOARD - 1),
            agent[:, 1] / (BOARD - 1),
        ],
        dim=-1,
    )


def _lcg_step(rng: torch.Tensor) -> torch.Tensor:
    return (rng * 6364136223846793005 + 1) & 0x7FFFFFFFFFFFFFFF


def env_step(
    agent: torch.Tensor,
    food: torch.Tensor,
    actions: torch.Tensor,
    rng_state: torch.Tensor,
):
    """Deterministic env step. rng_state: (N,) int64 for food respawns."""
    delta = torch.zeros_like(agent)
    delta[:, 1] = torch.where(actions == 0, -torch.ones_like(delta[:, 1]), delta[:, 1])
    delta[:, 1] = torch.where(actions == 1, torch.ones_like(delta[:, 1]), delta[:, 1])
    delta[:, 0] = torch.where(actions == 2, -torch.ones_like(delta[:, 0]), delta[:, 0])
    delta[:, 0] = torch.where(actions == 3, torch.ones_like(delta[:, 0]), delta[:, 0])
    agent = (agent + delta).clamp(0, BOARD - 1)
    hit = (agent == food).all(dim=-1)
    reward = hit.float()
    rng_state = rng_state.clone()
    if hit.any():
        rng_state = _lcg_step(rng_state)
        fx = (rng_state % BOARD).to(agent.dtype)
        rng_state = _lcg_step(rng_state)
        fy = (rng_state % BOARD).to(agent.dtype)
        new_food = torch.stack([fx, fy], dim=-1)
        food = food.clone()
        food[hit] = new_food[hit]
    return agent, food, reward, rng_state


_CUDA_SRC = r'''
#include <cstdint>
#include <cuda_runtime.h>

static const unsigned long long LCG_A = 6364136223846793005ULL;
static const unsigned long long LCG_M = 0x7FFFFFFFFFFFFFFFULL;

// Food respawn (from previous step's hits) + obs build + encoder
// Linear(4->256) fused in (K=4 dots, weights L2-resident).
// One warp (32 threads) per env so the 256 hh writes are fully coalesced
// (one thread per env would stride by 256 floats across the warp).
__global__ void k_respawn_obs(
    const float* __restrict__ agent,
    float* __restrict__ food,
    int64_t* __restrict__ rng,
    const unsigned char* __restrict__ mask,
    const int* __restrict__ any_hit,
    float* __restrict__ hh,
    const float* __restrict__ w_enc,
    const float* __restrict__ b_enc,
    int N) {
  int lane = (int)threadIdx.x & 31;
  int e = (int)blockIdx.x * 8 + ((int)threadIdx.x >> 5);
  bool ok = (e < N);
  float ax = 0.f, ay = 0.f, fx = 0.f, fy = 0.f;
  if (ok) {
    ax = agent[e * 2]; ay = agent[e * 2 + 1];
    fx = food[e * 2]; fy = food[e * 2 + 1];
  }
  if (lane == 0 && ok && *any_hit) {
    unsigned long long r = (unsigned long long)rng[e];
    r = r * LCG_A + 1ULL; r &= LCG_M;
    int nfx = (int)(r % 11ULL);
    r = r * LCG_A + 1ULL; r &= LCG_M;
    int nfy = (int)(r % 11ULL);
    rng[e] = (int64_t)r;
    if (mask[e]) {
      fx = (float)nfx; fy = (float)nfy;
      food[e * 2] = fx; food[e * 2 + 1] = fy;
    }
  }
  __syncwarp(0xffffffff);
  float o0 = (fx - ax) / 11.0f;
  float o1 = (fy - ay) / 11.0f;
  float o2 = ax / 10.0f;
  float o3 = ay / 10.0f;
  if (ok) {
#pragma unroll
    for (int c = 0; c < 8; ++c) {
      int o = lane + c * 32;
      hh[(size_t)e * 256u + (size_t)o] = b_enc[o]
          + w_enc[o * 4 + 0] * o0 + w_enc[o * 4 + 1] * o1
          + w_enc[o * 4 + 2] * o2 + w_enc[o * 4 + 3] * o3;
    }
  }
}

// MinGRU gate math for one layer (gates come from a cuBLAS GEMM).
// One thread per (env, hidden). State layout: (N,3,256) contiguous;
// l selects the layer slice.
__global__ void k_math(
    const float* __restrict__ gates,
    float* __restrict__ hh,
    float* __restrict__ state,
    int l, int N) {
  int idx = (int)blockIdx.x * (int)blockDim.x + (int)threadIdx.x;
  int e = idx >> 8, h = idx & 255;
  if (e >= N) return;
  size_t gb = (size_t)e * 768u;
  float zh = gates[gb + (size_t)h];
  float zg = gates[gb + 256u + (size_t)h];
  float zp = gates[gb + 512u + (size_t)h];
  size_t sb = (size_t)e * 768u + (size_t)l * 256u;
  float stv = state[sb + (size_t)h];
  float hcur = hh[(size_t)e * 256u + (size_t)h];
  float cand = tanhf(zh);
  float gz = 1.0f / (1.0f + expf(-zg));
  float out = stv + gz * (cand - stv);
  float gp = 1.0f / (1.0f + expf(-zp));
  state[sb + (size_t)h] = out;
  hh[(size_t)e * 256u + (size_t)h] = gp * out + (1.0f - gp) * hcur;
}

// Action head Linear(256->4) fused in + greedy action + env step +
// reward + cross-block any-hit reduction. Logits are also stored every
// step (the final step's values are returned).
__global__ void k_act_env(
    const float* __restrict__ hh,
    const float* __restrict__ w_a,
    const float* __restrict__ b_a,
    float* __restrict__ last_logits,
    float* __restrict__ agent,
    float* __restrict__ food,
    float* __restrict__ rewards,
    int64_t* __restrict__ positions,
    unsigned char* __restrict__ mask,
    int* __restrict__ any_hit,
    int* __restrict__ blk_flags,
    int* __restrict__ blk_cnt,
    int N) {
  // One warp (32 threads) per env: the 256-wide head dots are strip-mined
  // across lanes (coalesced hh/weight reads) and reduced with shuffles.
  int lane = (int)threadIdx.x & 31;
  int e = (int)blockIdx.x * 8 + ((int)threadIdx.x >> 5);
  bool ok = (e < N);
  __shared__ int s_any;
  if ((int)threadIdx.x == 0) s_any = 0;
  float p0 = 0.f, p1 = 0.f, p2 = 0.f, p3 = 0.f;
  if (ok) {
#pragma unroll
    for (int c = 0; c < 8; ++c) {
      int j = lane + c * 32;
      float hv = hh[(size_t)e * 256u + (size_t)j];
      p0 += w_a[j] * hv;
      p1 += w_a[256 + j] * hv;
      p2 += w_a[512 + j] * hv;
      p3 += w_a[768 + j] * hv;
    }
  }
#pragma unroll
  for (int off = 16; off > 0; off >>= 1) {
    p0 += __shfl_down_sync(0xffffffff, p0, off);
    p1 += __shfl_down_sync(0xffffffff, p1, off);
    p2 += __shfl_down_sync(0xffffffff, p2, off);
    p3 += __shfl_down_sync(0xffffffff, p3, off);
  }
  __syncthreads();
  if (lane == 0 && ok) {
    float lg0 = b_a[0] + p0, lg1 = b_a[1] + p1;
    float lg2 = b_a[2] + p2, lg3 = b_a[3] + p3;
    last_logits[e * 4] = lg0;
    last_logits[e * 4 + 1] = lg1;
    last_logits[e * 4 + 2] = lg2;
    last_logits[e * 4 + 3] = lg3;
    int act = 0;
    float best = lg0;
    if (lg1 > best) { best = lg1; act = 1; }
    if (lg2 > best) { best = lg2; act = 2; }
    if (lg3 > best) { best = lg3; act = 3; }
    int aix = (int)agent[e * 2];
    int aiy = (int)agent[e * 2 + 1];
    int fxi = (int)food[e * 2];
    int fyi = (int)food[e * 2 + 1];
    if (act == 0) aiy -= 1;
    else if (act == 1) aiy += 1;
    else if (act == 2) aix -= 1;
    else aix += 1;
    if (aix < 0) aix = 0; else if (aix > 10) aix = 10;
    if (aiy < 0) aiy = 0; else if (aiy > 10) aiy = 10;
    agent[e * 2] = (float)aix;
    agent[e * 2 + 1] = (float)aiy;
    int hit = (aix == fxi && aiy == fyi) ? 1 : 0;
    rewards[e] += (float)hit;
    mask[e] = (unsigned char)hit;
    positions[e * 2] = (int64_t)aix;
    positions[e * 2 + 1] = (int64_t)aiy;
    if (hit) atomicOr(&s_any, 1);
  }
  __syncthreads();
  if ((int)threadIdx.x == 0) {
    blk_flags[blockIdx.x] = s_any;
    if (atomicAdd(blk_cnt, 1) == (int)gridDim.x - 1) {
      int a = 0;
      for (int b = 0; b < (int)gridDim.x; ++b) a |= blk_flags[b];
      *any_hit = a;
      *blk_cnt = 0;
    }
  }
}
'''

_CPP_SRC = r'''
#include <torch/extension.h>
#include <cstdint>

void k_respawn_obs_wrap(
    torch::Tensor agent, torch::Tensor food, torch::Tensor rng,
    torch::Tensor mask, torch::Tensor any_hit, torch::Tensor hh,
    torch::Tensor w_enc, torch::Tensor b_enc, int64_t N);
void k_math_wrap(
    torch::Tensor gates, torch::Tensor hh, torch::Tensor state,
    int64_t l, int64_t N);
void k_act_env_wrap(
    torch::Tensor hh, torch::Tensor w_a, torch::Tensor b_a,
    torch::Tensor last_logits, torch::Tensor agent, torch::Tensor food,
    torch::Tensor rewards, torch::Tensor positions, torch::Tensor mask,
    torch::Tensor any_hit, torch::Tensor blk_flags, torch::Tensor blk_cnt,
    int64_t N);
'''

_CU_WRAPPERS = r'''
#include <torch/extension.h>
#include <c10/cuda/CUDAStream.h>

void k_respawn_obs_wrap(
    torch::Tensor agent, torch::Tensor food, torch::Tensor rng,
    torch::Tensor mask, torch::Tensor any_hit, torch::Tensor hh,
    torch::Tensor w_enc, torch::Tensor b_enc, int64_t N) {
  int threads = 256;
  int nb = (int)((N + 8 - 1) / 8);
  if (nb == 0) return;
  k_respawn_obs<<<nb, threads, 0, c10::cuda::getCurrentCUDAStream().stream()>>>(
      agent.data_ptr<float>(), food.data_ptr<float>(), rng.data_ptr<int64_t>(),
      mask.data_ptr<unsigned char>(), any_hit.data_ptr<int>(),
      hh.data_ptr<float>(), w_enc.data_ptr<float>(), b_enc.data_ptr<float>(),
      (int)N);
}

void k_math_wrap(
    torch::Tensor gates, torch::Tensor hh, torch::Tensor state,
    int64_t l, int64_t N) {
  int threads = 256;
  int64_t total = N * 256;
  int nb = (int)((total + threads - 1) / threads);
  if (nb == 0) return;
  k_math<<<nb, threads, 0, c10::cuda::getCurrentCUDAStream().stream()>>>(
      gates.data_ptr<float>(), hh.data_ptr<float>(), state.data_ptr<float>(),
      (int)l, (int)N);
}

void k_act_env_wrap(
    torch::Tensor hh, torch::Tensor w_a, torch::Tensor b_a,
    torch::Tensor last_logits, torch::Tensor agent, torch::Tensor food,
    torch::Tensor rewards, torch::Tensor positions, torch::Tensor mask,
    torch::Tensor any_hit, torch::Tensor blk_flags, torch::Tensor blk_cnt,
    int64_t N) {
  int threads = 256;
  int nb = (int)((N + 8 - 1) / 8);
  if (nb == 0) return;
  k_act_env<<<nb, threads, 0, c10::cuda::getCurrentCUDAStream().stream()>>>(
      hh.data_ptr<float>(), w_a.data_ptr<float>(), b_a.data_ptr<float>(),
      last_logits.data_ptr<float>(), agent.data_ptr<float>(),
      food.data_ptr<float>(), rewards.data_ptr<float>(),
      positions.data_ptr<int64_t>(), mask.data_ptr<unsigned char>(),
      any_hit.data_ptr<int>(), blk_flags.data_ptr<int>(),
      blk_cnt.data_ptr<int>(), (int)N);
}
'''

_mod = None


def _get_mod():
    global _mod
    if _mod is not None:
        return _mod
    from torch.utils.cpp_extension import load_inline

    if torch.cuda.is_available():
        try:
            cap = torch.cuda.get_device_capability()
            os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "%d.%d" % (cap[0], cap[1]))
        except Exception:
            pass
    os.environ.setdefault("MAX_JOBS", "8")
    try:
        _mod = load_inline(
            name="grid_mingru_hybrid",
            cpp_sources=_CPP_SRC,
            cuda_sources=_CUDA_SRC + _CU_WRAPPERS,
            functions=["k_respawn_obs_wrap", "k_math_wrap", "k_act_env_wrap"],
            extra_cflags=["-O3"],
            extra_cuda_cflags=["-O3"],
        )
    except Exception:
        # Fallback: default arch list (e.g. older nvcc without the native arch).
        os.environ.pop("TORCH_CUDA_ARCH_LIST", None)
        _mod = load_inline(
            name="grid_mingru_hybrid_fb",
            cpp_sources=_CPP_SRC,
            cuda_sources=_CUDA_SRC + _CU_WRAPPERS,
            functions=["k_respawn_obs_wrap", "k_math_wrap", "k_act_env_wrap"],
            extra_cflags=["-O3"],
            extra_cuda_cflags=["-O3"],
        )
    return _mod


_GCACHE = {}


def _init_buffers(b, seed):
    n = b["agent"].shape[0]
    dev = b["agent"].device
    g = torch.Generator(device="cpu")
    g.manual_seed(seed)
    # CPU randint keeps the exact reference init sequence. Draws go into
    # persistent pinned staging (same op/shape/dtype as the reference, so
    # the sequence is identical); the pin cost is paid once at alloc.
    torch.randint(0, BOARD, (n, 2), generator=g, out=b["pin_ai"])
    b["gi"].copy_(b["pin_ai"], non_blocking=True)
    b["agent"].copy_(b["gi"])
    torch.randint(0, BOARD, (n, 2), generator=g, out=b["pin_fi"])
    b["gf"].copy_(b["pin_fi"], non_blocking=True)
    b["food"].copy_(b["gf"])
    torch.arange(n, device=dev, dtype=torch.int64, out=b["rng"])
    b["rng"].add_(seed * 10007)
    b["state"].zero_()
    b["rewards"].zero_()
    b["mask"].zero_()
    b["any_hit"].zero_()
    b["blk_cnt"].zero_()
    b["blk_flags"].zero_()


def run(num_envs: int, horizon: int, seed: int, model: Model | None = None) -> dict:
    device = torch.device("cuda", torch.cuda.current_device())
    if model is None:
        model = Model()
    n, h = int(num_envs), int(horizon)
    topts = {"device": device, "dtype": torch.float32}

    if n == 0 or h == 0:
        g = torch.Generator(device="cpu")
        g.manual_seed(seed)
        agent = torch.randint(0, BOARD, (n, 2), generator=g).float().to(device)
        return {
            "rewards": torch.zeros(n, **topts),
            "positions": agent.round().long(),
            "last_logits": torch.zeros(n, NUM_ACTIONS, **topts),
        }

    key = (
        n,
        h,
        model.w_enc.data_ptr(),
        model.w_gru.data_ptr(),
        model.w_a.data_ptr(),
        model.b_enc.data_ptr(),
        model.b_a.data_ptr(),
    )
    ent = _GCACHE.get(key)
    if ent is None:
        mod = _get_mod()
        b = {}
        b["agent"] = torch.empty(n, 2, **topts)
        b["food"] = torch.empty(n, 2, **topts)
        b["rng"] = torch.empty(n, device=device, dtype=torch.int64)
        b["state"] = torch.empty(n, GRU_LAYERS, HIDDEN, **topts)
        b["rewards"] = torch.empty(n, **topts)
        b["positions"] = torch.empty(n, 2, device=device, dtype=torch.int64)
        b["mask"] = torch.empty(n, device=device, dtype=torch.uint8)
        b["hh"] = torch.empty(n, HIDDEN, **topts)
        b["gates"] = torch.empty(n, GRU_OUT, **topts)
        b["logits"] = torch.empty(n, NUM_ACTIONS, **topts)
        b["any_hit"] = torch.zeros(1, device=device, dtype=torch.int32)
        b["blk_cnt"] = torch.zeros(1, device=device, dtype=torch.int32)
        b["blk_flags"] = torch.zeros((n + 7) // 8, device=device, dtype=torch.int32)
        b["pin_ai"] = torch.empty(n, 2, dtype=torch.int64, pin_memory=True)
        b["pin_fi"] = torch.empty(n, 2, dtype=torch.int64, pin_memory=True)
        b["gi"] = torch.empty(n, 2, device=device, dtype=torch.int64)
        b["gf"] = torch.empty(n, 2, device=device, dtype=torch.int64)
        # Pre-transposed gru weights for mm with out= (no allocs in graph).
        # Encoder/head linears are fused into the elem kernels; their
        # original-layout weights are used directly.
        b["w_gru_t"] = model.w_gru.detach().to(**topts).transpose(1, 2).contiguous()
        b["w_enc"] = model.w_enc.detach().to(**topts).contiguous()
        b["b_enc"] = model.b_enc.detach().to(**topts).contiguous()
        b["w_a"] = model.w_a.detach().to(**topts).contiguous()
        b["b_a"] = model.b_a.detach().to(**topts).contiguous()

        def step():
            mod.k_respawn_obs_wrap(
                b["agent"], b["food"], b["rng"], b["mask"], b["any_hit"],
                b["hh"], b["w_enc"], b["b_enc"], n,
            )
            for l in range(GRU_LAYERS):
                torch.mm(b["hh"], b["w_gru_t"][l], out=b["gates"])
                mod.k_math_wrap(b["gates"], b["hh"], b["state"], l, n)
            mod.k_act_env_wrap(
                b["hh"], b["w_a"], b["b_a"], b["logits"], b["agent"],
                b["food"], b["rewards"], b["positions"], b["mask"],
                b["any_hit"], b["blk_flags"], b["blk_cnt"], n,
            )

        with torch.no_grad():
            _init_buffers(b, seed)
            step()
            step()
            step()  # warmup (instruction/data caches etc.)
            try:
                gr = torch.cuda.CUDAGraph()
                with torch.cuda.graph(gr):
                    for _ in range(h):
                        step()
                ent = (b, gr, step)
            except Exception:
                ent = (b, None, step)
            _GCACHE[key] = ent

    b, gr, step = ent
    with torch.no_grad():
        _init_buffers(b, seed)
        if gr is not None:
            gr.replay()
        else:
            for _ in range(h):
                step()
    return {
        "rewards": b["rewards"].detach(),
        "positions": b["positions"].detach(),
        "last_logits": b["logits"].detach(),
    }