reference PyTorch source (the prompt the agent sees)
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 []
final submitted kernel (152 lines, 5,649 chars) — full source, untruncated
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 <c10/cuda/CUDAException.h>
#include <cstdint>
template <typename index_t>
__global__ __launch_bounds__(256, 2)
void edge_dot128_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,
int E) {
constexpr int WARPS_PER_BLOCK = 8;
int tid = threadIdx.x;
int lane = tid & 31;
int warp = tid >> 5;
int edge = blockIdx.x * WARPS_PER_BLOCK + warp;
if (edge >= E) return;
int s = static_cast<int>(src_idx[edge]);
int d = static_cast<int>(dst_idx[edge]);
const float* __restrict__ sp = src_feat + (static_cast<long long>(s) << 7); // *128
const float* __restrict__ dp = dst_feat + (static_cast<long long>(d) << 7);
float sum = 0.0f;
float a0 = sp[lane]; float b0 = dp[lane];
float a1 = sp[lane + 32]; float b1 = dp[lane + 32];
float a2 = sp[lane + 64]; float b2 = dp[lane + 64];
float a3 = sp[lane + 96]; float b3 = dp[lane + 96];
sum = fmaf(a0, b0, sum);
sum = fmaf(a1, b1, sum);
sum = fmaf(a2, b2, sum);
sum = fmaf(a3, b3, sum);
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[edge] = sum;
}
template <typename index_t>
__global__ __launch_bounds__(256, 2)
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,
int E,
int D) {
constexpr int WARPS_PER_BLOCK = 8;
int tid = threadIdx.x;
int lane = tid & 31;
int warp = tid >> 5;
int edge = blockIdx.x * WARPS_PER_BLOCK + warp;
if (edge >= E) return;
int s = static_cast<int>(src_idx[edge]);
int d = static_cast<int>(dst_idx[edge]);
const float* __restrict__ sp = src_feat + static_cast<long long>(s) * D;
const float* __restrict__ dp = dst_feat + static_cast<long long>(d) * D;
float sum = 0.0f;
for (int k = lane; k < D; k += 32) {
sum = fmaf(sp[k], dp[k], sum);
}
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[edge] = sum;
}
torch::Tensor edge_dot_forward(torch::Tensor src_idx,
torch::Tensor dst_idx,
torch::Tensor src_feat,
torch::Tensor dst_feat) {
const int E = static_cast<int>(src_idx.numel());
const int D = static_cast<int>(src_feat.size(1));
auto out = torch::empty({E}, src_feat.options());
if (E == 0) return out;
constexpr int THREADS = 256;
constexpr int WARPS_PER_BLOCK = THREADS / 32;
dim3 block(THREADS);
dim3 grid((E + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK);
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
if (src_idx.scalar_type() == at::kInt) {
const int* sidx = src_idx.data_ptr<int>();
const int* didx = dst_idx.data_ptr<int>();
if (D == 128) {
edge_dot128_kernel<int><<<grid, block, 0, stream>>>(sidx, didx,
src_feat.data_ptr<float>(), dst_feat.data_ptr<float>(), out.data_ptr<float>(), E);
} else {
edge_dot_generic_kernel<int><<<grid, block, 0, stream>>>(sidx, didx,
src_feat.data_ptr<float>(), dst_feat.data_ptr<float>(), out.data_ptr<float>(), E, D);
}
} else { // int64 indices
const int64_t* sidx = src_idx.data_ptr<int64_t>();
const int64_t* didx = dst_idx.data_ptr<int64_t>();
if (D == 128) {
edge_dot128_kernel<int64_t><<<grid, block, 0, stream>>>(sidx, didx,
src_feat.data_ptr<float>(), dst_feat.data_ptr<float>(), out.data_ptr<float>(), E);
} else {
edge_dot_generic_kernel<int64_t><<<grid, block, 0, stream>>>(sidx, didx,
src_feat.data_ptr<float>(), dst_feat.data_ptr<float>(), out.data_ptr<float>(), E, D);
}
}
C10_CUDA_KERNEL_LAUNCH_CHECK();
return out;
}
'''
_ext = load_inline(
name="edge_sampled_dot_ext_v1",
cpp_sources=_cpp_src,
cuda_sources=_cuda_src,
functions=["edge_dot_forward"],
extra_cflags=["-O3"],
extra_cuda_cflags=["-O3"],
with_cuda=True,
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)