import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Edge-only sampled dense-dense matmul.

    For each graph edge (src, dst), compute a dot product between source and
    destination node embeddings. This pattern appears in link prediction,
    sampled attention, and graph contrastive learning.
    """

    def __init__(self):
        super().__init__()

    def forward(
        self,
        src_idx: torch.Tensor,
        dst_idx: torch.Tensor,
        src_feat: torch.Tensor,
        dst_feat: torch.Tensor,
    ) -> torch.Tensor:
        src = src_feat[src_idx.long()]
        dst = dst_feat[dst_idx.long()]
        return (src * dst).sum(dim=-1)


num_src_nodes = 1024
num_dst_nodes = 768
num_edges = 16384
feat_dim = 128


def get_inputs():
    p = popcorn_pri
    n_edges = p.jitter_int(num_edges)
    fd = p.jitter_int(feat_dim, align=8)
    ns = p.jitter_int(num_src_nodes)
    nd = p.jitter_int(num_dst_nodes)
    src_idx = torch.randint(0, ns, (n_edges,), dtype=torch.int32)
    dst_idx = torch.randint(0, nd, (n_edges,), dtype=torch.int32)
    src_feat = torch.randn(ns, fd, dtype=torch.float32)
    dst_feat = torch.randn(nd, fd, dtype=torch.float32)
    return [src_idx, dst_idx, src_feat, dst_feat]





def get_init_inputs():
    return []

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_src = r"""
#include <torch/extension.h>
torch::Tensor edge_dot_forward(torch::Tensor src_idx, torch::Tensor dst_idx, torch::Tensor src_feat, torch::Tensor dst_feat);
"""

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

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

__inline__ __device__ float group8_reduce_sum(float v, unsigned mask) {
    v += __shfl_down_sync(mask, v, 4, 8);
    v += __shfl_down_sync(mask, v, 2, 8);
    v += __shfl_down_sync(mask, v, 1, 8);
    return v;
}

template <typename index_t, int GROUPS_PER_BLOCK>
__global__ __launch_bounds__(GROUPS_PER_BLOCK * 8, 4)
void edge_dot128_group8_kernel(const index_t* __restrict__ src_idx,
                               const index_t* __restrict__ dst_idx,
                               const float* __restrict__ src_feat,
                               const float* __restrict__ dst_feat,
                               float* __restrict__ out,
                               int64_t E) {
    int tid = threadIdx.x;
    int lane = tid & 31;
    int glane = tid & 7;
    int g_in_warp = lane >> 3;
    int64_t e = (int64_t)blockIdx.x * GROUPS_PER_BLOCK + (tid >> 3);
    if (e >= E) return;
    unsigned mask = 0xffu << (g_in_warp * 8);

    int s = 0, d = 0;
    if (glane == 0) {
        s = idx_to_int(src_idx[e]);
        d = idx_to_int(dst_idx[e]);
    }
    s = __shfl_sync(mask, s, 0, 8);
    d = __shfl_sync(mask, d, 0, 8);

    const float4* __restrict__ sp = reinterpret_cast<const float4*>(src_feat + ((int64_t)s << 7));
    const float4* __restrict__ dp = reinterpret_cast<const float4*>(dst_feat + ((int64_t)d << 7));

    float sum = 0.0f;
    float4 a, b;
    a = sp[glane];      b = dp[glane];      sum += a.x*b.x + a.y*b.y + a.z*b.z + a.w*b.w;
    a = sp[glane + 8];  b = dp[glane + 8];  sum += a.x*b.x + a.y*b.y + a.z*b.z + a.w*b.w;
    a = sp[glane + 16]; b = dp[glane + 16]; sum += a.x*b.x + a.y*b.y + a.z*b.z + a.w*b.w;
    a = sp[glane + 24]; b = dp[glane + 24]; sum += a.x*b.x + a.y*b.y + a.z*b.z + a.w*b.w;
    sum = group8_reduce_sum(sum, mask);
    if (glane == 0) out[e] = sum;
}

template <typename index_t, int WARPS_PER_BLOCK>
__global__ __launch_bounds__(WARPS_PER_BLOCK * 32, 4)
void edge_dot_generic_kernel(const index_t* __restrict__ src_idx,
                             const index_t* __restrict__ dst_idx,
                             const float* __restrict__ src_feat,
                             const float* __restrict__ dst_feat,
                             float* __restrict__ out,
                             int64_t E, int64_t K) {
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    int64_t e = (int64_t)blockIdx.x * WARPS_PER_BLOCK + warp;
    if (e >= E) return;
    int64_t s = (int64_t)idx_to_int(src_idx[e]);
    int64_t d = (int64_t)idx_to_int(dst_idx[e]);
    const float* __restrict__ sp = src_feat + s * K;
    const float* __restrict__ dp = dst_feat + d * K;
    float sum = 0.0f;
    for (int64_t k = lane; k < K; k += 32) sum += sp[k] * dp[k];
    sum += __shfl_down_sync(0xffffffff, sum, 16);
    sum += __shfl_down_sync(0xffffffff, sum, 8);
    sum += __shfl_down_sync(0xffffffff, sum, 4);
    sum += __shfl_down_sync(0xffffffff, sum, 2);
    sum += __shfl_down_sync(0xffffffff, sum, 1);
    if (lane == 0) out[e] = sum;
}

template <typename index_t>
void launch_edge_dot(torch::Tensor src_idx, torch::Tensor dst_idx, torch::Tensor src_feat, torch::Tensor dst_feat, torch::Tensor out, int64_t E, int64_t K, cudaStream_t stream) {
    if (K == 128) {
        constexpr int GPB = 32;
        dim3 block(GPB * 8);
        dim3 grid((E + GPB - 1) / GPB);
        edge_dot128_group8_kernel<index_t, GPB><<<grid, block, 0, stream>>>(src_idx.data_ptr<index_t>(), dst_idx.data_ptr<index_t>(), src_feat.data_ptr<float>(), dst_feat.data_ptr<float>(), out.data_ptr<float>(), E);
    } else {
        constexpr int WPB = 8;
        dim3 block(WPB * 32);
        dim3 grid((E + WPB - 1) / WPB);
        edge_dot_generic_kernel<index_t, WPB><<<grid, block, 0, stream>>>(src_idx.data_ptr<index_t>(), dst_idx.data_ptr<index_t>(), src_feat.data_ptr<float>(), dst_feat.data_ptr<float>(), out.data_ptr<float>(), E, K);
    }
}

torch::Tensor edge_dot_forward(torch::Tensor src_idx, torch::Tensor dst_idx, torch::Tensor src_feat, torch::Tensor dst_feat) {
    int64_t E = src_idx.numel();
    int64_t K = src_feat.size(1);
    auto out = torch::empty({E}, src_feat.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (src_idx.scalar_type() == at::kInt) launch_edge_dot<int>(src_idx, dst_idx, src_feat, dst_feat, out, E, K, stream);
    else if (src_idx.scalar_type() == at::kLong) launch_edge_dot<int64_t>(src_idx, dst_idx, src_feat, dst_feat, out, E, K, stream);
    else launch_edge_dot<float>(src_idx, dst_idx, src_feat, dst_feat, out, E, K, stream);
    return out;
}
"""

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

class ModelNew(nn.Module):
    def __init__(self): super().__init__()
    def forward(self, src_idx: torch.Tensor, dst_idx: torch.Tensor, src_feat: torch.Tensor, dst_feat: torch.Tensor) -> torch.Tensor:
        return _ext.edge_dot_forward(src_idx, dst_idx, src_feat, dst_feat)