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

# Define CUDA source with vectorized operations
cuda_source = """
#include <cuda.h>
#include <cuda_runtime.h>
#include <cstdio>
#include <ATen/ATen.h>

// Vectorized scatter-add kernel using float4 for better memory throughput
__global__ void scatter_add_kernel_vec4(
    const int* __restrict__ dst_idx,
    const float* __restrict__ edge_feat,
    float* __restrict__ output,
    int num_edges,
    int feat_dim_vec4,
    int num_nodes
) {
    int edge = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (edge < num_edges) {
        int dst = dst_idx[edge];
        const float4* feat_ptr = reinterpret_cast<const float4*>(edge_feat + edge * feat_dim_vec4 * 4);
        float4* out_ptr = reinterpret_cast<float4*>(output + dst * feat_dim_vec4 * 4);
        
        #pragma unroll 4
        for (int d = 0; d < feat_dim_vec4; d++) {
            float4 val = feat_ptr[d];
            atomicAdd(&out_ptr[d].x, val.x);
            atomicAdd(&out_ptr[d].y, val.y);
            atomicAdd(&out_ptr[d].z, val.z);
            atomicAdd(&out_ptr[d].w, val.w);
        }
    }
}

__global__ void scatter_add_kernel(
    const int* __restrict__ dst_idx,
    const float* __restrict__ edge_feat,
    float* __restrict__ output,
    int num_edges,
    int feat_dim,
    int num_nodes
) {
    int edge = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (edge < num_edges) {
        int dst = dst_idx[edge];
        const float* feat_ptr = edge_feat + edge * feat_dim;
        float* out_ptr = output + dst * feat_dim;
        
        for (int d = 0; d < feat_dim; d++) {
            atomicAdd(&out_ptr[d], feat_ptr[d]);
        }
    }
}

void scatter_add_cuda(
    at::Tensor dst_idx,
    at::Tensor edge_feat,
    at::Tensor output,
    int64_t num_edges,
    int64_t feat_dim,
    int64_t num_nodes
) {
    at::Tensor dst_idx_int = dst_idx.to(at::ScalarType::Int);
    const int* dst_ptr = dst_idx_int.data_ptr<int>();
    const float* feat_ptr = edge_feat.data_ptr<float>();
    float* out_ptr = output.data_ptr<float>();
    
    int block_size = 512;
    int num_blocks = (num_edges + block_size - 1) / block_size;
    
    if (feat_dim % 4 == 0) {
        scatter_add_kernel_vec4<<<num_blocks, block_size>>>(
            dst_ptr, feat_ptr, out_ptr, num_edges, feat_dim / 4, num_nodes
        );
    } else {
        scatter_add_kernel<<<num_blocks, block_size>>>(
            dst_ptr, feat_ptr, out_ptr, num_edges, feat_dim, num_nodes
        );
    }
}
"""

scatter_add_ext = load_inline(
    name='scatter_add_ext',
    cpp_sources="void scatter_add_cuda(at::Tensor dst_idx, at::Tensor edge_feat, at::Tensor output, int64_t num_edges, int64_t feat_dim, int64_t num_nodes);",
    cuda_sources=cuda_source,
    functions=['scatter_add_cuda'],
    with_cuda=True,
    extra_cuda_cflags=['-O3'],
)


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_add_ext.scatter_add_cuda(dst_idx, edge_feat, out, edge_feat.shape[0], feat_dim, self.num_nodes)
        return out