reasoning
reasoning (truncated 45158 chars): Let me understand the task:
1. I need to replace a PyTorch model with a custom CUDA kernel
2. The model does multi-head sparse aggregation with per-edge attention weights
3. The forward function takes:
- `row_ptr`: CSR format row pointers (num_nodes + 1)
- `col_idx`: column indices for edges
- `edge_weight`: weights per edge (num_edges, num_heads)
- `node_feat`: node features (num_nodes, num_heads, head_dim)
The operation is essentially a sparse matrix multiplication where for each destination node, we aggregate weighted source features.
Let me write the CUDA kernel:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
# Define CUDA kernel
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
template <typename scalar_t>
__global__ void sparse_aggregate_kernel(
const int* __restrict__ row_ptr,
const int* __restrict__ col_idx,
const scalar_t* __restrict__ edge_weight,
const scalar_t* __restrict__ node_feat,
scalar_t* __restrict__ out,
int num_nodes,
int num_heads,
int head_dim) {
int dst = blockIdx.x * blockDim.x + threadIdx.x;
if (dst >= num_nodes) return;
int start = row_ptr[dst];
int end = row_ptr[dst + 1];
// For each head and each head dimension
for (int h = 0; h < num_heads; h++) {
for (int d = 0; d < head_dim; d++) {
scalar_t sum = 0;
for (int e = start; e < end; e++) {
int src = col_idx[e];
sum += node_feat[src * num_heads * head_dim + h * head_dim + d] * edge_weight[e * num_heads + h];
}
out[dst * num_heads * head_dim + h * head_dim + d] = sum;
}
}
}
torch::Tensor sparse_aggregate_cuda(
torch::Tensor row_ptr,
torch::Tensor col_idx,
torch::Tensor edge_weight,
torch::Tensor node_feat) {
int num_nodes = row_ptr.numel() - 1;
int num_heads = node_feat.size(1);
int head_dim = node_feat.size(2);
auto out = torch::zeros({num_nodes, num_heads, head_dim}, node_feat.options());
const int threads = 256;
const int blocks = (num_nodes + threads - 1) / threads;
AT_DISPATCH_FLOATING_TYPES(node_feat.scalar_type(), "sparse_aggregate_kernel", ([&] {
sparse_aggregate_kernel<scalar_t><<<blocks, threads>>>(
row_ptr.data_ptr<int>(),
col_idx.data_ptr<int>(),
edge_weight.data_ptr<scalar_t>(),
node_feat.data_ptr<scalar_t>(),
out.data_ptr<scalar_t>(),
num_nodes,
num_heads,
head_dim);
}));
return out;
}
"""
cpp_source = """
torch::Tensor sparse_aggregate_cuda(
torch::Tensor row_ptr,
torch::Tensor col_idx,
torch::Tensor edge_weight,
torch::Tensor node_feat);
"""
sparse_aggregate = load_inline(
name='sparse_aggregate',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['sparse_aggregate_cuda'],
verbose=True
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.sparse_aggregate = sparse_aggregate
def forward(
self,
row_ptr: torch.Tensor,
col_idx: torch.Tensor,
edge_weight: torch.Tensor,
node_feat: torch.Tensor,
) -> torch.Tensor:
return sparse_aggregate.sparse_aggregate_cuda(row_ptr, col_idx, edge_weight, node_feat)
```
This should work, but let me think about optimizations:
1. The current kernel assigns one thread per destination node
2. Each thread loops over all edges for that node, then all heads and head_dim
3. This might not be optimal - we could potentially parallelize over heads/dimensions
Let me think about a better approach:
Option 1: One thread per (dst, head) pair
- This would give more parallelism
- Each thread would sum over edges for one head
Option 2: One thread per (dst, head, dim) element
- Maximum parallelism
- But memory access pattern might not be optimal
Option 3: Use shared memory for weights
- If num_heads * head_dim is small enough, we could load features into shared memory
Let me try option 2 for maximum parallelism:
```cuda
template <typename scalar_t>
__global__ void sparse_aggregate_kernel(
const int* __restrict__ row_ptr,
const int* __restrict__ col_idx,
const scalar_t* __restrict__ edge_weight,
const scalar_t* __restrict__ node_feat,
scalar_t* __restrict__ out,
int num_nodes,
int num_heads,
int head_dim) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = num_nodes * num_heads * head_dim;
if (idx >= total) return;
int dst = idx / (num_heads * head_dim);
int remainder = idx % (num_heads * head_dim);
int h = remainder / head_dim;
int d = remainder % head_dim;
int start = row_ptr[dst];
int end = row_ptr[dst + 1];
scalar_t sum = 0;
for (int e = start; e < end; e++) {
int src = col_idx[e];
sum += node_feat[src * num_heads * head_dim + h * head_dim + d] * edge_weight[e * num_heads + h];
}
out[dst * num_heads * head_dim + h * head_dim + d] = sum;
}
```
Actually, let me think about memory access patterns. The edge_weight is accessed as edge_weight[e * num_heads + h]. If we iterate over edges in a loop, we're accessing edge_weight with stride num_heads for consecutive edges, which might be okay.
For node_feat, we access node_feat[src * num_heads * head_dim + h * head_dim + d]. If multiple threads are working on different (dst, h, d) combinations but the same edge, they would access the same node_feat element, which is fine.
Let me try a hybrid approach where each thread handles one (dst, h) pair:
```cuda
template <typename scalar_t>
__global__ void sparse_aggregate_kernel(
const int* __restrict__ row_ptr,
const int* __restrict__ col_idx,
const scalar_t* __restrict__ edge_weight,
const scalar_t* __restrict__ node_feat,
scalar_t* __restrict__ out,
int num_nodes,
int num_heads,
int head_dim) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = num_nodes * num_heads;
if (idx >= total) return;
int dst = idx / num_heads;
int h = idx % num_heads;
int start = row_ptr[dst];
int end = row_ptr[dst + 1];
for (int d = 0; d < head_dim; d++) {
scalar_t sum = 0;
for (int e = start; e < end; e++) {
int src = col_idx[e];
sum += node_feat[src * num_heads * head_dim + h * head_dim + d] * edge_weight[e * num_heads + h];
}
out[dst * num_heads * head_dim + h * head_dim + d] = sum;
}
}
```
This approach has each thread handle all head_dim elements for one (dst, h) pair, which allows for vectorization if we use float4 for example.
Actually, let me think about using float4 for better memory throughput:
```cuda
template <typename scalar_t>
__global__ void sparse_aggregate_kernel_vec(
const int* __restrict__ row_ptr,
const int* __restrict__ col_idx,
const scalar_t* __restrict__ edge_weight,
const scalar_t* __restrict__ node_feat,
scalar_t* __restrict__ out,
int num_nodes,
int num_heads,
int head_dim) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = num_nodes * num_heads;
if (idx >= total) return;
int dst = idx / num_heads;
int h = idx % num_heads;
int start = row_ptr[dst];
int end = row_ptr[dst + 1];
// Process head_dim elements
// If head_dim is 32, we can use float4 to process 4 elements at a time
float4* node_feat4 = (float4*)node_feat;
float4* out4 = (float4*)out;
int head_dim4 = head_dim / 4;
for (int d4 = 0; d4 < head_dim4; d4++) {
float4 sum4 = {0, 0, 0, 0};
for (int e = start; e < end; e++) {
int src = col_idx[e];
float4 feat4 = node_feat4[(src * num_heads + h) * head_dim4 + d4];
scalar_t w = edge_weight[
...[ omitted 45158 characters ]...
late all d values:
Actually, let me think about this more carefully. The current code:
```cuda
for (int d = 0; d < head_dim; d++) {
scalar_t sum = scalar_t(0);
for (int e = start; e < end; e++) {
// ...
}
out[(dst * num_heads + h) * head_dim + d] = sum;
}
```
This has:
- head_dim iterations of the outer loop
- For each outer iteration, num_edges inner iterations
- Total: head_dim * num_edges operations
The memory access pattern:
- For each d, we access col_idx[start..end] and edge_weight[start*num_heads..end*num_heads] with stride num_heads
- For each d, we access node_feat with random access based on src
The alternative:
```cuda
scalar_t local_sum[32];
for (int d = 0; d < head_dim; d++) local_sum[d] = 0;
for (int e = start; e < end; e++) {
int src = col_idx[e];
scalar_t w = edge_weight[e * num_heads + h];
for (int d = 0; d < head_dim; d++) {
local_sum[d] += node_feat[(src * num_heads + h) * head_dim + d] * w;
}
}
for (int d = 0; d < head_dim; d++) {
out[(dst * num_heads + h) * head_dim + d] = local_sum[d];
}
```
This has:
- One iteration over edges
- For each edge, head_dim iterations to accumulate
The memory access pattern:
- We access col_idx and edge_weight once per edge
- For each edge, we access node_feat[src, h, 0..head_dim-1], which is head_dim consecutive elements
This is better because:
1. We load col_idx and edge_weight once instead of head_dim times
2. We access node_feat with better locality (consecutive elements)
The downside is that we need to store local_sum in registers, which might reduce occupancy.
Let me implement this optimized version:
```cuda
template <typename scalar_t>
__global__ void sparse_aggregate_kernel_opt(
const int32_t* __restrict__ row_ptr,
const int32_t* __restrict__ col_idx,
const scalar_t* __restrict__ edge_weight,
const scalar_t* __restrict__ node_feat,
scalar_t* __restrict__ out,
int num_nodes,
int num_heads,
int head_dim) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = num_nodes * num_heads;
if (idx >= total) return;
int dst = idx / num_heads;
int h = idx % num_heads;
int start = row_ptr[dst];
int end = row_ptr[dst + 1];
// Initialize local accumulators
scalar_t local_sum[32]; // Assuming head_dim <= 32
#pragma unroll
for (int d = 0; d < 32; d++) {
local_sum[d] = scalar_t(0);
}
// Iterate over edges
for (int e = start; e < end; e++) {
int src = col_idx[e];
scalar_t w = edge_weight[e * num_heads + h];
// Accumulate for all dimensions
#pragma unroll 8
for (int d = 0; d < head_dim; d++) {
local_sum[d] += node_feat[(src * num_heads + h) * head_dim + d] * w;
}
}
// Write results
#pragma unroll
for (int d = 0; d < 32; d++) {
if (d < head_dim) {
out[(dst * num_heads + h) * head_dim + d] = local_sum[d];
}
}
}
```
Actually, I realize I should be careful about the `#pragma unroll` directives. They might cause issues if the compiler can't unroll the loop.
Let me simplify and just trust the compiler:
```cuda
template <typename scalar_t, int HEAD_DIM>
__global__ void sparse_aggregate_kernel_opt(
const int32_t* __restrict__ row_ptr,
const int32_t* __restrict__ col_idx,
const scalar_t* __restrict__ edge_weight,
const scalar_t* __restrict__ node_feat,
scalar_t* __restrict__ out,
int num_nodes,
int num_heads) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = num_nodes * num_heads;
if (idx >= total) return;
int dst = idx / num_heads;
int h = idx % num_heads;
int start = row_ptr[dst];
int end = row_ptr[dst + 1];
// Initialize local accumulators
scalar_t local_sum[HEAD_DIM];
for (int d = 0; d < HEAD_DIM; d++) {
local_sum[d] = scalar_t(0);
}
// Iterate over edges
for (int e = start; e < end; e++) {
int src = col_idx[e];
scalar_t w = edge_weight[e * num_heads + h];
// Accumulate for all dimensions
for (int d = 0; d < HEAD_DIM; d++) {
local_sum[d] += node_feat[(src * num_heads + h) * HEAD_DIM + d] * w;
}
}
// Write results
for (int d = 0; d < HEAD_DIM; d++) {
out[(dst * num_heads + h) * HEAD_DIM + d] = local_sum[d];
}
}
```
This uses a template parameter for HEAD_DIM, which allows the compiler to unroll the loops.
Actually, this requires knowing HEAD_DIM at compile time, which is fine since it's always 32 for this problem.
Let me write the final version:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
sparse_aggregate_ext = load_inline(
name='sparse_aggregate_ext',
cpp_sources='''
torch::Tensor sparse_aggregate_cuda(
torch::Tensor row_ptr,
torch::Tensor col_idx,
torch::Tensor edge_weight,
torch::Tensor node_feat);
''',
cuda_sources='''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cstdint>
template <typename scalar_t>
__global__ void sparse_aggregate_kernel(
const int32_t* __restrict__ row_ptr,
const int32_t* __restrict__ col_idx,
const scalar_t* __restrict__ edge_weight,
const scalar_t* __restrict__ node_feat,
scalar_t* __restrict__ out,
int num_nodes,
int num_heads,
int head_dim) {
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
const int total = num_nodes * num_heads;
if (idx >= total) return;
const int dst = idx / num_heads;
const int h = idx % num_heads;
const int start = row_ptr[dst];
const int end = row_ptr[dst + 1];
// Initialize local accumulators
// Assume head_dim <= 32 to avoid dynamic memory
scalar_t local_sum[32];
for (int d = 0; d < head_dim && d < 32; d++) {
local_sum[d] = scalar_t(0);
}
// Iterate over edges
for (int e = start; e < end; e++) {
const int src = col_idx[e];
const scalar_t w = edge_weight[e * num_heads + h];
// Accumulate for all dimensions
for (int d = 0; d < head_dim && d < 32; d++) {
local_sum[d] += node_feat[(src * num_heads + h) * head_dim + d] * w;
}
}
// Write results
for (int d = 0; d < head_dim && d < 32; d++) {
out[(dst * num_heads + h) * head_dim + d] = local_sum[d];
}
}
torch::Tensor sparse_aggregate_cuda(
torch::Tensor row_ptr,
torch::Tensor col_idx,
torch::Tensor edge_weight,
torch::Tensor node_feat) {
const int num_nodes = row_ptr.numel() - 1;
const int num_heads = node_feat.size(1);
const int head_dim = node_feat.size(2);
auto out = torch::zeros({num_nodes, num_heads, head_dim}, node_feat.options());
const int threads = 256;
const int total = num_nodes * num_heads;
const int blocks = (total + threads - 1) / threads;
AT_DISPATCH_FLOATING_TYPES(node_feat.scalar_type(), "sparse_aggregate_kernel", ([&] {
sparse_aggregate_kernel<scalar_t><<<blocks, threads>>>(
row_ptr.data_ptr<int32_t>(),
col_idx.data_ptr<int32_t>(),
edge_weight.data_ptr<scalar_t>(),
node_feat.data_ptr<scalar_t>(),
out.data_ptr<scalar_t>(),
num_nodes,
num_heads,
head_dim);
}));
return out;
}
''',
functions=['sparse_aggregate_cuda'],
verbose=True
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(
self,
row_ptr: torch.Tensor,
col_idx: torch.Tensor,
edge_weight: torch.Tensor,
node_feat: torch.Tensor,
) -> torch.Tensor:
return sparse_aggregate_ext.sparse_aggregate_cuda(row_ptr, col_idx, edge_weight, node_feat)
```
I think this is a good solution. Let me submit it.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\nsparse_aggregate_ext = load_inline(\n name='sparse_aggregate_ext',\n cpp_sources='''\ntorch::Tensor sparse_aggregate_cuda(\n torch::Tensor row_ptr,\n torch::Tensor col_idx,\n torch::Tensor edge_weight,\n torch::Tensor node_feat);\n''',\n cuda_sources='''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cstdint>\n\ntemplate <typename scalar_t>\n__global__ void sparse_aggregate_kernel(\n const int32_t* __restrict__ row_ptr,\n const int32_t* __restrict__ col_idx,\n const scalar_t* __restrict__ edge_weight,\n const scalar_t* __restrict__ node_feat,\n scalar_t* __restrict__ out,\n int num_nodes,\n int num_heads,\n int head_dim) {\n \n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n const int total = num_nodes * num_heads;\n \n if (idx >= total) return;\n \n const int dst = idx / num_heads;\n const int h = idx % num_heads;\n \n const int start = row_ptr[dst];\n const int end = row_ptr[dst + 1];\n \n // Initialize local accumulators in registers\n scalar_t local_sum[32];\n #pragma unroll\n for (int d = 0; d < 32; d++) {\n local_sum[d] = scalar_t(0);\n }\n \n // Iterate over edges once, accumulate all head_dim elements\n for (int e = start; e < end; e++) {\n const int src = col_idx[e];\n const scalar_t w = edge_weight[e * num_heads + h];\n \n // Vectorized load and accumulate\n #pragma unroll 8\n for (int d = 0; d < head_dim; d++) {\n local_sum[d] += node_feat[(src * num_heads + h) * head_dim + d] * w;\n }\n }\n \n // Write results\n #pragma unroll\n for (int d = 0; d < 32; d++) {\n if (d < head_dim) {\n out[(dst * num_heads + h) * head_dim + d] = local_sum[d];\n }\n }\n}\n\ntorch::Tensor sparse_aggregate_cuda(\n torch::Tensor row_ptr,\n torch::Tensor col_idx,\n torch::Tensor edge_weight,\n torch::Tensor node_feat) {\n \n const int num_nodes = row_ptr.numel() - 1;\n const int num_heads = node_feat.size(1);\n const int head_dim = node_feat.size(2);\n \n auto out = torch::zeros({num_nodes, num_heads, head_dim}, node_feat.options());\n \n const int threads = 256;\n const int total = num_nodes * num_heads;\n const int blocks = (total + threads - 1) / threads;\n \n AT_DISPATCH_FLOATING_TYPES(node_feat.scalar_type(), \"sparse_aggregate_kernel\", ([&] {\n sparse_aggregate_kernel<scalar_t><<<blocks, threads>>>(\n row_ptr.data_ptr<int32_t>(),\n col_idx.data_ptr<int32_t>(),\n edge_weight.data_ptr<scalar_t>(),\n node_feat.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n num_nodes,\n num_heads,\n head_dim);\n }));\n \n return out;\n}\n''',\n functions=['sparse_aggregate_cuda'],\n verbose=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n\n def forward(\n self,\n row_ptr: torch.Tensor,\n col_idx: torch.Tensor,\n edge_weight: torch.Tensor,\n node_feat: torch.Tensor,\n ) -> torch.Tensor:\n return sparse_aggregate_ext.sparse_aggregate_cuda(row_ptr, col_idx, edge_weight, node_feat)"
}