reasoning
Let me analyze this task:
1. The reference model does edge-only sampled dense-dense matmul: for each edge, it indexes into src_feat and dst_feat, then computes dot product.
2. The dimensions are: num_src_nodes=1024, num_dst_nodes=768, num_edges=16384, feat_dim=128.
3. This is essentially a gather + element-wise multiply + reduce operation.
A simple CUDA kernel can be more efficient than PyTorch's native operations because:
- PyTorch does 3 separate ops: gather src, gather dst, then dot product
- A fused kernel can do all three in one go, reducing memory bandwidth
The operation per edge:
- Read feat_dim floats from src_feat[src_idx]
- Read feat_dim floats from dst_feat[dst_idx]
- Multiply and sum
Memory access pattern:
- src_idx and dst_idx are random accesses (gather)
- src_feat and dst_feat are 2D tensors
- The output is 1D tensor of size num_edges
For 16384 edges and 128 features:
- Total output: 16384 floats = 64KB
- Total reads from features: 2 * 16384 * 128 = 4,194,304 floats = 16MB
This is bandwidth-bound but PyTorch's separate operations have overhead.
Let me write a simple CUDA kernel that processes each edge with one thread, and each thread computes a dot product of length 128.
Actually, with feat_dim=128, we could use vectorized loads (float4) to load 4 floats at a time, which would be more efficient.
Let me write a kernel:
```cuda
__global__ void edge_dot_kernel(
const int* src_idx,
const int* dst_idx,
const float* src_feat,
const float* dst_feat,
float* out,
int num_edges,
int feat_dim,
int num_src_nodes,
int num_dst_nodes
) {
int edge = blockIdx.x * blockDim.x + threadIdx.x;
if (edge >= num_edges) return;
int s = src_idx[edge];
int d = dst_idx[edge];
const float* src_ptr = src_feat + s * feat_dim;
const float* dst_ptr = dst_feat + d * feat_dim;
float sum = 0.0f;
// Unroll by 4 using float4
int vec_len = feat_dim / 4;
for (int i = 0; i < vec_len; i++) {
float4 s4 = ((const float4*)src_ptr)[i];
float4 d4 = ((const float4*)dst_ptr)[i];
sum += s4.x * d4.x;
sum += s4.y * d4.y;
sum += s4.z * d4.z;
sum += s4.w * d4.w;
}
out[edge] = sum;
}
```
Wait, let me make sure alignment is correct. The feature dimension is 128, which is divisible by 4, so float4 should work fine. But are the row pointers aligned to 16 bytes? PyTorch tensors are typically aligned, but each row starting at `s * feat_dim` - for feat_dim=128, s*128*sizeof(float) = s*512, which is divisible by 16 for any integer s, since 512 is divisible by 16. So yes, float4 loads are safe.
Let me also think about block size and grid size. 16384 edges with 256 threads per block = 64 blocks. This might not be enough to saturate the GPU. Let me use a larger block size or more edges per thread. Actually 64 blocks on H100 with 132 SMs might be okay but we could also process multiple edges per thread or use a 2D grid.
Actually with 16384 edges and 128 features, each edge does 256 float reads + 128 multiply-adds. The computation is very light. The key is to saturate memory bandwidth.
PyTorch's approach:
1. `src_feat[src_idx]` - indexed select, outputs [16384, 128]
2. `dst_feat[dst_idx]` - indexed select, outputs [16384, 128]
3. `src * dst` - elementwise mul, reads and writes [16384, 128]
4. `.sum(dim=-1)` - reduction, reads [16384, 128], writes [16384]
Total memory traffic (roughly):
- Step 1: reads src_feat (reads ~16MB from global, but actually much less due to caching since only 1024*128=512KB unique), writes 8MB
- Step 2: reads dst_feat (~12MB unique, 768*128=384KB), writes 8MB
- Step 3: reads 16MB, writes 8MB
- Step 4: reads 8MB, writes 64KB
Total: lots of temporary memory traffic!
The fused kernel:
- Reads src_idx, dst_idx (small)
- Reads src_feat (gather, ~512KB unique but accessed randomly)
- Reads dst_feat (gather, ~384KB unique but accessed randomly)
- Writes output (64KB)
This should be much faster due to less memory traffic. The feature tables are small and will likely stay in cache.
Let me write this up in Python with load_inline.
Actually, I should be careful about int64 vs int32. The reference uses `.long()` for indexing, but the inputs are int32. In the reference forward, `src_idx` and `dst_idx` are int32 tensors, and then `.long()` converts them to int64 before indexing.
In my CUDA kernel, I should accept int32 indices since that's what get_inputs provides. Wait, let me check the reference again:
```python
def forward(self, src_idx, dst_idx, src_feat, dst_feat):
src = src_feat[src_idx.long()]
dst = dst_feat[dst_idx.long()]
return (src * dst).sum(dim=-1)
```
So the inputs are int32 but converted to long in the reference. For correctness, I need to make sure I handle whatever type is passed. Let me accept int32 since that's what get_inputs returns, and that's what gets passed to ModelNew.
I should write a host wrapper function that takes torch tensors and calls the kernel.
Let me also consider using template parameters for feat_dim to allow loop unrolling. But feat_dim=128 so I can just pass it as a runtime parameter.
Actually, I can make the kernel more efficient by having each thread process 1 edge with vectorized loads. Let me write it carefully.
One concern: float4 loads require 16-byte alignment. Let me verify: src_feat rows are at offset s*128*4 = s*512 bytes. 512 is divisible by 16, so all rows are 16-byte aligned. Similarly for dst_feat with d*128*4 = d*512. So float4 is safe.
Let me write the kernel code.
```python
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = '''
#include <torch/extension.h>
__global__ void edge_dot_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__ out,
int num_edges,
int feat_dim
) {
int edge = blockIdx.x * blockDim.x + threadIdx.x;
if (edge >= num_edges) return;
int s = src_idx[edge];
int d = dst_idx[edge];
const float4* src_vec = reinterpret_cast<const float4*>(src_feat + s * feat_dim);
const float4* dst_vec = reinterpret_cast<const float4*>(dst_feat + d * feat_dim);
float sum = 0.0f;
int vec_len = feat_dim >> 2;
#pragma unroll 16
for (int i = 0; i < vec_len; i++) {
float4 s4 = src_vec[i];
float4 d4 = dst_vec[i];
sum += s4.x * d4.x;
sum += s4.y * d4.y;
sum += s4.z * d4.z;
sum += s4.w * d4.w;
}
out[edge] = sum;
}
torch::Tensor edge_dot_forward(
torch::Tensor src_idx,
torch::Tensor dst_idx,
torch::Tensor src_feat,
torch::Tensor dst_feat
) {
int num_edges = src_idx.numel();
int feat_dim = src_feat.size(1);
auto out = torch::empty({num_edges}, src_feat.options());
const int threads = 256;
const int blocks = (num_edges + threads - 1) / threads;
edge_dot_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>(),
out.data_ptr<float>(),
num_edges,
feat_dim
);
return out;
}
'''
```
Wait, but what if src_idx or dst_idx are int64? The reference uses `.long()` which converts to int64. But get_inputs returns int32. Let me check what actually gets passed to ModelNew.
Looking at get_inputs:
```python
src_idx = torch.randint(0, num_src_nodes, (num_edges,), dtype=torch.int32)
```
So it returns int32. But the reference does `.long()` to convert to int64 for indexing.
For my kernel, I should handle both int32 and int64 to be safe, or cast to int32 if needed. Let me make it accept int64_t as well, or just cast in the wrapper.
Actually, the simplest is to check the dtype and handle both. Or cast to int32 if it's int64, since num_nodes is small.
Let me modify to handle both:
```cpp
if (src_idx.scalar_type() == torch::kInt32) {
// call kernel with int32
} else {
// call kernel with int64 and cast
}
```
Actually, to keep it simple, let me just cast to int32 in the host wrapper if needed:
```cpp
if (src_idx.dtype() != torch::kInt32) {
src_idx = src_idx.to(torch::kInt32);
}
```
But that adds an extra copy. Since the inputs are int32, this shouldn't happen in the benchmark.
Actually, let me be more robust and have two kernel instantiations. But that's more code. Let me just handle int32 since that's what get_inputs returns, and the evaluation matches get_inputs.
Wait, looking more carefully: the task says "Produce numerically equivalent outputs at fp32 precision." It doesn't say I need to handle all dtypes. The reference takes any tensor type and casts to long. But get_inputs produces int32.
To be safe, let me handle both int32 and int64. I can do this with a template:
```cpp
template <typename IndexT>
__global__ void edge_dot_kernel(...) {
IndexT s = src_idx[edge];
...
}
```
And dispatch in the host function.
Actually wait, I should also make sure the output is on the same device. The reference operations will work on CUDA if inputs are CUDA. But get_inputs returns CPU tensors? Let me check - the task doesn't explicitly say, but typically these run on CPU inputs and the model moves to GPU or stays on CPU.
Actually, looking at this more carefully, these benchmarks run on GPU. The reference model operations will run on whatever device the inputs are on. If on CPU, they'd be slow. So I assume inputs are on GPU or should be moved to GPU. But since the reference is PyTorch and it's being timed, they're probably on GPU already.
Actually, let me check: in get_inputs(), there's no `.cuda()`, so by default the tensors are on CPU. But the timing measurement would run them on CPU which would be slow. However, the evaluation typically moves them to GPU before timing. Or they stay on CPU and GPU kernel evaluation also runs on CPU? No, the CUDA kernel needs GPU inputs.
I think the evaluation framework handles device placement. Or the inputs might already be on GPU. Let me just write the kernel assuming all tensors are on the same CUDA device, which is standard.
Actually, I can add `TORCH_CHECK(src_idx.is_cuda(), ...)` or just let it fail naturally.
Let me now write the complete ModelNew class.
One more optimization: I can use `__ldg` (load through read-only data cache) for the feature reads to improve caching behavior. Or use `const __restrict__` which helps the compiler.
For small node counts (1024 and 768 nodes) with 128-dim features, the feature tables are tiny (512KB and 384KB), so they should easily fit in L2 cache. The gather pattern benefits from cache locality.
Let me make sure to use `__restrict__` to tell the compiler there's no aliasing.
Also, I should tune the block size. 256 is standard, but maybe 128 or 512 would be better. With 16384 edges:
- 128 threads/block -> 128 blocks
- 256 threads/block -> 64 blocks
- 512 threads/block -> 32 blocks
H100 has 132 SMs. More blocks is better for occupancy, but each block is very lightweight. 128 blocks should give good distribution. Let me try 256 first and see.
Actually, since each thread is doing a nontrivial amount of work (32 float4 loads = 32 memory instructions + 32*4 FMAs = 128 FMAs), 64 blocks of 256 threads might be fine. Each SM can handle multiple blocks.
Wait, I don't need to overthink this. Let me just write a clean kernel and see if it compiles and passes.
```python
import torch
from torch.utils.cpp_extension import load_inline
CUDA_SOURCE = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
template <typename IndexT>
__global__ void edge_dot_kernel(
const IndexT* __restrict__ src_idx,
const IndexT* __restrict__ dst_idx,
const float* __restrict__ src_feat,
const float* __restrict__ dst_feat,
float* __restrict__ out,
int num_edges,
int feat_dim
) {
int edge = blockIdx.x * blockDim.x + threadIdx.x;
if (edge >= num_edges) return;
int s = static_cast<int>(src_idx[edge]);
int d = static_cast<int>(dst_idx[edge]);
const float4* src_vec = reinterpret_cast<const float4*>(src_feat + s * feat_dim);
const float4* dst_vec = reinterpret_cast<const float4*>(dst_feat + d * feat_dim);
float sum = 0.0f;
#pragma unroll 16
for (int i = 0; i < feat_dim / 4; ++i) {
float4 s4 = src_vec[i];
float4 d4 = dst_vec[i];
sum += s4.x * d4.x;
sum += s4.y * d4.y;
sum += s4.z * d4.z;
sum += s4.w * d4.w;
}
out[edge] = sum;
}
torch::Tensor edge_dot_forward(
torch::Tensor src_idx,
torch::Tensor dst_idx,
torch::Tensor src_feat,
torch::Tensor dst_feat
) {
int num_edges = src_idx.numel();
int feat_dim = src_feat.size(1);
auto out = torch::empty({num_edges}, src_feat.options());
const int threads = 256;
const int blocks = (num_edges + threads - 1) / threads;
AT_DISPATCH_INDEX_TYPES(src_idx.scalar_type(), "edge_dot_forward", [&] {
edge_dot_kernel<<<blocks, threads>>>(
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>(),
num_edges,
feat_dim
);
});
return out;
}
"""
class ModelNew(torch.nn.Module):
def __init__(self):
super().__init__()
self.edge_dot = load_inline(
name="edge_dot",
cpp_sources="",
cuda_sources=CUDA_SOURCE,
functions=["edge_dot_forward"],
extra_cuda_cflags=["-O3"],
verbose=False,
).edge_dot_forward
def forward(self, src_idx, dst_idx, src_feat, dst_feat):
return self.edge_dot(src_idx, dst_idx, src_feat, dst_feat)
```
Hmm, wait. `AT_DISPATCH_INDEX_TYPES` only handles int64_t. For int32, PyTorch uses `AT_DISPATCH_INTEGRAL_TYPES` or I need to handle it differently.
Actually, PyTorch's `AT_DISPATCH_INDEX_TYPES` dispatches `int64_t` (Long) and maybe others. For int32, I should use `AT_DISPATCH_INTEGRAL_TYPES` which handles both int32 and int64.
Let me fix that.
```cpp
AT_DISPATCH_INTEGRAL_TYPES(src_idx.scalar_type(), "edge_dot_forward", [&] {
edge_dot_kernel<<<blocks, threads>>>(
src_idx.data_ptr<scalar_t>(),
...
);
});
```
Yes, `AT_DISPATCH_INTEGRAL_TYPES` handles int32 (kInt) and int64 (kLong). Perfect.
Wait, there's another issue. `load_inline` needs proper setup. Let me make sure the C++ code works with load_inline. The `functions` parameter tells it which functions to expose.
Also, I should add `#include <ATen/ATen.h>` or make sure torch/extension.h includes what we need. It should.
Let me also think about whether I should handle CPU tensors. The inputs from get_inputs are CPU tensors. If the evaluation doesn't move them to GPU, my kernel will fail. But typically CUDA kernel evaluations assume CUDA tensors. Actually, let me just check if I need to handle device placement.
The reference timing will use whatever device the inputs are on. If inputs are CPU, PyTorch will run on CPU. For a fair comparison, I should probably run on the same device. But a CUDA kernel can only run on GPU.
I suspect the evaluation framework moves inputs to GPU before running the model. This is standard practice in these benchmarks. Let me proceed with that assumption.
Let me write a clean version and test compile.
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nfrom torch.utils.cpp_extension import load_inline\n\nCUDA_SOURCE = r\"\"\"\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n\ntemplate <typename IndexT>\n__global__ void edge_dot_kernel(\n const IndexT* __restrict__ src_idx,\n const IndexT* __restrict__ dst_idx,\n const float* __restrict__ src_feat,\n const float* __restrict__ dst_feat,\n float* __restrict__ out,\n int num_edges,\n int feat_dim\n) {\n int edge = blockIdx.x * blockDim.x + threadIdx.x;\n if (edge >= num_edges) return;\n\n int s = static_cast<int>(src_idx[edge]);\n int d = static_cast<int>(dst_idx[edge]);\n\n const float4* src_vec = reinterpret_cast<const float4*>(src_feat + s * feat_dim);\n const float4* dst_vec = reinterpret_cast<const float4*>(dst_feat + d * feat_dim);\n\n float sum = 0.0f;\n #pragma unroll 16\n for (int i = 0; i < feat_dim / 4; ++i) {\n float4 s4 = src_vec[i];\n float4 d4 = dst_vec[i];\n sum += s4.x * d4.x;\n sum += s4.y * d4.y;\n sum += s4.z * d4.z;\n sum += s4.w * d4.w;\n }\n\n out[edge] = sum;\n}\n\ntorch::Tensor edge_dot_forward(\n torch::Tensor src_idx,\n torch::Tensor dst_idx,\n torch::Tensor src_feat,\n torch::Tensor dst_feat\n) {\n int num_edges = src_idx.numel();\n int feat_dim = src_feat.size(1);\n\n auto out = torch::empty({num_edges}, src_feat.options());\n\n const int threads = 256;\n const int blocks = (num_edges + threads - 1) / threads;\n\n AT_DISPATCH_INTEGRAL_TYPES(src_idx.scalar_type(), \"edge_dot_forward\", [&] {\n edge_dot_kernel<<<blocks, threads>>>(\n src_idx.data_ptr<scalar_t>(),\n dst_idx.data_ptr<scalar_t>(),\n src_feat.data_ptr<float>(),\n dst_feat.data_ptr<float>(),\n out.data_ptr<float>(),\n num_edges,\n feat_dim\n );\n });\n\n return out;\n}\n\"\"\"\n\n# Pre-compile at module load\n_edge_dot = load_inline(\n name=\"edge_dot\",\n cpp_sources=\"\",\n cuda_sources=CUDA_SOURCE,\n functions=[\"edge_dot_forward\"],\n extra_cuda_cflags=[\"-O3\"],\n verbose=False,\n)\n\n\nclass ModelNew(torch.nn.Module):\n def __init__(self):\n super().__init__()\n self.edge_dot = _edge_dot.edge_dot_forward\n\n def forward(self, src_idx, dst_idx, src_feat, dst_feat):\n return self.edge_dot(src_idx, dst_idx, src_feat, dst_feat)\n"
}