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

cuda_source = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void edge_dot_product_kernel(
    const int32_t* __restrict__ src_idx,
    const int32_t* __restrict__ dst_idx,
    const float* __restrict__ src_feat,
    const float* __restrict__ dst_feat,
    float* __restrict__ output,
    int num_edges,
    int feat_dim
) {
    int edge = blockIdx.x * blockDim.x + threadIdx.x;
    if (edge >= num_edges) return;

    int32_t src = src_idx[edge];
    int32_t dst = dst_idx[edge];

    float sum = 0.0f;
    const float* src_ptr = src_feat + (int64_t)src * feat_dim;
    const float* dst_ptr = dst_feat + (int64_t)dst * feat_dim;

    for (int f = 0; f < feat_dim; ++f) {
        sum += src_ptr[f] * dst_ptr[f];
    }

    output[edge] = sum;
}

torch::Tensor edge_dot_product_cuda(
    torch::Tensor src_idx,
    torch::Tensor dst_idx,
    torch::Tensor src_feat,
    torch::Tensor dst_feat
) {
    TORCH_CHECK(src_idx.is_cuda(), "src_idx must be CUDA");
    TORCH_CHECK(dst_idx.is_cuda(), "dst_idx must be CUDA");
    TORCH_CHECK(src_feat.is_cuda(), "src_feat must be CUDA");
    TORCH_CHECK(dst_feat.is_cuda(), "dst_feat must be CUDA");

    int num_edges = src_idx.size(0);
    int feat_dim = src_feat.size(1);

    auto output = torch::empty({num_edges}, src_feat.options());

    const int threads = 256;
    const int blocks = (num_edges + threads - 1) / threads;

    edge_dot_product_kernel<<<blocks, threads>>>(
        src_idx.data_ptr<int32_t>(),
        dst_idx.data_ptr<int32_t>(),
        src_feat.data_ptr<float>(),
        dst_feat.data_ptr<float>(),
        output.data_ptr<float>(),
        num_edges,
        feat_dim
    );

    return output;
}
'''

cpp_source = '''
torch::Tensor edge_dot_product_cuda(
    torch::Tensor src_idx,
    torch::Tensor dst_idx,
    torch::Tensor src_feat,
    torch::Tensor dst_feat
);
'''

module = load_inline(
    name='edge_dot_product',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['edge_dot_product_cuda'],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.module = module

    def forward(
        self,
        src_idx: torch.Tensor,
        dst_idx: torch.Tensor,
        src_feat: torch.Tensor,
        dst_feat: torch.Tensor,
    ) -> torch.Tensor:
        src_idx = src_idx.contiguous().to(torch.int32);
        dst_idx = dst_idx.contiguous().to(torch.int32);
        return self.module.edge_dot_product_cuda(src_idx, dst_idx, src_feat, dst_feat);
```