reasoning
reasoning (truncated 32380 chars): Let me analyze this problem. We need to implement a GCN-style sparse aggregation with degree normalization in CUDA.
The reference implementation does:
1. For each destination node `dst`:
- Get neighbors from CSR format (row_ptr, col_idx)
- For each source node `src` in neighbors:
- Compute normalization: `1.0 / sqrt(deg[dst] * deg[src])`
- Aggregate: `out[dst] += node_feat[src] * norm`
This is essentially a sparse matrix-vector multiplication with a specific normalization pattern.
Let me write a CUDA kernel for this:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
# Define the CUDA kernel
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>
__global__ void gcn_aggregate_kernel(
const int* __restrict__ row_ptr,
const int* __restrict__ col_idx,
const float* __restrict__ node_feat,
const float* __restrict__ degrees,
float* __restrict__ out,
int num_nodes,
int feat_dim
) {
int dst = blockIdx.x;
int tid = threadIdx.x;
if (dst >= num_nodes) return;
int start = row_ptr[dst];
int end = row_ptr[dst + 1];
// Each thread processes a subset of the feature dimension
for (int f = tid; f < feat_dim; f += blockDim.x) {
float sum = 0.0f;
for (int i = start; i < end; i++) {
int src = col_idx[i];
float norm = 1.0f / sqrtf(degrees[dst] * degrees[src]);
sum += node_feat[src * feat_dim + f] * norm;
}
out[dst * feat_dim + f] = sum;
}
}
torch::Tensor gcn_aggregate_cuda(
torch::Tensor row_ptr,
torch::Tensor col_idx,
torch::Tensor node_feat,
torch::Tensor degrees
) {
int num_nodes = row_ptr.numel() - 1;
int feat_dim = node_feat.size(1);
auto out = torch::zeros({num_nodes, feat_dim}, node_feat.options());
const int threads = 256;
const int blocks = num_nodes;
gcn_aggregate_kernel<<<blocks, threads>>>(
row_ptr.data_ptr<int>(),
col_idx.data_ptr<int>(),
node_feat.data_ptr<float>(),
degrees.data_ptr<float>(),
out.data_ptr<float>(),
num_nodes,
feat_dim
);
return out;
}
"""
cpp_source = """
torch::Tensor gcn_aggregate_cuda(
torch::Tensor row_ptr,
torch::Tensor col_idx,
torch::Tensor node_feat,
torch::Tensor degrees
);
"""
# Load the extension
gcn_ext = load_inline(
name='gcn_aggregate',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['gcn_aggregate_cuda'],
extra_cuda_cflags=['-O3'],
verbose=True
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(
self,
row_ptr: torch.Tensor,
col_idx: torch.Tensor,
node_feat: torch.Tensor,
degrees: torch.Tensor,
) -> torch.Tensor:
return gcn_ext.gcn_aggregate_cuda(row_ptr, col_idx, node_feat, degrees)
```
Wait, I need to be careful about the data types. The row_ptr and col_idx are int32, but CUDA typically uses int for indexing. Let me check the types again...
Looking at the reference:
- `row_ptr`: torch.int32
- `col_idx`: torch.int32
- `node_feat`: torch.float32
- `degrees`: torch.float32
I should use int32 pointers in the kernel. Let me also consider optimizations:
1. **Thread block assignment**: One block per destination node
2. **Thread assignment**: Each thread handles multiple features (stride loop)
3. **Shared memory**: Could precompute norms for each neighbor, but for avg_degree=16, overhead might not be worth it
Let me also think about using shared memory for the normalization factors since they're reused across all feature dimensions.
Actually, let me reconsider the kernel design:
- num_nodes = 640
- avg_degree = 16
- feat_dim = 80
With 256 threads per block and 80 features, we have more threads than work if each thread handles one feature. So the stride loop is fine.
But I could also use shared memory to store the normalization factors to avoid recomputing them for each feature dimension.
Let me optimize further:
```cuda
__global__ void gcn_aggregate_kernel(
const int* __restrict__ row_ptr,
const int* __restrict__ col_idx,
const float* __restrict__ node_feat,
const float* __restrict__ degrees,
float* __restrict__ out,
int num_nodes,
int feat_dim,
int max_degree // for shared memory size
) {
int dst = blockIdx.x;
int tid = threadIdx.x;
if (dst >= num_nodes) return;
int start = row_ptr[dst];
int end = row_ptr[dst + 1];
int degree = end - start;
extern __shared__ float shared_norms[];
// Compute norms in shared memory
for (int i = tid; i < degree; i += blockDim.x) {
int src = col_idx[start + i];
shared_norms[i] = 1.0f / sqrtf(degrees[dst] * degrees[src]);
}
__syncthreads();
// Each thread processes a subset of the feature dimension
for (int f = tid; f < feat_dim; f += blockDim.x) {
float sum = 0.0f;
for (int i = 0; i < degree; i++) {
int src = col_idx[start + i];
sum += node_feat[src * feat_dim + f] * shared_norms[i];
}
out[dst * feat_dim + f] = sum;
}
}
```
But wait, with avg_degree=16, the shared memory usage is small. The real question is whether this is faster than recomputing.
Actually, for such small degrees, recomputing the sqrt might be fine. Let me think about the memory access patterns...
Another approach: use warp-level parallelism to reduce across neighbors. But with varying degrees, this gets complex.
Let me try a simpler approach first - just parallelize over destination nodes and feature dimensions:
Actually, let me think about the data access pattern more carefully:
- node_feat is accessed with `src * feat_dim + f`
- For each (dst, f), we access `degree` consecutive elements from col_idx and corresponding rows from node_feat
The node_feat accesses are strided, which could be slow. But with feat_dim=80 and coalesced access within each thread block, it should be okay.
Let me also consider using more threads for feature parallelism. With 80 features and 256 threads, we can have good coverage.
Let me also use vectorized loads if possible. But for simplicity, let me start with a basic kernel and see if it works.
Here's my implementation plan:
1. Basic kernel: one block per node, threads parallelize over features
2. Use shared memory for normalization factors
3. Make sure to handle empty neighborhoods
Let me write the complete solution:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
__global__ void gcn_aggregate_kernel(
const int* __restrict__ row_ptr,
const int* __restrict__ col_idx,
const float* __restrict__ node_feat,
const float* __restrict__ degrees,
float* __restrict__ out,
int num_nodes,
int feat_dim
) {
const int dst = blockIdx.x;
const int tid = threadIdx.x;
const int block_size = blockDim.x;
if (dst >= num_nodes) return;
const int start = row_ptr[dst];
const int end = row_ptr[dst + 1];
const int degree = end - start;
if (degree == 0) {
// No neighbors, output is already zero
for (int f = tid; f < feat_dim; f += block_size) {
out[dst * feat_dim + f] = 0.0f;
}
return;
}
// Shared memory for normalization factors
extern __shared__ float shared_data[];
float* norms = shared_data;
// Compute normalization factors cooperatively
for (int i = tid; i < degree; i += block_size) {
const int src = col_idx[start + i];
norms[i] = rsqrtf(degrees[dst] * degrees[src]); // 1/sqrt(x)
}
__syncthreads();
// Aggregate features
for (int f = tid; f < feat_d
...[ omitted 32380 characters ]...
odes) return;
const int start = row_ptr[dst];
const int end = row_ptr[dst + 1];
const int degree = end - start;
if (degree == 0) return;
const float dst_deg = degrees[dst];
// Shared memory for source indices and norms
__shared__ int shared_src[128];
__shared__ float shared_norm[128];
// Cooperatively load data into shared memory
for (int i = tid; i < degree; i += nthreads) {
const int src = col_idx[start + i];
shared_src[i] = src;
shared_norm[i] = 1.0f / sqrtf(dst_deg * degrees[src]);
}
__syncthreads();
// Process features - each thread handles one or more feature dimensions
for (int f = tid; f < feat_dim; f += nthreads) {
float sum = 0.0f;
for (int i = 0; i < degree; i++) {
const int src = shared_src[i];
sum += node_feat[src * feat_dim + f] * shared_norm[i];
}
out[dst * feat_dim + f] = sum;
}
}
torch::Tensor gcn_aggregate_cuda(
torch::Tensor row_ptr,
torch::Tensor col_idx,
torch::Tensor node_feat,
torch::Tensor degrees
) {
row_ptr = row_ptr.contiguous();
col_idx = col_idx.contiguous();
node_feat = node_feat.contiguous();
degrees = degrees.contiguous();
const int num_nodes = row_ptr.numel() - 1;
const int feat_dim = node_feat.size(1);
auto out = torch::zeros({num_nodes, feat_dim}, node_feat.options());
const int threads = 128;
const int blocks = num_nodes;
gcn_aggregate_kernel<<<blocks, threads>>>(
row_ptr.data_ptr<int>(),
col_idx.data_ptr<int>(),
node_feat.data_ptr<float>(),
degrees.data_ptr<float>(),
out.data_ptr<float>(),
num_nodes,
feat_dim
);
return out;
}
"""
cpp_source = r"""
torch::Tensor gcn_aggregate_cuda(
torch::Tensor row_ptr,
torch::Tensor col_idx,
torch::Tensor node_feat,
torch::Tensor degrees
);
"""
gcn_ext = load_inline(
name='gcn_aggregate',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['gcn_aggregate_cuda'],
extra_cuda_cflags=['-O3'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(
self,
row_ptr: torch.Tensor,
col_idx: torch.Tensor,
node_feat: torch.Tensor,
degrees: torch.Tensor,
) -> torch.Tensor:
return gcn_ext.gcn_aggregate_cuda(row_ptr, col_idx, node_feat, degrees)
```
Wait, I should check the CUDA error and add cudaDeviceSynchronize to ensure the kernel completes before returning. But for performance measurement, this might not be necessary if the framework handles it.
Actually, torch.utils.cpp_extension should handle synchronization automatically. Let me not add explicit synchronization.
Let me also think about potential issues:
1. Empty neighborhoods: handled
2. Very large degrees: handled with 128-element shared memory
3. Type mismatches: int32 for indices, float for values
I think this is good. Let me submit.
Actually, let me also consider a simpler version without shared memory to see if it's faster. The shared memory adds a synchronization overhead, and with only 16 neighbors on average, the recomputation cost is small.
Let me write both versions and pick the simpler one:
```cuda
__global__ void gcn_aggregate_kernel_simple(
const int* __restrict__ row_ptr,
const int* __restrict__ col_idx,
const float* __restrict__ node_feat,
const float* __restrict__ degrees,
float* __restrict__ out,
int num_nodes,
int feat_dim
) {
const int dst = blockIdx.x;
const int tid = threadIdx.x;
const int nthreads = blockDim.x;
if (dst >= num_nodes) return;
const int start = row_ptr[dst];
const int end = row_ptr[dst + 1];
const int degree = end - start;
if (degree == 0) return;
const float dst_deg = degrees[dst];
// Process features
for (int f = tid; f < feat_dim; f += nthreads) {
float sum = 0.0f;
for (int i = 0; i < degree; i++) {
const int src = col_idx[start + i];
const float norm = 1.0f / sqrtf(dst_deg * degrees[src]);
sum += node_feat[src * feat_dim + f] * norm;
}
out[dst * feat_dim + f] = sum;
}
}
```
The simple version:
- No shared memory
- No __syncthreads()
- Recomputes norm for each feature
With 16 neighbors and 80 features:
- 16 * 80 = 1280 sqrt computations per node
- Each sqrt is about 10-20 cycles
- Total: 12800-25600 cycles per node
This might be faster or slower depending on memory latency hiding.
I'll go with the shared memory version since it avoids redundant sqrt computations and the __syncthreads overhead is minimal for small degrees.
Let me finalize:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
__global__ void __launch_bounds__(128) gcn_aggregate_kernel(
const int* __restrict__ row_ptr,
const int* __restrict__ col_idx,
const float* __restrict__ node_feat,
const float* __restrict__ degrees,
float* __restrict__ out,
int num_nodes,
int feat_dim
) {
const int dst = blockIdx.x;
const int tid = threadIdx.x;
const int nthreads = blockDim.x;
if (dst >= num_nodes) return;
const int start = row_ptr[dst];
const int end = row_ptr[dst + 1];
const int degree = end - start;
if (degree == 0) return;
const float dst_deg = degrees[dst];
// Shared memory for source indices and norms
__shared__ int shared_src[128];
__shared__ float shared_norm[128];
// Cooperatively load data into shared memory
for (int i = tid; i < degree; i += nthreads) {
const int src = col_idx[start + i];
shared_src[i] = src;
shared_norm[i] = 1.0f / sqrtf(dst_deg * degrees[src]);
}
__syncthreads();
// Process features - each thread handles one or more feature dimensions
for (int f = tid; f < feat_dim; f += nthreads) {
float sum = 0.0f;
for (int i = 0; i < degree; i++) {
const int src = shared_src[i];
sum += node_feat[src * feat_dim + f] * shared_norm[i];
}
out[dst * feat_dim + f] = sum;
}
}
torch::Tensor gcn_aggregate_cuda(
torch::Tensor row_ptr,
torch::Tensor col_idx,
torch::Tensor node_feat,
torch::Tensor degrees
) {
row_ptr = row_ptr.contiguous();
col_idx = col_idx.contiguous();
node_feat = node_feat.contiguous();
degrees = degrees.contiguous();
const int num_nodes = row_ptr.numel() - 1;
const int feat_dim = node_feat.size(1);
auto out = torch::zeros({num_nodes, feat_dim}, node_feat.options());
const int threads = 128;
const int blocks = num_nodes;
gcn_aggregate_kernel<<<blocks, threads>>>(
row_ptr.data_ptr<int>(),
col_idx.data_ptr<int>(),
node_feat.data_ptr<float>(),
degrees.data_ptr<float>(),
out.data_ptr<float>(),
num_nodes,
feat_dim
);
return out;
}
"""
cpp_source = r"""
torch::Tensor gcn_aggregate_cuda(
torch::Tensor row_ptr,
torch::Tensor col_idx,
torch::Tensor node_feat,
torch::Tensor degrees
);
"""
gcn_ext = load_inline(
name='gcn_aggregate',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['gcn_aggregate_cuda'],
extra_cuda_cflags=['-O3'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(
self,
row_ptr: torch.Tensor,
col_idx: torch.Tensor,
node_feat: torch.Tensor,
degrees: torch.Tensor,
) -> torch.Tensor:
return gcn_ext.gcn_aggregate_cuda(row_ptr, col_idx, node_feat, degrees)
```
This looks good. Let me submit.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = r\"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n__global__ void __launch_bounds__(128) gcn_aggregate_kernel(\n const int* __restrict__ row_ptr,\n const int* __restrict__ col_idx,\n const float* __restrict__ node_feat,\n const float* __restrict__ degrees,\n float* __restrict__ out,\n int num_nodes,\n int feat_dim\n) {\n const int dst = blockIdx.x;\n const int tid = threadIdx.x;\n const int nthreads = blockDim.x;\n \n if (dst >= num_nodes) return;\n \n const int start = row_ptr[dst];\n const int end = row_ptr[dst + 1];\n const int degree = end - start;\n \n if (degree == 0) return;\n \n const float dst_deg = degrees[dst];\n \n // Shared memory for source indices and norms\n __shared__ int shared_src[128];\n __shared__ float shared_norm[128];\n \n // Cooperatively load data into shared memory\n for (int i = tid; i < degree; i += nthreads) {\n const int src = col_idx[start + i];\n shared_src[i] = src;\n shared_norm[i] = 1.0f / sqrtf(dst_deg * degrees[src]);\n }\n __syncthreads();\n \n // Process features - each thread handles one or more feature dimensions\n for (int f = tid; f < feat_dim; f += nthreads) {\n float sum = 0.0f;\n for (int i = 0; i < degree; i++) {\n const int src = shared_src[i];\n sum += node_feat[src * feat_dim + f] * shared_norm[i];\n }\n out[dst * feat_dim + f] = sum;\n }\n}\n\ntorch::Tensor gcn_aggregate_cuda(\n torch::Tensor row_ptr,\n torch::Tensor col_idx,\n torch::Tensor node_feat,\n torch::Tensor degrees\n) {\n row_ptr = row_ptr.contiguous();\n col_idx = col_idx.contiguous();\n node_feat = node_feat.contiguous();\n degrees = degrees.contiguous();\n \n const int num_nodes = row_ptr.numel() - 1;\n const int feat_dim = node_feat.size(1);\n \n auto out = torch::zeros({num_nodes, feat_dim}, node_feat.options());\n \n const int threads = 128;\n const int blocks = num_nodes;\n \n gcn_aggregate_kernel<<<blocks, threads>>>(\n row_ptr.data_ptr<int>(),\n col_idx.data_ptr<int>(),\n node_feat.data_ptr<float>(),\n degrees.data_ptr<float>(),\n out.data_ptr<float>(),\n num_nodes,\n feat_dim\n );\n \n return out;\n}\n\"\"\"\n\ncpp_source = r\"\"\"\ntorch::Tensor gcn_aggregate_cuda(\n torch::Tensor row_ptr,\n torch::Tensor col_idx,\n torch::Tensor node_feat,\n torch::Tensor degrees\n);\n\"\"\"\n\ngcn_ext = load_inline(\n name='gcn_aggregate',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gcn_aggregate_cuda'],\n extra_cuda_cflags=['-O3'],\n verbose=False\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 node_feat: torch.Tensor,\n degrees: torch.Tensor,\n ) -> torch.Tensor:\n return gcn_ext.gcn_aggregate_cuda(row_ptr, col_idx, node_feat, degrees)"
}