import math

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

# Default ModelNew is PyTorch-only. run_and_check: use check_kernel=False with backend=cuda
# until the candidate file contains a real CUDA kernel (see kernel_static_checker).


class Model(nn.Module):
    """
    Graph-Fused VLA (GF-VLA) style **scene graph relation update** in one forward.

    Objects are nodes with vision-language aligned embeddings; directed pairs
    (i → j) carry pairwise conditioning (geometry, language grounding, or fused
    VLA features). This reference fuses:

      LayerNorm(nodes) → multi-head attention over objects with **pairwise
      relation MLP bias** on edges → residual write-back.

    Intended as a fusion target for navigation / manipulation stacks that maintain
    explicit scene graphs and update them from multimodal observations.
    """

    def __init__(self, dim: int, pair_dim: int, num_heads: int):
        super().__init__()
        assert dim % num_heads == 0
        self.dim = dim
        self.pair_dim = pair_dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = 1.0 / math.sqrt(self.head_dim)

        self.norm = nn.LayerNorm(dim)
        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.rel_mlp = nn.Sequential(
            nn.Linear(pair_dim, pair_dim),
            nn.SiLU(),
            nn.Linear(pair_dim, num_heads),
        )
        self.out_proj = nn.Linear(dim, dim, bias=False)

    def forward(
        self,
        nodes: torch.Tensor,
        pair_feats: torch.Tensor,
        edge_mask: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            nodes: (B, N, D) object / region embeddings.
            pair_feats: (B, N, N, P) pairwise VLA or geometric conditioning.
            edge_mask: (B, N, N) float in {0, 1}; 1 keeps attention over column j.
        Returns:
            (B, N, D) updated node embeddings (residual around attention block).
        """
        B, N, D = nodes.shape
        _, N2, N3, P = pair_feats.shape
        assert N == N2 == N3 and P == self.pair_dim
        h = self.norm(nodes)

        q = self.q_proj(h).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(h).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(h).view(B, N, self.num_heads, self.head_dim).transpose(1, 2)

        logits = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        rel_bias = self.rel_mlp(pair_feats).permute(0, 3, 1, 2)
        logits = logits + rel_bias

        logits = logits.masked_fill(edge_mask.unsqueeze(1) < 0.5, float("-inf"))
        attn = F.softmax(logits, dim=-1)
        ctx = torch.matmul(attn, v).transpose(1, 2).reshape(B, N, D)
        return nodes + self.out_proj(ctx)


class ModelNew(nn.Module):
    """KernelBench candidate entry point; default delegates to ``Model``."""

    def __init__(self, dim: int, pair_dim: int, num_heads: int):
        super().__init__()
        self._impl = Model(dim, pair_dim, num_heads)

    def forward(
        self,
        nodes: torch.Tensor,
        pair_feats: torch.Tensor,
        edge_mask: torch.Tensor,
    ) -> torch.Tensor:
        return self._impl(nodes, pair_feats, edge_mask)


batch_size = 4
num_objects = 24
dim = 96
pair_dim = 16
num_heads = 4


def get_inputs():
    p = popcorn_pri
    nodes = torch.randn(p.jitter_int(batch_size), p.jitter_int(num_objects), dim)
    pair_feats = torch.randn(p.jitter_int(batch_size), p.jitter_int(num_objects), p.jitter_int(num_objects), pair_dim)
    edge_mask = torch.ones(p.jitter_int(batch_size), p.jitter_int(num_objects), p.jitter_int(num_objects))
    return [nodes, pair_feats, edge_mask]


def get_init_inputs():
    return [dim, pair_dim, num_heads]

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

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

torch::Tensor gf_vla_forward_cuda(torch::Tensor nodes, torch::Tensor pair_feats, torch::Tensor edge_mask,
                                  torch::Tensor ln_w, torch::Tensor ln_b,
                                  torch::Tensor q_w, torch::Tensor k_w, torch::Tensor v_w,
                                  torch::Tensor r1_w, torch::Tensor r1_b, torch::Tensor r2_w, torch::Tensor r2_b,
                                  torch::Tensor out_w);
'''

cuda_source = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

#define BATCH 4
#define NOBJ 24
#define DIM 96
#define HEADS 4
#define HDIM 24
#define PDIM 16

__global__ void ln_qkv_kernel(const float* __restrict__ nodes,
                              const float* __restrict__ ln_w,
                              const float* __restrict__ ln_b,
                              const float* __restrict__ qw,
                              const float* __restrict__ kw,
                              const float* __restrict__ vw,
                              float* __restrict__ q,
                              float* __restrict__ k,
                              float* __restrict__ v) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    __shared__ float mean_s, rstd_s;
    __shared__ float h[DIM];
    __shared__ float red[256];
    float sum = 0.0f;
    for (int d = tid; d < DIM; d += blockDim.x) sum += nodes[row*DIM + d];
    red[tid] = sum;
    __syncthreads();
    for (int stride = blockDim.x/2; stride > 0; stride >>= 1) { if (tid < stride) red[tid] += red[tid + stride]; __syncthreads(); }
    if (tid == 0) mean_s = red[0] * (1.0f/(float)DIM);
    __syncthreads();
    float mean = mean_s;
    float vsum = 0.0f;
    for (int d = tid; d < DIM; d += blockDim.x) { float x = nodes[row*DIM + d] - mean; vsum += x*x; }
    red[tid] = vsum;
    __syncthreads();
    for (int stride = blockDim.x/2; stride > 0; stride >>= 1) { if (tid < stride) red[tid] += red[tid + stride]; __syncthreads(); }
    if (tid == 0) rstd_s = rsqrtf(red[0] * (1.0f/(float)DIM) + 1.0e-5f);
    __syncthreads();
    float rstd = rstd_s;
    for (int d = tid; d < DIM; d += blockDim.x) h[d] = (nodes[row*DIM + d] - mean) * rstd * ln_w[d] + ln_b[d];
    __syncthreads();
    for (int t = tid; t < 3*DIM; t += blockDim.x) {
        int which = t / DIM;
        int o = t - which*DIM;
        const float* w = (which == 0 ? qw : (which == 1 ? kw : vw));
        float acc = 0.0f;
        #pragma unroll
        for (int d = 0; d < DIM; ++d) acc += h[d] * w[o*DIM + d];
        int head = o / HDIM;
        int dd = o - head*HDIM;
        int b = row / NOBJ;
        int n = row - b*NOBJ;
        int idx = ((b*NOBJ + n)*HEADS + head)*HDIM + dd;
        if (which == 0) q[idx] = acc; else if (which == 1) k[idx] = acc; else v[idx] = acc;
    }
}

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

__global__ void attn_rel_kernel(const float* __restrict__ q,
                                const float* __restrict__ k,
                                const float* __restrict__ v,
                                const float* __restrict__ pair_feats,
                                const float* __restrict__ edge_mask,
                                const float* __restrict__ r1w,
                                const float* __restrict__ r1b,
                                const float* __restrict__ r2w,
                                const float* __restrict__ r2b,
                                float* __restrict__ ctx) {
    int bid = blockIdx.x;
    int i = bid % NOBJ;
    int tmp = bid / NOBJ;
    int hidx = tmp % HEADS;
    int b = tmp / HEADS;
    int tid = threadIdx.x;
    __shared__ float logits[NOBJ];
    __shared__ float probs[NOBJ];
    __shared__ float sum_s;
    if (tid < NOBJ) {
        int j = tid;
        float logit = -INFINITY;
        if (edge_mask[(b*NOBJ + i)*NOBJ + j] >= 0.5f) {
            float dot = 0.0f;
            const float* qptr = &q[((b*NOBJ + i)*HEADS + hidx)*HDIM];
            const float* kptr = &k[((b*NOBJ + j)*HEADS + hidx)*HDIM];
            #pragma unroll
            for (int d = 0; d < HDIM; ++d) dot += qptr[d] * kptr[d];
            float rb = r2b[hidx];
            const float* pptr = &pair_feats[((b*NOBJ + i)*NOBJ + j)*PDIM];
            #pragma unroll
            for (int m = 0; m < PDIM; ++m) {
                float hv = r1b[m];
                #pragma unroll
                for (int p = 0; p < PDIM; ++p) hv += pptr[p] * r1w[m*PDIM + p];
                hv = silu_f(hv);
                rb += hv * r2w[hidx*PDIM + m];
            }
            logit = dot * 0.2041241452319315f + rb;
        }
        logits[j] = logit;
    }
    __syncthreads();
    if (tid == 0) {
        float m = logits[0];
        #pragma unroll
        for (int j = 1; j < NOBJ; ++j) if (logits[j] > m) m = logits[j];
        float s = 0.0f;
        #pragma unroll
        for (int j = 0; j < NOBJ; ++j) { float e = expf(logits[j] - m); probs[j] = e; s += e; }
        sum_s = s;
    }
    __syncthreads();
    float invs = 1.0f / sum_s;
    if (tid < HDIM) {
        int d = tid;
        float acc = 0.0f;
        #pragma unroll
        for (int j = 0; j < NOBJ; ++j) acc += (probs[j] * invs) * v[((b*NOBJ + j)*HEADS + hidx)*HDIM + d];
        ctx[(b*NOBJ + i)*DIM + hidx*HDIM + d] = acc;
    }
}

__global__ void out_res_kernel(const float* __restrict__ nodes,
                               const float* __restrict__ ctx,
                               const float* __restrict__ outw,
                               float* __restrict__ out) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    for (int o = tid; o < DIM; o += blockDim.x) {
        float acc = 0.0f;
        #pragma unroll
        for (int d = 0; d < DIM; ++d) acc += ctx[row*DIM + d] * outw[o*DIM + d];
        out[row*DIM + o] = nodes[row*DIM + o] + acc;
    }
}

torch::Tensor gf_vla_forward_cuda(torch::Tensor nodes, torch::Tensor pair_feats, torch::Tensor edge_mask,
                                  torch::Tensor ln_w, torch::Tensor ln_b,
                                  torch::Tensor q_w, torch::Tensor k_w, torch::Tensor v_w,
                                  torch::Tensor r1_w, torch::Tensor r1_b, torch::Tensor r2_w, torch::Tensor r2_b,
                                  torch::Tensor out_w) {
    auto opts = nodes.options();
    auto q = torch::empty({BATCH, NOBJ, HEADS, HDIM}, opts);
    auto k = torch::empty({BATCH, NOBJ, HEADS, HDIM}, opts);
    auto v = torch::empty({BATCH, NOBJ, HEADS, HDIM}, opts);
    auto ctx = torch::empty({BATCH, NOBJ, DIM}, opts);
    auto out = torch::empty_like(nodes);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    ln_qkv_kernel<<<BATCH*NOBJ, 256, 0, stream>>>(nodes.data_ptr<float>(), ln_w.data_ptr<float>(), ln_b.data_ptr<float>(), q_w.data_ptr<float>(), k_w.data_ptr<float>(), v_w.data_ptr<float>(), q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>());
    attn_rel_kernel<<<BATCH*HEADS*NOBJ, 64, 0, stream>>>(q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), pair_feats.data_ptr<float>(), edge_mask.data_ptr<float>(), r1_w.data_ptr<float>(), r1_b.data_ptr<float>(), r2_w.data_ptr<float>(), r2_b.data_ptr<float>(), ctx.data_ptr<float>());
    out_res_kernel<<<BATCH*NOBJ, 128, 0, stream>>>(nodes.data_ptr<float>(), ctx.data_ptr<float>(), out_w.data_ptr<float>(), out.data_ptr<float>());
    return out;
}
'''

_ext = load_inline(name='gf_vla_ext_v2', cpp_sources=cpp_source, cuda_sources=cuda_source,
                   functions=['gf_vla_forward_cuda'], verbose=False,
                   extra_cuda_cflags=['-O3'], extra_cflags=['-O3'])

class ModelNew(nn.Module):
    def __init__(self, dim: int, pair_dim: int, num_heads: int):
        super().__init__()
        assert dim == 96 and pair_dim == 16 and num_heads == 4
        self.dim = dim; self.pair_dim = pair_dim; self.num_heads = num_heads; self.head_dim = dim // num_heads
        self.scale = 1.0 / math.sqrt(self.head_dim)
        self.norm = nn.LayerNorm(dim)
        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.rel_mlp = nn.Sequential(nn.Linear(pair_dim, pair_dim), nn.SiLU(), nn.Linear(pair_dim, num_heads))
        self.out_proj = nn.Linear(dim, dim, bias=False)
    def forward(self, nodes: torch.Tensor, pair_feats: torch.Tensor, edge_mask: torch.Tensor) -> torch.Tensor:
        return _ext.gf_vla_forward_cuda(nodes, pair_feats, edge_mask, self.norm.weight, self.norm.bias, self.q_proj.weight, self.k_proj.weight, self.v_proj.weight, self.rel_mlp[0].weight, self.rel_mlp[0].bias, self.rel_mlp[2].weight, self.rel_mlp[2].bias, self.out_proj.weight)