turn 01255.1s turn · 16.3s LLM · 1238.8s tools · tokens in=1,259 · out=1,272 · total=3,091 · FINAL
reasoning
The task is to replace the forward computation of a PyTorch model with a custom CUDA kernel. The model performs per-destination max aggregation over neighbor features in CSR format.
I need to replace the forward computation of a PyTorch model with a custom CUDA kernel for per-destination max aggregation over neighbor features in CSR format. The reference is a Python loop that finds neighbors using row_ptr and col_idx, computes max over neighbor features for each destination node.
For ModelNew, I'll use torch.
assistant```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 <cfloat>
__global__ void max_aggregate_kernel(
const int* __restrict__ row_ptr,
const int* __restrict__ col_idx,
const float* __restrict__ node_feat,
float* __restrict__ out,
int num_nodes,
int feat_dim
) {
int dst = blockIdx.x;
if (dst >= num_nodes) return;
int start = row_ptr[dst];
int end = row_ptr[dst + 1];
int degree = end - start;
if (degree == 0) {
for (int f = threadIdx.x; f < feat_dim; f += blockDim.x) {
out[dst * feat_dim + f] = -FLT_MAX;
}
return;
}
extern __shared__ float sdata[];
for (int f = threadIdx.x; f < feat_dim; f += blockDim.x) {
float local_max = -FLT_MAX;
for (int i = 0; i < degree; i++) {
int src = col_idx[start + i];
float val = node_feat[src * feat_dim + f];
if (val > local_max) local_max = val;
}
out[dst * feat_dim + f] = local_max;
}
}
torch::Tensor max_aggregate_cuda(
torch::Tensor row_ptr,
torch::Tensor col_idx,
torch::Tensor node_feat
) {
int num_nodes = row_ptr.size(0) - 1;
int feat_dim = node_feat.size(1);
auto out = torch::empty({num_nodes, feat_dim}, node_feat.options());
dim3 grid(num_nodes);
dim3 block(256);
int shared_mem = 0; // not used in this version
max_aggregate_kernel<<<grid, block, shared_mem>>>(
row_ptr.data_ptr<int>(),
col_idx.data_ptr<int>(),
node_feat.data_ptr<float>(),
out.data_ptr<float>(),
num_nodes,
feat_dim
);
return out;
}
'''
cpp_source = r'''
torch::Tensor max_aggregate_cuda(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor node_feat);
'''
module = load_inline(
name='max_aggregate',
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=['max_aggregate_cuda'],
verbose=False,
extra_cflags=['-O3'],
extra_ldflags=[],
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.module = module
def forward(
self,
row_ptr: torch.Tensor,
col_idx: torch.Tensor,
node_feat: torch.Tensor,
) -> torch.Tensor:
return self.module.max_aggregate_cuda(row_ptr, col_idx, node_feat)
```
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 <cfloat>\n\n__global__ void max_aggregate_kernel(\n const int* __restrict__ row_ptr,\n const int* __restrict__ col_idx,\n const float* __restrict__ node_feat,\n float* __restrict__ out,\n int num_nodes,\n int feat_dim\n) {\n int dst = blockIdx.x;\n if (dst >= num_nodes) return;\n\n int start = row_ptr[dst];\n int end = row_ptr[dst + 1];\n int degree = end - start;\n\n if (degree == 0) {\n for (int f = threadIdx.x; f < feat_dim; f += blockDim.x) {\n out[dst * feat_dim + f] = -FLT_MAX;\n }\n return;\n }\n\n for (int f = threadIdx.x; f < feat_dim; f += blockDim.x) {\n float local_max = -FLT_MAX;\n for (int i = 0; i < degree; i++) {\n int src = col_idx[start + i];\n float val = node_feat[src * feat_dim + f];\n if (val > local_max) local_max = val;\n }\n out[dst * feat_dim + f] = local_max;\n }\n}\n\ntorch::Tensor max_aggregate_cuda(\n torch::Tensor row_ptr,\n torch::Tensor col_idx,\n torch::Tensor node_feat\n) {\n int num_nodes = row_ptr.size(0) - 1;\n int feat_dim = node_feat.size(1);\n\n auto out = torch::empty({num_nodes, feat_dim}, node_feat.options());\n\n dim3 grid(num_nodes);\n dim3 block(256);\n int shared_mem = 0;\n\n max_aggregate_kernel<<<grid, block, shared_mem>>>(\n row_ptr.data_ptr<int>(),\n col_idx.data_ptr<int>(),\n node_feat.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 max_aggregate_cuda(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor node_feat);\n'''\n\nmodule = load_inline(\n name='max_aggregate',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['max_aggregate_cuda'],\n verbose=False,\n extra_cflags=['-O3'],\n extra_ldflags=[],\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.module = module\n\n def forward(\n self,\n row_ptr: torch.Tensor,\n col_idx: torch.Tensor,\n node_feat: torch.Tensor,\n ) -> torch.Tensor:\n return self.module.max_aggregate_cuda(row_ptr, col_idx, node_feat)\n"
}
result ← submit_kernel (fail)
failsubmit_kernel
submit_kernel FAILED: correctness check did not pass (? trials).
runtime_error: expected scalar type Int but found Float