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

__global__ void scatter_add_kernel(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 float* src = edge_feat + edge * feat_dim;
        float* dst_ptr = out + dst * feat_dim;
        for (int f = 0; f < feat_dim; ++f) {
            atomicAdd(dst_ptr + f, src[f]);
        }
    }
}

torch::Tensor scatter_add_cuda(torch::Tensor dst_idx, torch::Tensor edge_feat, int num_nodes) {
    auto dst_idx_long = dst_idx.to(torch::kInt64);
    int num_edges = dst_idx_long.size(0);
    int feat_dim = edge_feat.size(1);
    auto out = torch::zeros({num_nodes, feat_dim}, edge_feat.options());
    const int threads = 256;
    const int blocks = (num_edges + threads - 1) / threads;
    scatter_add_kernel<<<blocks, threads>>>(
        dst_idx_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);
"""

scatter_add = load_inline(
    name="scatter_add",
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=["scatter_add_cuda"],
    verbose=True,
    with_cuda=True
)

class ModelNew(nn.Module):
    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:
        return scatter_add.scatter_add_cuda(dst_idx, edge_feat, self.num_nodes)