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


_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor scatter_add_forward(torch::Tensor dst_idx, torch::Tensor edge_feat, int64_t num_nodes);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

// One CUDA thread processes four feature values.  The benchmark shape has
// feat_dim == 64, so this is 16 vector lanes per edge.  This keeps global reads
// vectorized/coalesced while avoiding the overhead of launching one thread per
// scalar element.
__global__ void scatter_add64_float4_kernel(const int* __restrict__ dst,
                                            const float4* __restrict__ x4,
                                            float* __restrict__ out,
                                            int64_t n_vec4) {
    int64_t tid = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (int64_t k = tid; k < n_vec4; k += stride) {
        int64_t e = k >> 4;          // 16 float4's per edge (64 floats)
        int f4 = (int)(k & 15);
        int d = dst[e];
        float4 v = x4[k];
        float* o = out + ((int64_t)d << 6) + (f4 << 2);
        atomicAdd(o + 0, v.x);
        atomicAdd(o + 1, v.y);
        atomicAdd(o + 2, v.z);
        atomicAdd(o + 3, v.w);
    }
}

__global__ void scatter_add_generic_kernel(const int* __restrict__ dst,
                                           const float* __restrict__ x,
                                           float* __restrict__ out,
                                           int64_t total,
                                           int feat_dim) {
    int64_t tid = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (int64_t k = tid; k < total; k += stride) {
        int64_t e = k / feat_dim;
        int f = (int)(k - e * feat_dim);
        int d = dst[e];
        atomicAdd(out + (int64_t)d * feat_dim + f, x[k]);
    }
}

torch::Tensor scatter_add_forward(torch::Tensor dst_idx, torch::Tensor edge_feat, int64_t num_nodes) {
    const auto E = edge_feat.size(0);
    const auto F = edge_feat.size(1);
    auto out = torch::empty({num_nodes, F}, edge_feat.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cudaMemsetAsync(out.data_ptr<float>(), 0, (size_t)num_nodes * (size_t)F * sizeof(float), stream);

    const int threads = 256;
    if (F == 64 && edge_feat.is_contiguous() && dst_idx.is_contiguous()) {
        int64_t n_vec4 = E * 16; // E * F / 4
        int blocks = (int)((n_vec4 + threads - 1) / threads);
        if (blocks > 4096) blocks = 4096;
        scatter_add64_float4_kernel<<<blocks, threads, 0, stream>>>(
            (const int*)dst_idx.data_ptr<int>(),
            (const float4*)edge_feat.data_ptr<float>(),
            out.data_ptr<float>(),
            n_vec4);
    } else {
        int64_t total = E * F;
        int blocks = (int)((total + threads - 1) / threads);
        if (blocks > 4096) blocks = 4096;
        scatter_add_generic_kernel<<<blocks, threads, 0, stream>>>(
            (const int*)dst_idx.data_ptr<int>(),
            edge_feat.data_ptr<float>(),
            out.data_ptr<float>(),
            total,
            (int)F);
    }
    return out;
}
"""

_ext = load_inline(
    name="scatter_add_coo_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["scatter_add_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)


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 _ext.scatter_add_forward(dst_idx, edge_feat, self.num_nodes)