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_sources = r'''
#include <torch/extension.h>
torch::Tensor scatter_add_forward(torch::Tensor a, torch::Tensor b, int64_t num_nodes);
'''
_cuda_sources = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>

__global__ void zero_kernel(float* __restrict__ out, int n) { int i=blockIdx.x*blockDim.x+threadIdx.x; if(i<n) out[i]=0.0f; }

template <typename IdxT>
__global__ void node_scan64_reg_kernel(const IdxT* __restrict__ dst, const float* __restrict__ edge, float* __restrict__ out, int E) {
    int node = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warp = tid >> 5;
    float s0 = 0.0f, s1 = 0.0f;
    const unsigned FULL = 0xffffffffu;
    #pragma unroll 1
    for (int e = tid; e < E; e += 256) {
        bool hit = ((int)dst[e] == node);
        unsigned mask = __ballot_sync(FULL, hit);
        while (mask) {
            int src_lane = __ffs(mask) - 1;
            int me = __shfl_sync(FULL, e, src_lane);
            const float* row = edge + ((int64_t)me << 6);
            s0 += row[lane];
            s1 += row[lane + 32];
            mask &= mask - 1;
        }
    }
    __shared__ float partial[8 * 64];
    partial[(warp << 6) + lane] = s0;
    partial[(warp << 6) + lane + 32] = s1;
    __syncthreads();
    if (tid < 64) {
        float sum = partial[tid] + partial[64 + tid] + partial[128 + tid] + partial[192 + tid]
                  + partial[256 + tid] + partial[320 + tid] + partial[384 + tid] + partial[448 + tid];
        out[((int64_t)node << 6) + tid] = sum;
    }
}

template <typename IdxT>
__global__ void scatter64_vec4_kernel(const IdxT* __restrict__ dst, const float* __restrict__ edge, float* __restrict__ out, int E) {
    int idx=blockIdx.x*blockDim.x+threadIdx.x; int total=E<<4;
    if(idx<total){ int e=idx>>4; int g=idx&15; int d=(int)dst[e]; const float4 v=reinterpret_cast<const float4*>(edge+(e<<6))[g]; float* o=out+((int64_t)d<<6)+(g<<2); atomicAdd(o,v.x); atomicAdd(o+1,v.y); atomicAdd(o+2,v.z); atomicAdd(o+3,v.w); }
}
template <typename IdxT>
__global__ void scatter_generic_kernel(const IdxT* __restrict__ dst, const float* __restrict__ edge, float* __restrict__ out, int E, int D) { int idx=blockIdx.x*blockDim.x+threadIdx.x; int total=E*D; if(idx<total){ int e=idx/D; int f=idx-e*D; atomicAdd(out+(int64_t)((int)dst[e])*D+f, edge[idx]); } }

torch::Tensor scatter_add_forward(torch::Tensor a, torch::Tensor b, int64_t num_nodes) {
    torch::Tensor dst_idx=a, edge_feat=b;
    if(edge_feat.dim()<2 && dst_idx.dim()>=2){ dst_idx=b; edge_feat=a; }
    const at::cuda::OptionalCUDAGuard device_guard(device_of(edge_feat));
    int64_t E64=edge_feat.size(0), D64=edge_feat.size(1);
    auto out=torch::empty({num_nodes,D64}, edge_feat.options()); float* outp=out.data_ptr<float>();
    cudaStream_t stream=at::cuda::getCurrentCUDAStream(); int E=(int)E64;
    if(D64==64 && num_nodes <= 4096) {
        dim3 grid((unsigned)num_nodes); const int threads=256;
        if(dst_idx.scalar_type()==at::kInt) node_scan64_reg_kernel<int><<<grid,threads,0,stream>>>(dst_idx.data_ptr<int>(), edge_feat.data_ptr<float>(), outp, E);
        else if(dst_idx.scalar_type()==at::kLong) node_scan64_reg_kernel<int64_t><<<grid,threads,0,stream>>>(dst_idx.data_ptr<int64_t>(), edge_feat.data_ptr<float>(), outp, E);
        else node_scan64_reg_kernel<float><<<grid,threads,0,stream>>>(dst_idx.data_ptr<float>(), edge_feat.data_ptr<float>(), outp, E);
    } else {
        const int threads=256; int nout=(int)(num_nodes*D64); zero_kernel<<<(nout+threads-1)/threads,threads,0,stream>>>(outp,nout);
        int total=(D64==64)?(E<<4):(int)(E64*D64); int blocks=(total+threads-1)/threads;
        if(D64==64){
            if(dst_idx.scalar_type()==at::kInt) scatter64_vec4_kernel<int><<<blocks,threads,0,stream>>>(dst_idx.data_ptr<int>(), edge_feat.data_ptr<float>(), outp, E);
            else if(dst_idx.scalar_type()==at::kLong) scatter64_vec4_kernel<int64_t><<<blocks,threads,0,stream>>>(dst_idx.data_ptr<int64_t>(), edge_feat.data_ptr<float>(), outp, E);
            else scatter64_vec4_kernel<float><<<blocks,threads,0,stream>>>(dst_idx.data_ptr<float>(), edge_feat.data_ptr<float>(), outp, E);
        } else {
            if(dst_idx.scalar_type()==at::kInt) scatter_generic_kernel<int><<<blocks,threads,0,stream>>>(dst_idx.data_ptr<int>(), edge_feat.data_ptr<float>(), outp, E, (int)D64);
            else if(dst_idx.scalar_type()==at::kLong) scatter_generic_kernel<int64_t><<<blocks,threads,0,stream>>>(dst_idx.data_ptr<int64_t>(), edge_feat.data_ptr<float>(), outp, E, (int)D64);
            else scatter_generic_kernel<float><<<blocks,threads,0,stream>>>(dst_idx.data_ptr<float>(), edge_feat.data_ptr<float>(), outp, E, (int)D64);
        }
    }
    return out;
}
'''
_ext=load_inline(name='scatter_add_ext_v6', cpp_sources=_cpp_sources, cuda_sources=_cuda_sources, functions=['scatter_add_forward'], extra_cuda_cflags=['-O3'], extra_cflags=['-O3'], with_cuda=True, verbose=False)
class ModelNew(nn.Module):
    def __init__(self, num_nodes): super().__init__(); self.num_nodes=int(num_nodes)
    def forward(self, dst_idx: torch.Tensor, edge_feat: torch.Tensor) -> torch.Tensor:
        if edge_feat.dim()<2 and dst_idx.dim()>=2: dst_idx, edge_feat = edge_feat, dst_idx
        return _ext.scatter_add_forward(dst_idx, edge_feat, self.num_nodes)