import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Scatter-add edge features into destination-node accumulators.

    This is a canonical graph primitive for COO-formatted workloads where edge
    messages are already materialized and must be reduced into node states.
    """

    def __init__(self, num_nodes):
        super().__init__()
        self.num_nodes = num_nodes

    def forward(
        self,
        dst_idx: torch.Tensor,
        edge_feat: torch.Tensor,
    ) -> torch.Tensor:
        feat_dim = edge_feat.shape[1]
        out = torch.zeros(self.num_nodes, feat_dim, dtype=edge_feat.dtype, device=edge_feat.device)
        out.index_add_(0, dst_idx.long(), edge_feat)
        return out


num_nodes = 768
num_edges = 16384
feat_dim = 64


def get_inputs():
    p = popcorn_pri
    n_edges = p.jitter_int(num_edges)
    fd = p.jitter_int(feat_dim, align=8)
    dst_idx = torch.randint(0, num_nodes, (n_edges,), dtype=torch.int32)
    edge_feat = torch.randn(n_edges, fd, dtype=torch.float32)
    return [dst_idx, edge_feat]





def get_init_inputs():
    return [num_nodes]

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

cuda_source = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cstdint>

__global__ void scatter_add_kernel_vec4(
    const int64_t* __restrict__ dst_idx,
    const float* __restrict__ edge_feat,
    float* __restrict__ out,
    int num_edges,
    int feat_dim
) {
    int edge = blockIdx.x * blockDim.x + threadIdx.x;
    if (edge < num_edges) {
        int64_t dst = dst_idx[edge];
        const float4* feat4 = reinterpret_cast<const float4*>(edge_feat + edge * feat_dim);
        float4* out_feat4 = reinterpret_cast<float4*>(out + dst * feat_dim);
        int vecs = feat_dim / 4;
        for (int v = 0; v < vecs; ++v) {
            float4 val = feat4[v];
            atomicAdd(&out_feat4[v].x, val.x);
            atomicAdd(&out_feat4[v].y, val.y);
            atomicAdd(&out_feat4[v].z, val.z);
            atomicAdd(&out_feat4[v].w, val.w);
        }
    }
}

torch::Tensor scatter_add_cuda(torch::Tensor dst_idx, torch::Tensor edge_feat, int num_nodes) {
    int num_edges = dst_idx.size(0);
    int feat_dim = edge_feat.size(1);
    auto out = torch::zeros({num_nodes, feat_dim}, edge_feat.options());
    auto dst_long = dst_idx.to(torch::kInt64).contiguous();
    dim3 block(256);
    dim3 grid((num_edges + 255) / 256);
    scatter_add_kernel_vec4<<<grid, block>>>(
        dst_long.data_ptr<int64_t>(),
        edge_feat.data_ptr<float>(),
        out.data_ptr<float>(),
        num_edges,
        feat_dim
    );
    return out;
}
'''

cpp_source = '''
torch::Tensor scatter_add_cuda(torch::Tensor dst_idx, torch::Tensor edge_feat, int num_nodes);
'''

module = load_inline(
    name='scatter_add_vec4',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['scatter_add_cuda'],
    verbose=True
)

class ModelNew(nn.Module):
    def __init__(self, num_nodes):
        super().__init__()
        self.num_nodes = num_nodes
        self.module = module

    def forward(self, dst_idx, edge_feat):
        return self.module.scatter_add_cuda(dst_idx, edge_feat, self.num_nodes)