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 = r"""
#include <cuda.h>
#include <cuda_runtime.h>
#include <cstdio>

// Optimized scatter-add kernel with vectorized loads
__global__ __launch_bounds__(128, 16)
void scatter_add_kernel_opt(
    const int* __restrict__ dst_idx,
    const float* __restrict__ edge_feat,
    float* __restrict__ out,
    const int num_edges,
    const int feat_dim
) {
    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
    const int stride = blockDim.x * gridDim.x;
    
    for (int e = tid; e < num_edges; e += stride) {
        const int dst = dst_idx[e];
        const float* feat_ptr = edge_feat + e * feat_dim;
        float* out_ptr = out + dst * feat_dim;
        
        const int num_vec4 = feat_dim >> 2;
        for (int d = 0; d < num_vec4; d++) {
            const float4 val = reinterpret_cast<const float4*>(feat_ptr)[d];
            float4* out_vec = reinterpret_cast<float4*>(out_ptr + d * 4);
            atomicAdd(&out_vec->x, val.x);
            atomicAdd(&out_vec->y, val.y);
            atomicAdd(&out_vec->z, val.z);
            atomicAdd(&out_vec->w, val.w);
        }
        
        for (int d = num_vec4 * 4; d < feat_dim; d++) {
            atomicAdd(&out_ptr[d], feat_ptr[d]);
        }
    }
}

void scatter_add_cuda(
    const at::Tensor& dst_idx,
    const at::Tensor& edge_feat,
    at::Tensor& out
) {
    const int num_edges = edge_feat.size(0);
    const int feat_dim = edge_feat.size(1);
    
    int block_size = 128;
    int num_blocks = (num_edges + block_size - 1) / block_size;
    num_blocks = std::min(num_blocks, 65535);
    
    scatter_add_kernel_opt<<<num_blocks, block_size>>>(
        dst_idx.data_ptr<int>(),
        edge_feat.data_ptr<float>(),
        out.data_ptr<float>(),
        num_edges,
        feat_dim
    );
}
"""

cpp_source = r"""
void scatter_add_cuda(const at::Tensor& dst_idx, const at::Tensor& edge_feat, at::Tensor& out);
"""

scatter_ops = load_inline(
    name='scatter_add_ops',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['scatter_add_cuda'],
    extra_cuda_cflags=['-O3', '-arch=sm_90', '-use_fast_math'],
    verbose=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:
        feat_dim = edge_feat.shape[1]
        out = torch.zeros(self.num_nodes, feat_dim, dtype=edge_feat.dtype, device=edge_feat.device)
        scatter_ops.scatter_add_cuda(dst_idx.to(torch.int32), edge_feat, out)
        return out