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>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime_api.h>
extern "C" void scatter_add64_i32_launch(const int* dst, const float* edge, float* out, int E, cudaStream_t stream);
extern "C" void scatter_add64_i64_launch(const long long* dst, const float* edge, float* out, int E, cudaStream_t stream);
extern "C" void scatter_add64_f32_launch(const float* dst, const float* edge, float* out, int E, cudaStream_t stream);

torch::Tensor scatter_add_forward(torch::Tensor dst_idx, torch::Tensor edge_feat, int64_t num_nodes) {
    auto out = torch::empty({num_nodes, edge_feat.size(1)}, edge_feat.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    int N = (int)num_nodes;
    int E = (int)edge_feat.size(0);
    cudaMemsetAsync(out.data_ptr<float>(), 0, (size_t)N * 64 * sizeof(float), stream);
    const float* edge = edge_feat.data_ptr<float>();
    float* op = out.data_ptr<float>();
    if (dst_idx.scalar_type() == at::kInt) {
        scatter_add64_i32_launch(dst_idx.data_ptr<int>(), edge, op, E, stream);
    } else if (dst_idx.scalar_type() == at::kLong) {
        scatter_add64_i64_launch((const long long*)dst_idx.data_ptr<int64_t>(), edge, op, E, stream);
    } else {
        scatter_add64_f32_launch(dst_idx.data_ptr<float>(), edge, op, E, stream);
    }
    return out;
}
'''

cuda_src = r'''
#include <cuda_runtime.h>

template <typename IDX>
__device__ __forceinline__ int idx_to_int(IDX x) { return (int)x; }

template <typename IDX>
__global__ void scatter_vec4_atomic_kernel(const IDX* __restrict__ dst, const float4* __restrict__ edge4, float4* __restrict__ out4, int total4) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < total4) {
        int e = tid >> 4;
        int g = tid & 15;
        int d = idx_to_int<IDX>(dst[e]);
        float4 v = edge4[tid];
#if __CUDA_ARCH__ >= 900
        atomicAdd(out4 + ((d << 4) + g), v);
#else
        float* o = (float*)(out4 + ((d << 4) + g));
        atomicAdd(o + 0, v.x);
        atomicAdd(o + 1, v.y);
        atomicAdd(o + 2, v.z);
        atomicAdd(o + 3, v.w);
#endif
    }
}

template <typename IDX>
static inline void launch_t(const IDX* dst, const float* edge, float* out, int E, cudaStream_t stream) {
    int total4 = E << 4;
    int threads = 256;
    int blocks = (total4 + threads - 1) / threads;
    scatter_vec4_atomic_kernel<IDX><<<blocks, threads, 0, stream>>>(dst, reinterpret_cast<const float4*>(edge), reinterpret_cast<float4*>(out), total4);
}

extern "C" void scatter_add64_i32_launch(const int* dst, const float* edge, float* out, int E, cudaStream_t stream) { launch_t<int>(dst, edge, out, E, stream); }
extern "C" void scatter_add64_i64_launch(const long long* dst, const float* edge, float* out, int E, cudaStream_t stream) { launch_t<long long>(dst, edge, out, E, stream); }
extern "C" void scatter_add64_f32_launch(const float* dst, const float* edge, float* out, int E, cudaStream_t stream) { launch_t<float>(dst, edge, out, E, stream); }
'''

_ext = load_inline(name='scatter_add_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['scatter_add_forward'], extra_cuda_cflags=['-O3'], extra_cflags=['-O3'], 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)