reasoning
Let me analyze this task:
1. We need to implement edge-only sampled dense-dense matmul
2. For each edge (src, dst), we compute a dot product between source and destination node embeddings
3. Input sizes:
- num_src_nodes = 1024
- num_dst_nodes = 768
- num_edges = 16384
- feat_dim = 128
The reference implementation:
1. Gathers source features using src_idx
2. Gathers destination features using dst_idx
3. Computes element-wise product
4. Sums along the feature dimension
Let me first get GPU specs to understand what I'm working with, then write an optimized CUDA kernel.
The key operations are:
- Gather: src_feat[src_idx] - shape (num_edges, feat_dim)
- Gather: dst_feat[dst_idx] - shape (num_edges, feat_dim)
- Element-wise multiply
- Sum reduction along dim -1
For optimization:
1. We can fuse all operations into one kernel
2. Each thread can handle one edge (compute dot product)
3. Load features, multiply, and sum in a loop over feat_dim
4. Potentially use shared memory for features if there's reuse
Let me think about memory access patterns:
- src_idx and dst_idx are random indices
- For each edge, we need to access one row from src_feat and one row from dst_feat
- There might be some locality if multiple edges share the same source or destination node
Given the sizes:
- num_edges = 16384
- feat_dim = 128
- Each feature vector is 128 * 4 = 512 bytes
A simple approach: each thread handles one edge, loops over feat_dim to compute dot product.
Let me write this kernel:
```cpp
__global__ void edge_dot_product_kernel(
const int* src_idx,
const int* dst_idx,
const float* src_feat,
const float* dst_feat,
float* output,
int num_edges,
int feat_dim,
int num_src_nodes,
int num_dst_nodes
) {
int edge_id = blockIdx.x * blockDim.x + threadIdx.x;
if (edge_id >= num_edges) return;
int src_id = src_idx[edge_id];
int dst_id = dst_idx[edge_id];
const float* src_vec = src_feat + src_id * feat_dim;
const float* dst_vec = dst_feat + dst_id * feat_dim;
float sum = 0.0f;
for (int i = 0; i < feat_dim; i++) {
sum += src_vec[i] * dst_vec[i];
}
output[edge_id] = sum;
}
```
This is a simple implementation. Let me think about optimizations:
1. Vectorized loads: Load 4 floats at a time using float4
2. Shared memory: Not much benefit here since each edge accesses random features
3. Coalesced memory access: The output is coalesced, but src/dst feature access is random
For better performance:
- Use vectorized loads (float4) for feature vectors
- Each thread handles one edge
- Loop unrolling
Let me check if feat_dim=128 is divisible by 4: 128/4 = 32, yes!
```cpp
__global__ void edge_dot_product_kernel(
const int* src_idx,
const int* dst_idx,
const float* src_feat,
const float* dst_feat,
float* output,
int num_edges,
int feat_dim,
int num_src_nodes,
int num_dst_nodes
) {
int edge_id = blockIdx.x * blockDim.x + threadIdx.x;
if (edge_id >= num_edges) return;
int src_id = src_idx[edge_id];
int dst_id = dst_idx[edge_id];
const float4* src_vec4 = (const float4*)(src_feat + src_id * feat_dim);
const float4* dst_vec4 = (const float4*)(dst_feat + dst_id * feat_dim);
float sum = 0.0f;
#pragma unroll
for (int i = 0; i < feat_dim / 4; i++) {
float4 s = src_vec4[i];
float4 d = dst_vec4[i];
sum += s.x * d.x + s.y * d.y + s.z * d.z + s.w * d.w;
}
output[edge_id] = sum;
}
```
Actually, let me think more carefully. The alignment of float4 might be an issue. The src_feat and dst_feat tensors should be properly aligned since they're created by PyTorch.
Let me also think about block size and grid size:
- num_edges = 16384
- If blockDim.x = 256, then we need 16384/256 = 64 blocks
- If blockDim.x = 128, then we need 16384/128 = 128 blocks
Let me also consider using warp-level primitives for reduction, but since each thread handles one edge independently, that's not needed here.
Actually, let me think about another optimization: what if multiple threads access the same source or destination node? We could potentially use shared memory to cache feature vectors, but with 1024 source nodes and 768 destination nodes, the total feature memory is (1024 + 768) * 128 * 4 = 920 KB, which exceeds shared memory limits (typically 48-164 KB per SM).
So the simple approach with vectorized loads should work well.
Let me write the full implementation:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
# Define the CUDA kernel
cuda_source = """
#include <cuda_runtime.h>
__global__ void edge_dot_product_kernel(
const int* __restrict__ src_idx,
const int* __restrict__ dst_idx,
const float* __restrict__ src_feat,
const float* __restrict__ dst_feat,
float* __restrict__ output,
int num_edges,
int feat_dim
) {
int edge_id = blockIdx.x * blockDim.x + threadIdx.x;
if (edge_id >= num_edges) return;
int src_id = src_idx[edge_id];
int dst_id = dst_idx[edge_id];
const float* src_vec = src_feat + src_id * feat_dim;
const float* dst_vec = dst_feat + dst_id * feat_dim;
float sum = 0.0f;
// Vectorized load using float4
const float4* src_vec4 = (const float4*)src_vec;
const float4* dst_vec4 = (const float4*)dst_vec;
for (int i = 0; i < feat_dim / 4; i++) {
float4 s = src_vec4[i];
float4 d = dst_vec4[i];
sum += s.x * d.x + s.y * d.y + s.z * d.z + s.w * d.w;
}
output[edge_id] = sum;
}
torch::Tensor edge_dot_product_cuda(
torch::Tensor src_idx,
torch::Tensor dst_idx,
torch::Tensor src_feat,
torch::Tensor dst_feat
) {
int num_edges = src_idx.size(0);
int feat_dim = src_feat.size(1);
auto output = torch::empty({num_edges}, src_feat.options());
int threads = 256;
int blocks = (num_edges + threads - 1) / threads;
edge_dot_product_kernel<<<blocks, threads>>>(
src_idx.data_ptr<int>(),
dst_idx.data_ptr<int>(),
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
);
"""
# Load the extension
edge_dot_product = load_inline(
name='edge_dot_product',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['edge_dot_product_cuda'],
verbose=True
)
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 edge_dot_product.edge_dot_product_cuda(src_idx, dst_idx, src_feat, dst_feat)
```
Wait, I should get the GPU specs first to understand the hardware better. Let me do that.
Also, I need to be careful about the data types. The reference uses src_idx.long() to convert to int64, but the input says src_idx is int32. Let me check the reference again...
The reference does:
```python
src = src_feat[src_idx.long()]
dst = dst_feat[dst_idx.long()]
return (src * dst).sum(dim=-1)
```
So src_idx and dst_idx are int32 tensors, but they're converted to long (int64) for indexing. The kernel should handle this properly.
Actually, looking at the inputs:
```python
src_idx = torch.randint(0, num_src_nodes, (num_edges,), dtype=torch.int32)
```
So src_idx is int32. My kernel uses `const int*` which matches int32 on most platforms, but I should double-check.
Let me compile and test this first.