reasoning
Let me analyze the problem:
1. This is a sparse attention/aggregation operation on a graph
2. Input:
- `row_ptr`: CSR row pointer (num_nodes+1,)
- `col_idx`: CSR column indices (num_edges,)
- `edge_weight`: per-edge per-head weights (num_edges, num_heads)
- `node_feat`: node features (num_nodes, num_heads, head_dim)
3. Output: (num_nodes, num_heads, head_dim)
4. For each destination node, we aggregate weighted source node features
The reference is a Python loop with PyTorch operations. For a CUDA kernel, I need to parallelize this efficiently.
Key parameters:
- num_nodes = 320
- avg_degree = 14
- num_heads = 8
- head_dim = 32
Total edges ≈ 320 * 14 = 4480
This is a relatively small graph. The total output size is 320 * 8 * 32 ≈ 82K elements.
Strategies:
1. **Simple approach**: One thread per output element (node, head, feature)
- Each thread iterates over incoming edges for that node
- Pros: Good parallelism (320*8*32 = 81920 threads)
- Cons: Memory access pattern may be suboptimal
2. **Shared memory approach**: Load node features into shared memory
- But the graph is small, so probably not needed
3. **Warp-level aggregation**: One warp per (node, head), each thread handles a feature dimension
- With 32 threads per warp and head_dim = 32, this maps perfectly
- Each warp processes one (node, head) pair
- Total warps needed: 320 * 8 = 2560
Given the small problem size, I think the most efficient approach is:
- **Option A**: One CUDA block per node, threads within block handle (head, feature)
- **Option B**: One warp per (node, head), using warp-level parallelism
Let me think about Option B more carefully:
- 320 nodes * 8 heads = 2560 (node, head) pairs
- Each pair needs to aggregate over ~14 edges, each edge reading head_dim=32 floats
- Total work: 2560 * 14 * 32 ≈ 1.1M multiply-adds
With H100, we have high memory bandwidth. The key is memory access pattern.
Actually, for such a small graph, we can use a persistent kernel or simple parallelization.
Let me design a kernel where each thread handles one (dst, head, feature_dim) element:
```cuda
__global__ void sparse_aggregate_kernel(
const int* row_ptr,
const int* col_idx,
const float* edge_weight,
const float* node_feat,
float* 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 f = idx % head_dim;
int h = (idx / head_dim) % num_heads;
int dst = idx / (num_heads * head_dim);
float sum = 0.0f;
int start = row_ptr[dst];
int end = row_ptr[dst + 1];
for (int e = start; e < end; e++) {
int src = col_idx[e];
float w = edge_weight[e * num_heads + h];
float feat = node_feat[(src * num_heads + h) * head_dim + f];
sum += w * feat;
}
out[idx] = sum;
}
```
This gives us perfect coalescing for `node_feat` and `out` since adjacent threads access adjacent feature dims. The `edge_weight` access is also coalesced when threads in a warp have the same head (which they do for head_dim consecutive threads).
Wait, let me check: for a fixed (dst, h), threads with different f values will access `edge_weight[e * num_heads + h]` - that's the same value for all threads in a block working on the same (dst, h). So there will be broadcast-style cache hits.
For `node_feat`, adjacent f values access consecutive memory: `[(src * num_heads + h) * head_dim + f]` - this is coalesced.
For `out`, adjacent f values write to consecutive memory - coalesced.
The main issue is that for small degree (avg 14), we only do 14 iterations per thread, and the memory loads from `col_idx` and `edge_weight` might not be fully utilized.
But given the small size, this should be very fast. Let's think about occupancy. With 81920 total threads, and if we use 256 threads per block: 320 blocks. Each block has 256 threads. On H100 with 132 SMs, we'd have good occupancy.
Actually, I realize that we can potentially be even more efficient. Since `head_dim = 32`, we could use warp-level primitives where each warp handles one (dst, head) and the 32 threads handle the 32 feature dimensions.
Actually the simple element-wise approach should be fine for this size. Let me first code it up and test.
Wait, there's another consideration. The reference implementation does:
```python
weights = edge_weight[start:end].unsqueeze(-1)
out[dst] = (node_feat[src] * weights).sum(dim=0)
```
Here `edge_weight` has shape (num_edges, num_heads), and `node_feat[src]` has shape (degree, num_heads, head_dim). The multiplication broadcasts: `weights` is (degree, num_heads, 1) and `node_feat[src]` is (degree, num_heads, head_dim).
So in my kernel, for edge e and head h, the weight is `edge_weight[e, h]`, same for all feature dims.
Let me write a clean kernel. Given the small graph, a simple CUDA kernel should be much faster than Python loops.
Actually, since num_nodes=320, num_heads=8, head_dim=32, total output elements = 320*8*32 = 81920.
If I launch with 1024 threads per block, I need ~80 blocks. Or 256 threads/block -> 320 blocks. Both should be fine.
Let me think about edge cases and memory layouts.
- `row_ptr` is int32
- `col_idx` is int32
- `edge_weight` is float32 with shape (num_edges, num_heads), so stride is num_heads for the first dim
- `node_feat` is float32 with shape (num_nodes, num_heads, head_dim)
- `out` is float32 with shape (num_nodes, num_heads, head_dim)
Memory layout assumption: PyTorch default is row-major (C-contiguous).
- `node_feat[(n * num_heads + h) * head_dim + f]`
- `edge_weight[e * num_heads + h]`
- `out[(n * num_heads + h) * head_dim + f]`
Let me write this carefully.
One more thought: for such a small graph, the degree variance matters. Let me think about whether I should use shared memory caching for the weights or features. Since degree is only ~14 on average, there's not much to gain from shared memory for a single (dst, head). But if a block processes multiple (dst, head) pairs, shared memory might help for `col_idx` and `edge_weight` if there are shared neighbors.
Actually given the problem is memory-bandwidth bound with such small compute per element, the simple element-wise kernel with good memory coalescing should be optimal or near-optimal.
Let me draft the kernel code:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = '''
__global__ void sparse_aggregate_kernel(
const int* __restrict__ row_ptr,
const int* __restrict__ col_idx,
const float* __restrict__ edge_weight,
const float* __restrict__ node_feat,
float* __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 * head_dim;
if (idx >= total) return;
const int f = idx % head_dim;
const int tmp = idx / head_dim;
const int h = tmp % num_heads;
const int dst = tmp / num_heads;
const int start = row_ptr[dst];
const int end = row_ptr[dst + 1];
float sum = 0.0f;
#pragma unroll 4
for (int e = start; e < end; e++) {
const int src = col_idx[e];
const float w = edge_weight[e * num_heads + h];
const float feat = node_feat[(src * num_heads + h) * head_dim + f];
sum += w * feat;
}
out[idx] = sum;
}
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 total = num_nodes * num_heads * head_dim;
const int threads = 256;
const int blocks = (total + threads - 1) / threads;
sparse_aggregate_kernel<<<blocks, threads>>>(
row_ptr.data_ptr<int>(),
col_idx.data_ptr<int>(),
edge_weight.data_ptr<float>(),
node_feat.data_ptr<float>(),
out.data_ptr<float>(),
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'],
extra_cuda_cflags=['-O3'],
)
```
Wait, I need to be more careful. The `load_inline` function signature:
- `cpp_sources` should contain only function declarations, not definitions (since pybind11 will generate bindings)
- `cuda_sources` should contain the actual CUDA implementation
- `functions` lists the names to bind
Actually, let me check the exact usage. In `load_inline`, `cpp_sources` is typically used for CPU/C++ wrapper code and `cuda_sources` for CUDA code. If I only have a CUDA function, I should put the declaration in `cpp_sources` and definition in `cuda_sources`.
But actually, I think the function can be declared and defined together in `cuda_sources` if it's marked with proper linkage, or I can put the torch::Tensor wrapper in `cpp_sources` and the kernel in `cuda_sources`.
Actually, for simplicity, let me put everything in `cuda_sources` and have an empty/matching `cpp_sources`. Or I could have the wrapper function in `cpp_sources` calling the kernel declared in `cuda_sources`.
Let me think more carefully. `load_inline` compiles both. Usually:
- `cpp_sources`: Contains the pybind11-free wrapper functions that will be auto-bound
- `cuda_sources`: Contains CUDA kernels and host functions that call them
If the wrapper function in `cpp_sources` calls a CUDA kernel, the kernel needs to be declared. But since they're compiled together, if the kernel is defined in `cuda_sources` and declared in `cpp_sources`, it should work.
Actually, looking at examples, a common pattern is:
```python
cpp_source = '''
torch::Tensor my_func(torch::Tensor a) {
// host-side code, can call CUDA kernels
...
}
'''
cuda_source = '''
__global__ void my_kernel(...) { ... }
'''
```
This should work because nvcc compiles both sources together. Wait no, `cuda_sources` is passed to nvcc and `cpp_sources` to c++. They might be compiled separately.
The correct approach is: put the `torch::Tensor` wrapper function in `cpp_sources` and have it call a `__global__` kernel declared in `cpp_sources` and defined in `cuda_sources`. Or, simpler: put everything CUDA-related in `cuda_sources` with `extern "C"` if needed, but since nvcc can compile C++ code with torch types...
Actually, I think the simplest approach that works is:
- Put the kernel definition in `cuda_sources`
- Put the host wrapper function that launches the kernel in `cpp_sources`
- Forward declare the kernel in `cpp_sources`
But wait, the kernel is `__global__`, so it can't be called from a .cpp file compiled by g++, it needs to be called from a .cu file compiled by nvcc. Actually, the `cpp_sources` might also be compiled by nvcc if it contains CUDA code...
Let me just use a single source file approach or put the host launcher in `cuda_sources` and declare it in `cpp_sources`.
Actually, looking at the PyTorch docs more carefully: `cpp_sources` is compiled with the C++ compiler, `cuda_sources` with NVCC. So `cpp_sources` cannot contain `__global__` or `__device__` code, but it can call `extern` CUDA functions if properly linked.
The standard pattern:
```python
cuda_source = '''
__global__ void kernel(...) { ... }
void launch_kernel(...) {
kernel<<<grid, block>>>(...);
}
'''
cpp_source = '''
void launch_kernel(...);
torch::Tensor my_func(torch::Tensor a) {
launch_kernel(...);
return a;
}
'''
```
But wait, `launch_kernel` defined in `cuda_sources` is compiled by nvcc and callable from `cpp_sources`. Hmm, but the declaration `void launch_kernel(...)` in `cpp_source` might not know about the definition in `cuda_source` at compile time... though the linker should resolve it.
Actually, for simplicity and to avoid linkage issues, let me put the `torch::Tensor` wrapper function directly in `cuda_sources` since nvcc can compile C++ code with libtorch:
```python
cuda_source = '''
#include <torch/extension.h>
__global__ void kernel(...) { ... }
torch::Tensor sparse_aggregate_cuda(...) {
kernel<<<...>>>(...);
return out;
}
'''
cpp_source = '''
#include <torch/extension.h>
torch::Tensor sparse_aggregate_cuda(...);
'''
```
Wait, but I need to include the same headers in both. Let me check if `load_inline` automatically adds `#include <torch/extension.h>`.
Actually, looking at PyTorch docs, `load_inline` does some automatic header inclusion. But to be safe, I'll include what we need in both sources or rely on defaults.
Hmm, I think the cleanest way is:
```python
cuda_source = '''
#include <torch/extension.h>
__global__ void sparse_aggregate_kernel(...) { ... }
torch::Tensor sparse_aggregate_cuda(torch::Tensor row_ptr, ...) {
...
sparse_aggregate_kernel<<<blocks, threads>>>(...);
return out;
}
'''
cpp_source = '''
#include <torch/extension.h>
torch::Tensor sparse_aggregate_cuda(torch::Tensor row_ptr, torch::Tensor col_idx,
torch::Tensor edge_weight, torch::Tensor node_feat);
'''
```
This should work. The function `sparse_aggregate_cuda` is defined in the CUDA compilation unit and declared in the C++ compilation unit. PyTorch's binding generator will use the declaration in `cpp_sources` to generate the pybind11 bindings.
Let me write the complete kernel and test it.
Actually, I want to double check: should `cpp_sources` contain only declarations, or can it also contain the wrapper? If it contains the wrapper but calls `<<< >>>`, then it needs to be compiled by nvcc. But `cpp_sources` is compiled by the C++ compiler by default. So the wrapper that uses `<<< >>>` should be in `cuda_sources`.
However, I just realized: actually in `torch.utils.cpp_extension.load_inline`, both `cpp_sources` and `cuda_sources` are passed to the same underlying `load` function. For mixed C++/CUDA extensions, the `cpp_sources` is indeed compiled with a C++ compiler, but some older versions or configurations might compile everything with nvcc. The safest pattern is:
- Put all CUDA-specific code (kernels and launchers) in `cuda_sources`
- Put declarations in `cpp_sources` for binding
Actually, another safe pattern that I've seen work:
```python
cuda_sources = '''
#include <torch/extension.h>
__global__ void kernel(...) { ... }
// Host wrapper callable from C++
torch::Tensor my_cuda_func(torch::Tensor x) {
kernel<<<...>>>(...);
return x;
}
'''
cpp_sources = '''
#include <torch/extension.h>
// Forward declaration
torch::Tensor my_cuda_func(torch::Tensor x);
'''
```
This should work because nvcc compiles `my_cuda_func` which contains CUDA-specific syntax. The C++ compiler sees only the declaration, and the linker resolves it.
Let me go with this pattern.
One more consideration: should I use `__restrict__`? Yes, it helps the compiler optimize.
Another optimization: since `head_dim = 32` and `num_heads = 8`, I could make the block size exactly `num_heads * head_dim = 256` and have one block per node. This means:
- Block 0 handles node 0, all heads and features
- Thread 0 handles (node 0, head 0, feature 0)
- Thread 1 handles (node 0, head 0, feature 1)
- ...
- Thread 31 handles (node 0, head 0, feature 31)
- Thread 32 handles (node 0, head 1, feature 0)
- etc.
This gives exact coalescing for all memory accesses. With one block per node, we have exactly 320 blocks. On H100 with 132 SMs, each SM gets ~2-3 blocks, and since each block has 256 threads, we get good utilization.
Wait, with 320 blocks and 132 SMs, each SM gets 2 or 3 blocks. 256 threads per block means 512-768 threads per SM. H100 supports 2048 threads per SM, so occupancy is ~25-37%. Is this enough? The kernel does very little compute per thread, so we might want higher occupancy.
Alternatively, we could process multiple nodes per block, or use 128 threads per block to get more blocks.
Actually, let me calculate occupancy more carefully. H100 has:
- 2048 threads per SM max
- 32 blocks per SM max
With 256 threads/block:
- Max blocks per SM = min(32, 2048/256) = 8 blocks per SM
- Total theoretical blocks = 132 * 8 = 1056
- We have 320 blocks, so ~2.4 blocks per SM
- Threads per SM = 2.4 * 256 = 614
- Occupancy = 614 / 2048 ≈ 30%
With 128 threads/block:
- Max blocks per SM = min(32, 2048/128) = 16
- Total theoretical = 2112
- We have 640 blocks, so ~4.85 blocks per SM
- Threads per SM = 4.85 * 128 = 621
- Occupancy = 621 / 2048 ≈ 30%
Hmm, actually comes out similar because we're just dividing the same total work.
Wait, but with 128 threads/block, we can have more blocks in flight if the work is divided such that more blocks can run concurrently. Actually no, the total number of threads is the same: 81920. Total SMs = 132. So average threads per SM is always 81920/132 ≈ 621 when all blocks are resident. The question is whether we can get all blocks resident.
With 256 threads/block: 320 blocks total. 320/132 = 2.42 blocks per SM. This is well within the 8 block limit. So all 320 blocks should be resident simultaneously. Occupancy = 320*256 / (132*2048) = 81920 / 270336 ≈ 30.3%.
The bottleneck for occupancy is likely not block count but register usage or shared memory. Since our kernel uses very few registers and no shared memory, we should be able to get full occupancy possible with 256 threads/block, which is 25% (since 2048 / 256 = 8 blocks per SM max, but if we have only ~2.4 blocks per SM on average, we get ~30%).
Actually wait, 30% occupancy for a memory-bound kernel might not be ideal. We want enough warps to hide memory latency. 621 threads per SM = ~19 warps per SM. For H100, 64 warps per SM max. So ~30% warp occupancy. H100 is quite good at latency hiding, but more warps would help.
How can we increase occupancy? We need more total threads or smaller blocks. The total threads is fixed at 81920 (unless we do redundant work). So occupancy is fundamentally limited by total_work / (num_SMs * max_threads_per_SM) = 81920 / 270336 ≈ 30%.
This means our kernel will run at ~30% occupancy no matter what (assuming no register/shared memory limits). That's actually fine for H100 on a small problem.
But wait - could we increase total parallelism by, say, having one thread per edge per head per feature? That would be 4480 * 8 * 32 ≈ 1.1M threads, which would give much higher occupancy. But then we'd need atomic adds to accumulate, which might be expensive.
Actually, let's reconsider. The edges average 14 per node. So total edges ≈ 4480. With 8 heads and 32 features, that's 1,146,880 edge-head-feature tuples. Launching 1.1M threads doing atomicAdd might actually be slower due to atomics.
But with such a low degree, maybe we can do something smarter. If degree is low (~14), we can load the edge weights and src indices into shared memory and do the computation efficiently.
Actually, let me think about this differently. Since each node has a small degree, if we assign one block per node with 256 threads (8 heads * 32 features), we can:
1. Have all threads in a block load the col_idx and edge_weight for that node's edges into shared memory
2. Then each thread computes its output by reading from shared memory
But degree is only ~14, so:
- col_idx: 14 * 4 bytes = 56 bytes
- edge_weight: 14 * 8 heads * 4 bytes = 448 bytes
Total shared memory: ~500 bytes per block - negligible!
This could significantly improve performance because:
1. Edge weights and column indices are loaded once per block and reused by all 256 threads
2. The `node_feat` loads are still from global memory but are coalesced
3. We eliminate redundant loads of `col_idx` and `edge_weight`
Actually, even better: we can load edge weights into shared memory. Each block processes one dst node. There are `degree` edges. For each edge, there are `num_heads` weights. Total weights = degree * num_heads.
With degree ~14 and num_heads = 8, that's 112 floats = 448 bytes in shared memory. Plus col_idx = 14 ints = 56 bytes. Very small.
Let me design a kernel with shared memory:
```cuda
__global__ void sparse_aggregate_kernel(
const int* __restrict__ row_ptr,
const int* __restrict__ col_idx,
const float* __restrict__ edge_weight,
const float* __restrict__ node_feat,
float* __restrict__ out,
int num_nodes,
int num_heads,
int head_dim
) {
__shared__ int s_col_idx[128]; // max degree? but we don't know max degree
__shared__ float s_edge_weight[128 * 8]; // too big if degree is large
```
Wait, I don't know the maximum degree at compile time. I can make it dynamic or use a large upper bound. But the problem says avg_degree = 14, and we can compute max degree. Actually, I can just allocate dynamically:
```cuda
extern __shared__ char shared_mem[];
int* s_col_idx = (int*)shared_mem;
float* s_edge_weight = (float*)&s_col_idx[degree];
```
Or use `__shared__` with a reasonable upper bound. Looking at the input generation:
```python
degree = torch.full((num_nodes,), avg_degree, dtype=torch.int32)
degree = torch.clamp(degree + ((torch.arange(num_nodes, dtype=torch.int32) % 5) - 2), min=1)
```
So degrees are in range [avg_degree - 2, avg_degree + 2] = [12, 16], clamped to min 1. So max degree is 16.
Wait, that's fixed! The degree for each node is:
- `avg_degree + ((i % 5) - 2) = 14 + ((i % 5) - 2) = 12 + (i % 5)`
- But then clamped to min 1
- So degrees are: 12, 13, 14, 15, 16, 12, 13, 14, 15, 16, ...
Maximum degree is 16. That's wonderful! We can use static shared memory with size 16 * 8 = 128 floats for edge weights, and 16 ints for col_idx.
But wait, the evaluator might use different inputs. The problem states specific values but the correctness check might use different random seeds or parameters. I should be safe and not hardcode max degree, OR I can query the max degree from row_ptr.
Actually, looking at the problem, `get_inputs()` generates inputs with those specific parameters, and the correctness check will use the same `get_inputs()`. But timing might use different sizes or distributions.
Hmm, to be safe, I'll use dynamic shared memory rather than hardcoding a bound. Or I can compute max degree from row_ptr on the CPU side and pass it as a template parameter or compute the shared memory size per block.
Actually, a simpler approach: I'll just write the kernel without shared memory but with good coalescing, and see if it's fast enough. Given the tiny problem size, it might already be very fast.
Let me do some back-of-the-envelope calculation:
- Memory reads:
- row_ptr: 321 ints = ~1.3 KB (negligible, cached)
- col_idx: ~4480 ints = ~18 KB
- edge_weight: 4480 * 8 = 35840 floats = ~143 KB
- node_feat: 320 * 8 * 32 = 81920 floats = ~327 KB
- Total feature reads per output: each output reads ~14 edges, each reading one float from node_feat and one float from edge_weight
- Total bytes read ≈ 14 * 81920 * 8 = 9.1 MB (for node_feat and edge_weight, ignoring col_idx)
At 3.35 TB/s bandwidth, 9.1 MB takes ~2.7 microseconds. The reference Python/PyTorch implementation will be much slower due to Python loop overhead and many small kernel launches. So even our simple kernel should give huge speedup.
Actually, wait. The PyTorch reference does `node_feat[src] * weights` where `src` is a tensor of indices. This calls `index_select` and broadcasting multiply, then `sum`. For each dst node, PyTorch launches small operations. The overhead of 320 iterations in Python + 320 kernel launches is probably the bottleneck.
Our single CUDA kernel will eliminate all that overhead. So a simple kernel will likely be 100x+ faster.
Let me just write the simple kernel first and test.
One more thing: I want to make sure I handle the data types correctly. The inputs are:
- row_ptr: torch.int32
- col_idx: torch.int32
- edge_weight: torch.float32
- node_feat: torch.float32
In my kernel, I'll use `int` for the index tensors and `float` for the float tensors.
Let me write the complete module and test it.