correct
1_Conv2D_ReLU_BiasAdd.py
started
2026-05-02T04:40:17Z
timingkernelbench.timing
kernel mean / median / std
2.73 / — / 0.88 μs
ref mean / median / std
4.38 / — / 0.04 μs
numerical precisioneval_kernel_against_ref
gpu memoryextended_metrics.measure_memory
kernel launches / fusiontorch.profiler
fusion ratio (ref / kernel)
2.50x
energyNVML (pynvml)
ratio (ref / kernel)
2.00x
SOL (speed-of-light)nsight_hardware_counters
roofline / occupancynsight_hardware_counters
dram bandwidth
6737.14 GB/s
warp stalls (top 5)nsight
reference PyTorch source (the prompt the agent sees)
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Simple model that performs a convolution, applies ReLU, and adds a bias term.
"""
def __init__(self, in_channels, out_channels, kernel_size, bias_shape):
super(Model, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
self.bias = nn.Parameter(torch.randn(bias_shape))
def forward(self, x):
x = self.conv(x)
x = torch.relu(x)
x = x + self.bias
return x
batch_size = 128
in_channels = 64
out_channels = 128
height = width = 128
kernel_size = 3
bias_shape = (out_channels, 1, 1)
def get_inputs():
return [torch.rand(batch_size, in_channels, height, width)]
def get_init_inputs():
return [in_channels, out_channels, kernel_size, bias_shape]turn 010.6s turn · 10.6s LLM · 0.0s tools
reasoning
**Designing ModelNew**
I'm focusing on crafting ModelNew with a Conv2d layer: 128 batch size, 64 input channels, and 128 output channels, using 128x128 input and a kernel size of 3, resulting in an output of 126x126. I need the constructor to create the Conv2d layer and bias parameters similarly to Model. For the forward pass, I might implement a custom solution that's faster than PyTorch's reference. The convolutions demand significant operations, potentially up to 30 GFLOPs. I'll consider using a custom CUDA kernel for efficiency.
**Implementing Custom Convolution**
I'm considering an approach where we use the `at::conv2d` for convolution, followed by an in-place operation for ReLU and bias addition. This should help avoid unnecessary intermediate steps and overhead from PyTorch, which seems likely to be faster. I need to confirm if `at::conv2d` is callable in C++ and link it correctly.
I’ll implement a function `fused_forward` that processes the input and returns the modified tensor. The kernel may face an issue since it could call `aten`, not utilizing cuDNN. I’ll include the appropriate headers and ensure everything links properly.
call → get_gpu_specs
callget_gpu_specs
{}
result ← get_gpu_specs (ok)
okget_gpu_specs
GPU specs for NVIDIA H100 80GB HBM3:
total memory (runtime): 79.1 GB
spec entry: H100
GPU Architecture: Hopper
GPU Memory: 80GB
Memory Bandwidth: 3.35 TB/s
FP64 TFLOPS: 34
FP64 Tensor Core TFLOPS: 67
FP32 TFLOPS: 67
TF32 Tensor Core TFLOPS: 989 with sparsity
BFLOAT16 Tensore Core TFLOPS: 1979 with sparsity
FP16 Tensor Core TFLOPS: 1979 with sparsity
FP8 Tensor Core TFLOPS: 3958 with sparsity
INT8 Tensor Core TOPS: 3958 with sparsity
Register File Size: 64K 32-bit registers per SM
Maximum number of registers per thread: 255
Maximum number of thread blocks per SM: 32
Shared memory capacity per SM: 228 KB
Maximum shared memory per thread block: 227 KB
turn 1 · 1 tool call(s) · 27.7s turn · 27.7s LLM · 11,753 chars (click to expand)
turn 127.7s turn · 27.7s LLM · 0.0s tools
reasoning
**Assessing kernel efficiency**
It seems simpler kernels could efficiently handle scalar cases, but I'm wondering if vector cases might not be necessary. Since the output is 1GB, vectors might enhance performance. We could write vector kernel launches when certain conditions like total%4==0 are met. Tensors likely are well aligned, but I'm considering that CUDA host wrappers use properties like y.size. I need to balance that against potential launch overhead versus memory use. I think benchmarking with cuDNN might help clarify this along with correctness checks.
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r'''\n#include <torch/extension.h>\n#include <vector>\n\ntorch::Tensor apply_relu_bias_cuda(torch::Tensor y, torch::Tensor bias);\n\ntorch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add_bias) {\n c10::optional<at::Tensor> b = conv_bias;\n auto y = at::conv2d(input, weight, b, {1,1}, {0,0}, {1,1}, 1);\n return apply_relu_bias_cuda(y, add_bias);\n}\n'''\n\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void relu_bias_scalar_kernel(float* __restrict__ y, const float* __restrict__ bias,\n long total, int C, int HW) {\n long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < total) {\n float v = y[idx];\n int c = (idx / HW) % C;\n y[idx] = (v > 0.0f ? v : 0.0f) + bias[c];\n }\n}\n\n__global__ void relu_bias_vec4_kernel(float4* __restrict__ y4, const float* __restrict__ bias,\n long total4, int C, int HW4) {\n long i = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (i < total4) {\n float4 v = y4[i];\n int c = (i / HW4) % C;\n float b = bias[c];\n v.x = (v.x > 0.0f ? v.x : 0.0f) + b;\n v.y = (v.y > 0.0f ? v.y : 0.0f) + b;\n v.z = (v.z > 0.0f ? v.z : 0.0f) + b;\n v.w = (v.w > 0.0f ? v.w : 0.0f) + b;\n y4[i] = v;\n }\n}\n\ntorch::Tensor apply_relu_bias_cuda(torch::Tensor y, torch::Tensor bias) {\n const long total = y.numel();\n if (total == 0) return y;\n const int C = (int)y.size(1);\n const int HW = (int)(y.size(2) * y.size(3));\n const int threads = 256;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if ((total % 4 == 0) && (HW % 4 == 0)) {\n long total4 = total / 4;\n int HW4 = HW / 4;\n int blocks = (int)((total4 + threads - 1) / threads);\n relu_bias_vec4_kernel<<<blocks, threads, 0, stream>>>(reinterpret_cast<float4*>(y.data_ptr<float>()), bias.data_ptr<float>(), total4, C, HW4);\n } else {\n int blocks = (int)((total + threads - 1) / threads);\n relu_bias_scalar_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), bias.data_ptr<float>(), total, C, HW);\n }\n return y;\n}\n'''\n\n_ext = load_inline(name='conv_relu_bias_fused_ext_v1',\n cpp_sources=[_cpp_src],\n cuda_sources=[_cuda_src],\n functions=['fused_forward'],\n extra_cuda_cflags=['-O3'],\n extra_cflags=['-O3'],\n verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, bias_shape):\n super(ModelNew, self).__init__()\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)\n self.bias = nn.Parameter(torch.randn(bias_shape))\n\n def forward(self, x):\n return _ext.fused_forward(x, self.conv.weight, self.conv.bias, self.bias)\n"
}
result ← compile_kernel (fail)
failcompile_kernel
compile_kernel FAILED: RuntimeError.
RuntimeError: Error building extension 'conv_relu_bias_fused_ext_v1': [1/3] /opt/conda/bin/x86_64-conda-linux-gnu-c++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=conv_relu_bias_fused_ext_v1 -DTORCH_API_INCLUDE_EXTENSION_H -isystem /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include -isystem /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /home/prabhune/.local/share/uv/python/cpython-3.10-linux-x86_64-gnu/include/python3.10 -fPIC -std=c++17 -O3 -c /scratch/tejas/PopcornBench/runs/l12_default/original/gpt-5.5-priority/level_2_problem_1_cache/k_c4e5f2f468ac/conv_relu_bias_fused_ext_v1/main.cpp -o main.o
FAILED: [code=1] main.o
/opt/conda/bin/x86_64-conda-linux-gnu-c++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=conv_relu_bias_fused_ext_v1 -DTORCH_API_INCLUDE_EXTENSION_H -isystem /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include -isystem /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /home/prabhune/.local/share/uv/python/cpython-3.10-linux-x86_64-gnu/include/python3.10 -fPIC -std=c++17 -O3 -c /scratch/tejas/PopcornBench/runs/l12_default/original/gpt-5.5-priority/level_2_problem_1_cache/k_c4e5f2f468ac/conv_relu_bias_fused_ext_v1/main.cpp -o main.o
/scratch/tejas/PopcornBench/runs/l12_default/original/gpt-5.5-priority/level_2_problem_1_cache/k_c4e5f2f468ac/conv_relu_bias_fused_ext_v1/main.cpp: In function 'at::Tensor fused_forward(at::Tensor, at::Tensor, at::Tensor, at::Tensor)':
/scratch/tejas/PopcornBench/runs/l12_default/original/gpt-5.5-priority/level_2_problem_1_cache/k_c4e5f2f468ac/conv_relu_bias_fused_ext_v1/main.cpp:10:24: error: call of overloaded 'conv2d(at::Tensor&, at::Tensor&, std::optional<at::Tensor>&, <brace-enclosed initializer list>, <brace-enclosed initializer list>, <brace-enclosed initializer list>, int)' is ambiguous
10 | auto y = at::conv2d(input, weight, b, {1,1}, {0,0}, {1,1}, 1);
| ~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
In file included from /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/ATen/Functions.h:607,
from /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/ATen/ExpandUtils.h:5,
from /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/torch/csrc/autograd/input_metadata.h:4,
from /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/torch/csrc/autograd/function.h:8,
from /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/torch/csrc/autograd/custom_function.h:8,
from /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/torch/csrc/api/include/torch/autograd.h:6,
from /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/torch/csrc/api/include/torch/all.h:8,
from /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/torch/extension.h:6,
from /scratch/tejas/PopcornBench/runs/l12_default/original/gpt-5.5-priority/level_2_problem_1_cache/k_c4e5f2f468ac/conv_relu_bias_fused_ext_v1/main.cpp:1:
/scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/ATen/ops/conv2d.h:28:19: note: candidate: 'at::Tensor at::conv2d(const Tensor&, const Tensor&, const std::optional<Tensor>&, IntArrayRef, IntArrayRef, IntArrayRef, int64_t)'
28 | inline at::Tensor conv2d(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias={}, at::IntArrayRef stride=1, at::IntArrayRef padding=0, at::IntArrayRef dilation=1, int64_t groups=1) {
| ^~~~~~
/scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/ATen/ops/conv2d.h:50:19: note: candidate: 'at::Tensor at::conv2d(const Tensor&, const Tensor&, const std::optional<Tensor>&, IntArrayRef, c10::string_view, IntArrayRef, int64_t)'
50 | inline at::Tensor conv2d(const at::Tensor & input, const at::Tensor & weight, const ::std::optional<at::Tensor> & bias, at::IntArrayRef stride, c10::string_view padding, at::IntArrayRef dilation=1, int64_t groups=1) {
| ^~~~~~
[2/3] /usr/local/cuda/bin/nvcc -MD -MF cuda.cuda.o.d -ccbin /opt/conda/bin/x86_64-conda-linux-gnu-cc -DTORCH_EXTENSION_NAME=conv_relu_bias_fused_ext_v1 -DTORCH_API_INCLUDE_EXTENSION_H -isystem /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include -isystem /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /home/prabhune/.local/share/uv/python/cpython-3.10-linux-x86_64-gnu/include/python3.10 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_90,code=compute_90 -gencode=arch=compute_90,code=sm_90 --compiler-options '-fPIC' -O3 -std=c++17 -c /scratch/tejas/PopcornBench/runs/l12_default/original/gpt-5.5-priority/level_2_problem_1_cache/k_c4e5f2f468ac/conv_relu_bias_fused_ext_v1/cuda.cu -o cuda.cuda.o
FAILED: [code=2] cuda.cuda.o
/usr/local/cuda/bin/nvcc -MD -MF cuda.cuda.o.d -ccbin /opt/conda/bin/x86_64-conda-linux-gnu-cc -DTORCH_EXTENSION_NAME=conv_relu_bias_fused_ext_v1 -DTORCH_API_INCLUDE_EXTENSION_H -isystem /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include -isystem /scratch/tejas/PopcornBench/.venv/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /home/prabhune/.local/share/uv/python/cpython-3.10-linux-x86_64-gnu/include/python3.10 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_90,code=compute_90 -gencode=arch=compute_90,code=sm_90 --compiler-options '-fPIC' -O3 -std=c++17 -c /scratch/tejas/PopcornBench/runs/l12_default/original/gpt-5.5-priority/level_2_problem_1_cache/k_c4e5f2f468ac/conv_relu_bias_fused_ext_v1/cuda.cu -o cuda.cuda.o
nvcc warning : incompatible redefinition for option 'compiler-bindir', the last value of this option was used
/scratch/tejas/PopcornBench/runs/l12_default/original/gpt-5.5-priority/level_2_problem_1_cache/k_c4e5f2f468ac/conv_relu_bias_fused_ext_v1/cuda.cu(40): error: namespace "at::cuda" has no member "getCurrentCUDAStream"
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
^
1 error detected in the compilation of "/scratch/tejas/PopcornBench/runs/l12_default/original/gpt-5.5-priority/level_2_problem_1_cache/k_c4e5f2f468ac/conv_relu_bias_fused_ext_v1/cuda.cu".
ninja: build stopped: subcommand failed.
turn 29.4s turn · 9.4s LLM · 0.0s tools
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r'''\n#include <torch/extension.h>\n#include <vector>\n\ntorch::Tensor apply_relu_bias_cuda(torch::Tensor y, torch::Tensor bias);\n\ntorch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add_bias) {\n std::optional<at::Tensor> b = conv_bias;\n std::vector<int64_t> stride = {1, 1};\n std::vector<int64_t> padding = {0, 0};\n std::vector<int64_t> dilation = {1, 1};\n auto y = at::conv2d(input, weight, b, stride, padding, dilation, 1);\n return apply_relu_bias_cuda(y, add_bias);\n}\n'''\n\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void relu_bias_scalar_kernel(float* __restrict__ y, const float* __restrict__ bias,\n long total, int C, int HW) {\n long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < total) {\n float v = y[idx];\n int c = (idx / HW) % C;\n y[idx] = (v > 0.0f ? v : 0.0f) + bias[c];\n }\n}\n\n__global__ void relu_bias_vec4_kernel(float4* __restrict__ y4, const float* __restrict__ bias,\n long total4, int C, int HW4) {\n long i = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (i < total4) {\n float4 v = y4[i];\n int c = (i / HW4) % C;\n float b = bias[c];\n v.x = (v.x > 0.0f ? v.x : 0.0f) + b;\n v.y = (v.y > 0.0f ? v.y : 0.0f) + b;\n v.z = (v.z > 0.0f ? v.z : 0.0f) + b;\n v.w = (v.w > 0.0f ? v.w : 0.0f) + b;\n y4[i] = v;\n }\n}\n\ntorch::Tensor apply_relu_bias_cuda(torch::Tensor y, torch::Tensor bias) {\n const long total = y.numel();\n if (total == 0) return y;\n const int C = (int)y.size(1);\n const int HW = (int)(y.size(2) * y.size(3));\n const int threads = 256;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if ((total % 4 == 0) && (HW % 4 == 0)) {\n long total4 = total / 4;\n int HW4 = HW / 4;\n int blocks = (int)((total4 + threads - 1) / threads);\n relu_bias_vec4_kernel<<<blocks, threads, 0, stream>>>(reinterpret_cast<float4*>(y.data_ptr<float>()), bias.data_ptr<float>(), total4, C, HW4);\n } else {\n int blocks = (int)((total + threads - 1) / threads);\n relu_bias_scalar_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), bias.data_ptr<float>(), total, C, HW);\n }\n return y;\n}\n'''\n\n_ext = load_inline(name='conv_relu_bias_fused_ext_v2',\n cpp_sources=[_cpp_src],\n cuda_sources=[_cuda_src],\n functions=['fused_forward'],\n extra_cuda_cflags=['-O3'],\n extra_cflags=['-O3'],\n verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, bias_shape):\n super(ModelNew, self).__init__()\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)\n self.bias = nn.Parameter(torch.randn(bias_shape))\n\n def forward(self, x):\n return _ext.fused_forward(x, self.conv.weight, self.conv.bias, self.bias)\n"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 37.8s turn · 7.8s LLM · 0.0s tools
call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r'''\n#include <torch/extension.h>\n#include <vector>\n\ntorch::Tensor apply_relu_bias_cuda(torch::Tensor y, torch::Tensor bias);\n\ntorch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add_bias) {\n std::optional<at::Tensor> b = conv_bias;\n std::vector<int64_t> stride = {1, 1};\n std::vector<int64_t> padding = {0, 0};\n std::vector<int64_t> dilation = {1, 1};\n auto y = at::conv2d(input, weight, b, stride, padding, dilation, 1);\n return apply_relu_bias_cuda(y, add_bias);\n}\n'''\n\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void relu_bias_scalar_kernel(float* __restrict__ y, const float* __restrict__ bias,\n long total, int C, int HW) {\n long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < total) {\n float v = y[idx];\n int c = (idx / HW) % C;\n y[idx] = (v > 0.0f ? v : 0.0f) + bias[c];\n }\n}\n\n__global__ void relu_bias_vec4_kernel(float4* __restrict__ y4, const float* __restrict__ bias,\n long total4, int C, int HW4) {\n long i = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (i < total4) {\n float4 v = y4[i];\n int c = (i / HW4) % C;\n float b = bias[c];\n v.x = (v.x > 0.0f ? v.x : 0.0f) + b;\n v.y = (v.y > 0.0f ? v.y : 0.0f) + b;\n v.z = (v.z > 0.0f ? v.z : 0.0f) + b;\n v.w = (v.w > 0.0f ? v.w : 0.0f) + b;\n y4[i] = v;\n }\n}\n\ntorch::Tensor apply_relu_bias_cuda(torch::Tensor y, torch::Tensor bias) {\n const long total = y.numel();\n if (total == 0) return y;\n const int C = (int)y.size(1);\n const int HW = (int)(y.size(2) * y.size(3));\n const int threads = 256;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if ((total % 4 == 0) && (HW % 4 == 0)) {\n long total4 = total / 4;\n int HW4 = HW / 4;\n int blocks = (int)((total4 + threads - 1) / threads);\n relu_bias_vec4_kernel<<<blocks, threads, 0, stream>>>(reinterpret_cast<float4*>(y.data_ptr<float>()), bias.data_ptr<float>(), total4, C, HW4);\n } else {\n int blocks = (int)((total + threads - 1) / threads);\n relu_bias_scalar_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), bias.data_ptr<float>(), total, C, HW);\n }\n return y;\n}\n'''\n\n_ext = load_inline(name='conv_relu_bias_fused_ext_v2',\n cpp_sources=[_cpp_src],\n cuda_sources=[_cuda_src],\n functions=['fused_forward'],\n extra_cuda_cflags=['-O3'],\n extra_cflags=['-O3'],\n verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, bias_shape):\n super(ModelNew, self).__init__()\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)\n self.bias = nn.Parameter(torch.randn(bias_shape))\n\n def forward(self, x):\n return _ext.fused_forward(x, self.conv.weight, self.conv.bias, self.bias)\n"
}
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (5 / 5) trials all matched the reference.
Numerical precision: max_abs_err=0.00e+00 mean_abs_err=0.00e+00 max_rel_err=0.00e+00
turn 49.8s turn · 9.8s LLM · 0.0s tools
call → static_check
callstatic_check
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r'''\n#include <torch/extension.h>\n#include <vector>\n\ntorch::Tensor apply_relu_bias_cuda(torch::Tensor y, torch::Tensor bias);\n\ntorch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add_bias) {\n std::optional<at::Tensor> b = conv_bias;\n std::vector<int64_t> stride = {1, 1};\n std::vector<int64_t> padding = {0, 0};\n std::vector<int64_t> dilation = {1, 1};\n auto y = at::conv2d(input, weight, b, stride, padding, dilation, 1);\n return apply_relu_bias_cuda(y, add_bias);\n}\n'''\n\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void relu_bias_scalar_kernel(float* __restrict__ y, const float* __restrict__ bias,\n long total, int C, int HW) {\n long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < total) {\n float v = y[idx];\n int c = (idx / HW) % C;\n y[idx] = (v > 0.0f ? v : 0.0f) + bias[c];\n }\n}\n\n__global__ void relu_bias_vec4_kernel(float4* __restrict__ y4, const float* __restrict__ bias,\n long total4, int C, int HW4) {\n long i = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (i < total4) {\n float4 v = y4[i];\n int c = (i / HW4) % C;\n float b = bias[c];\n v.x = (v.x > 0.0f ? v.x : 0.0f) + b;\n v.y = (v.y > 0.0f ? v.y : 0.0f) + b;\n v.z = (v.z > 0.0f ? v.z : 0.0f) + b;\n v.w = (v.w > 0.0f ? v.w : 0.0f) + b;\n y4[i] = v;\n }\n}\n\ntorch::Tensor apply_relu_bias_cuda(torch::Tensor y, torch::Tensor bias) {\n const long total = y.numel();\n if (total == 0) return y;\n const int C = (int)y.size(1);\n const int HW = (int)(y.size(2) * y.size(3));\n const int threads = 256;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if ((total % 4 == 0) && (HW % 4 == 0)) {\n long total4 = total / 4;\n int HW4 = HW / 4;\n int blocks = (int)((total4 + threads - 1) / threads);\n relu_bias_vec4_kernel<<<blocks, threads, 0, stream>>>(reinterpret_cast<float4*>(y.data_ptr<float>()), bias.data_ptr<float>(), total4, C, HW4);\n } else {\n int blocks = (int)((total + threads - 1) / threads);\n relu_bias_scalar_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), bias.data_ptr<float>(), total, C, HW);\n }\n return y;\n}\n'''\n\n_ext = load_inline(name='conv_relu_bias_fused_ext_v2',\n cpp_sources=[_cpp_src],\n cuda_sources=[_cuda_src],\n functions=['fused_forward'],\n extra_cuda_cflags=['-O3'],\n extra_cflags=['-O3'],\n verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, bias_shape):\n super(ModelNew, self).__init__()\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)\n self.bias = nn.Parameter(torch.randn(bias_shape))\n\n def forward(self, x):\n return _ext.fused_forward(x, self.conv.weight, self.conv.bias, self.bias)\n"
}
result ← static_check (ok)
okstatic_check
static_check PASSED: no violations or warnings detected.
turn 520.1s turn · 20.1s LLM · 0.0s tools
reasoning
**Investigating convolution biases**
I'm considering whether the conv2d function has fused the bias. It seems like cuDNN handles convolution without bias in the cudnnConvolutionForward call. Then, it likely uses cudnnAddTensor separately to apply the bias. I want to dig deeper into this to see how they interact and if optimization is happening. There’s definitely more to explore here, especially in how these functions work together!
call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r'''\n#include <torch/extension.h>\n#include <vector>\n\ntorch::Tensor apply_bias_relu_bias_cuda(torch::Tensor y, torch::Tensor conv_bias, torch::Tensor add_bias);\n\ntorch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add_bias) {\n std::optional<at::Tensor> b = std::nullopt;\n std::vector<int64_t> stride = {1, 1};\n std::vector<int64_t> padding = {0, 0};\n std::vector<int64_t> dilation = {1, 1};\n auto y = at::conv2d(input, weight, b, stride, padding, dilation, 1);\n return apply_bias_relu_bias_cuda(y, conv_bias, add_bias);\n}\n'''\n\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void brb_scalar_kernel(float* __restrict__ y, const float* __restrict__ cb, const float* __restrict__ ab,\n long total, int C, int HW) {\n long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < total) {\n int c = (idx / HW) % C;\n float v = y[idx] + cb[c];\n y[idx] = (v > 0.0f ? v : 0.0f) + ab[c];\n }\n}\n\n__global__ void brb_vec4_kernel(float4* __restrict__ y4, const float* __restrict__ cb, const float* __restrict__ ab,\n long total4, int C, int HW4) {\n long i = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (i < total4) {\n float4 v = y4[i];\n int c = (i / HW4) % C;\n float b0 = cb[c];\n float b1 = ab[c];\n v.x += b0; v.y += b0; v.z += b0; v.w += b0;\n v.x = (v.x > 0.0f ? v.x : 0.0f) + b1;\n v.y = (v.y > 0.0f ? v.y : 0.0f) + b1;\n v.z = (v.z > 0.0f ? v.z : 0.0f) + b1;\n v.w = (v.w > 0.0f ? v.w : 0.0f) + b1;\n y4[i] = v;\n }\n}\n\ntorch::Tensor apply_bias_relu_bias_cuda(torch::Tensor y, torch::Tensor conv_bias, torch::Tensor add_bias) {\n const long total = y.numel();\n if (total == 0) return y;\n const int C = (int)y.size(1);\n const int HW = (int)(y.size(2) * y.size(3));\n const int threads = 256;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if ((total % 4 == 0) && (HW % 4 == 0)) {\n long total4 = total / 4;\n int HW4 = HW / 4;\n int blocks = (int)((total4 + threads - 1) / threads);\n brb_vec4_kernel<<<blocks, threads, 0, stream>>>(reinterpret_cast<float4*>(y.data_ptr<float>()), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), total4, C, HW4);\n } else {\n int blocks = (int)((total + threads - 1) / threads);\n brb_scalar_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), total, C, HW);\n }\n return y;\n}\n'''\n\n_ext = load_inline(name='conv_relu_bias_fused_ext_v3_nobias',\n cpp_sources=[_cpp_src],\n cuda_sources=[_cuda_src],\n functions=['fused_forward'],\n extra_cuda_cflags=['-O3'],\n extra_cflags=['-O3'],\n verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, bias_shape):\n super(ModelNew, self).__init__()\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)\n self.bias = nn.Parameter(torch.randn(bias_shape))\n\n def forward(self, x):\n return _ext.fused_forward(x, self.conv.weight, self.conv.bias, self.bias)\n"
}
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (5 / 5) trials all matched the reference.
Numerical precision: max_abs_err=0.00e+00 mean_abs_err=0.00e+00 max_rel_err=0.00e+00
turn 68.5s turn · 8.5s LLM · 0.0s tools
call → static_check
callstatic_check
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r'''\n#include <torch/extension.h>\n#include <vector>\n\ntorch::Tensor apply_bias_relu_bias_cuda(torch::Tensor y, torch::Tensor conv_bias, torch::Tensor add_bias);\n\ntorch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add_bias) {\n std::optional<at::Tensor> b = std::nullopt;\n std::vector<int64_t> stride = {1, 1};\n std::vector<int64_t> padding = {0, 0};\n std::vector<int64_t> dilation = {1, 1};\n auto y = at::conv2d(input, weight, b, stride, padding, dilation, 1);\n return apply_bias_relu_bias_cuda(y, conv_bias, add_bias);\n}\n'''\n\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void brb_scalar_kernel(float* __restrict__ y, const float* __restrict__ cb, const float* __restrict__ ab,\n long total, int C, int HW) {\n long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < total) {\n int c = (idx / HW) % C;\n float v = y[idx] + cb[c];\n y[idx] = (v > 0.0f ? v : 0.0f) + ab[c];\n }\n}\n\n__global__ void brb_vec4_kernel(float4* __restrict__ y4, const float* __restrict__ cb, const float* __restrict__ ab,\n long total4, int C, int HW4) {\n long i = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (i < total4) {\n float4 v = y4[i];\n int c = (i / HW4) % C;\n float b0 = cb[c];\n float b1 = ab[c];\n v.x += b0; v.y += b0; v.z += b0; v.w += b0;\n v.x = (v.x > 0.0f ? v.x : 0.0f) + b1;\n v.y = (v.y > 0.0f ? v.y : 0.0f) + b1;\n v.z = (v.z > 0.0f ? v.z : 0.0f) + b1;\n v.w = (v.w > 0.0f ? v.w : 0.0f) + b1;\n y4[i] = v;\n }\n}\n\ntorch::Tensor apply_bias_relu_bias_cuda(torch::Tensor y, torch::Tensor conv_bias, torch::Tensor add_bias) {\n const long total = y.numel();\n if (total == 0) return y;\n const int C = (int)y.size(1);\n const int HW = (int)(y.size(2) * y.size(3));\n const int threads = 256;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if ((total % 4 == 0) && (HW % 4 == 0)) {\n long total4 = total / 4;\n int HW4 = HW / 4;\n int blocks = (int)((total4 + threads - 1) / threads);\n brb_vec4_kernel<<<blocks, threads, 0, stream>>>(reinterpret_cast<float4*>(y.data_ptr<float>()), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), total4, C, HW4);\n } else {\n int blocks = (int)((total + threads - 1) / threads);\n brb_scalar_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), total, C, HW);\n }\n return y;\n}\n'''\n\n_ext = load_inline(name='conv_relu_bias_fused_ext_v3_nobias',\n cpp_sources=[_cpp_src],\n cuda_sources=[_cuda_src],\n functions=['fused_forward'],\n extra_cuda_cflags=['-O3'],\n extra_cflags=['-O3'],\n verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, bias_shape):\n super(ModelNew, self).__init__()\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)\n self.bias = nn.Parameter(torch.randn(bias_shape))\n\n def forward(self, x):\n return _ext.fused_forward(x, self.conv.weight, self.conv.bias, self.bias)\n"
}
result ← static_check (ok)
okstatic_check
static_check PASSED: no violations or warnings detected.
turn 726.5s turn · 26.5s LLM · 0.0s tools
call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r'''\n#include <torch/extension.h>\n#include <optional>\n\ntorch::Tensor apply_bias_relu_bias_cuda(torch::Tensor y, torch::Tensor conv_bias, torch::Tensor add_bias);\n\ntorch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add_bias) {\n static const int64_t stride_data[2] = {1, 1};\n static const int64_t padding_data[2] = {0, 0};\n static const int64_t dilation_data[2] = {1, 1};\n at::IntArrayRef stride(stride_data, 2), padding(padding_data, 2), dilation(dilation_data, 2);\n std::optional<at::Tensor> b = std::nullopt;\n auto y = at::conv2d(input, weight, b, stride, padding, dilation, 1);\n return apply_bias_relu_bias_cuda(y, conv_bias, add_bias);\n}\n'''\n\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void brb_vec4_2d_c128_kernel(float4* __restrict__ y4, const float* __restrict__ cb, const float* __restrict__ ab,\n int HW4) {\n int x = blockIdx.x * blockDim.x + threadIdx.x;\n int nc = blockIdx.y;\n if (x < HW4) {\n int c = nc & 127;\n long i = (long)nc * HW4 + x;\n float4 v = y4[i];\n float b0 = cb[c];\n float b1 = ab[c];\n v.x += b0; v.y += b0; v.z += b0; v.w += b0;\n v.x = (v.x > 0.0f ? v.x : 0.0f) + b1;\n v.y = (v.y > 0.0f ? v.y : 0.0f) + b1;\n v.z = (v.z > 0.0f ? v.z : 0.0f) + b1;\n v.w = (v.w > 0.0f ? v.w : 0.0f) + b1;\n y4[i] = v;\n }\n}\n\n__global__ void brb_vec4_kernel(float4* __restrict__ y4, const float* __restrict__ cb, const float* __restrict__ ab,\n long total4, int C, int HW4) {\n long i = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (i < total4) {\n float4 v = y4[i];\n int c = (i / HW4) % C;\n float b0 = cb[c];\n float b1 = ab[c];\n v.x += b0; v.y += b0; v.z += b0; v.w += b0;\n v.x = (v.x > 0.0f ? v.x : 0.0f) + b1;\n v.y = (v.y > 0.0f ? v.y : 0.0f) + b1;\n v.z = (v.z > 0.0f ? v.z : 0.0f) + b1;\n v.w = (v.w > 0.0f ? v.w : 0.0f) + b1;\n y4[i] = v;\n }\n}\n\n__global__ void brb_scalar_kernel(float* __restrict__ y, const float* __restrict__ cb, const float* __restrict__ ab,\n long total, int C, int HW) {\n long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < total) {\n int c = (idx / HW) % C;\n float v = y[idx] + cb[c];\n y[idx] = (v > 0.0f ? v : 0.0f) + ab[c];\n }\n}\n\ntorch::Tensor apply_bias_relu_bias_cuda(torch::Tensor y, torch::Tensor conv_bias, torch::Tensor add_bias) {\n const long total = y.numel();\n if (total == 0) return y;\n const int N = (int)y.size(0);\n const int C = (int)y.size(1);\n const int HW = (int)(y.size(2) * y.size(3));\n const int threads = 256;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if ((total % 4 == 0) && (HW % 4 == 0)) {\n int HW4 = HW / 4;\n if (C == 128 && N * C <= 65535) {\n dim3 blocks((HW4 + threads - 1) / threads, N * C);\n brb_vec4_2d_c128_kernel<<<blocks, threads, 0, stream>>>(reinterpret_cast<float4*>(y.data_ptr<float>()), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), HW4);\n } else {\n long total4 = total / 4;\n int blocks = (int)((total4 + threads - 1) / threads);\n brb_vec4_kernel<<<blocks, threads, 0, stream>>>(reinterpret_cast<float4*>(y.data_ptr<float>()), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), total4, C, HW4);\n }\n } else {\n int blocks = (int)((total + threads - 1) / threads);\n brb_scalar_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), total, C, HW);\n }\n return y;\n}\n'''\n\n_ext = load_inline(name='conv_relu_bias_fused_ext_v4_2d',\n cpp_sources=[_cpp_src],\n cuda_sources=[_cuda_src],\n functions=['fused_forward'],\n extra_cuda_cflags=['-O3'],\n extra_cflags=['-O3'],\n verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, bias_shape):\n super(ModelNew, self).__init__()\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)\n self.bias = nn.Parameter(torch.randn(bias_shape))\n\n def forward(self, x):\n return _ext.fused_forward(x, self.conv.weight, self.conv.bias, self.bias)\n"
}
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (5 / 5) trials all matched the reference.
Numerical precision: max_abs_err=0.00e+00 mean_abs_err=0.00e+00 max_rel_err=0.00e+00
turn 815.8s turn · 15.8s LLM · 0.0s tools · FINAL
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n_cpp_src = r'''\n#include <torch/extension.h>\n#include <optional>\n\ntorch::Tensor apply_bias_relu_bias_cuda(torch::Tensor y, torch::Tensor conv_bias, torch::Tensor add_bias);\n\ntorch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add_bias) {\n static const int64_t stride_data[2] = {1, 1};\n static const int64_t padding_data[2] = {0, 0};\n static const int64_t dilation_data[2] = {1, 1};\n at::IntArrayRef stride(stride_data, 2), padding(padding_data, 2), dilation(dilation_data, 2);\n std::optional<at::Tensor> b = std::nullopt;\n auto y = at::conv2d(input, weight, b, stride, padding, dilation, 1);\n return apply_bias_relu_bias_cuda(y, conv_bias, add_bias);\n}\n'''\n\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void brb_vec4_2d_c128_kernel(float4* __restrict__ y4, const float* __restrict__ cb, const float* __restrict__ ab,\n int HW4) {\n int x = blockIdx.x * blockDim.x + threadIdx.x;\n int nc = blockIdx.y;\n if (x < HW4) {\n int c = nc & 127;\n long i = (long)nc * HW4 + x;\n float4 v = y4[i];\n float b0 = cb[c];\n float b1 = ab[c];\n v.x += b0; v.y += b0; v.z += b0; v.w += b0;\n v.x = (v.x > 0.0f ? v.x : 0.0f) + b1;\n v.y = (v.y > 0.0f ? v.y : 0.0f) + b1;\n v.z = (v.z > 0.0f ? v.z : 0.0f) + b1;\n v.w = (v.w > 0.0f ? v.w : 0.0f) + b1;\n y4[i] = v;\n }\n}\n\n__global__ void brb_vec4_kernel(float4* __restrict__ y4, const float* __restrict__ cb, const float* __restrict__ ab,\n long total4, int C, int HW4) {\n long i = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (i < total4) {\n float4 v = y4[i];\n int c = (i / HW4) % C;\n float b0 = cb[c];\n float b1 = ab[c];\n v.x += b0; v.y += b0; v.z += b0; v.w += b0;\n v.x = (v.x > 0.0f ? v.x : 0.0f) + b1;\n v.y = (v.y > 0.0f ? v.y : 0.0f) + b1;\n v.z = (v.z > 0.0f ? v.z : 0.0f) + b1;\n v.w = (v.w > 0.0f ? v.w : 0.0f) + b1;\n y4[i] = v;\n }\n}\n\n__global__ void brb_scalar_kernel(float* __restrict__ y, const float* __restrict__ cb, const float* __restrict__ ab,\n long total, int C, int HW) {\n long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < total) {\n int c = (idx / HW) % C;\n float v = y[idx] + cb[c];\n y[idx] = (v > 0.0f ? v : 0.0f) + ab[c];\n }\n}\n\ntorch::Tensor apply_bias_relu_bias_cuda(torch::Tensor y, torch::Tensor conv_bias, torch::Tensor add_bias) {\n const long total = y.numel();\n if (total == 0) return y;\n const int N = (int)y.size(0);\n const int C = (int)y.size(1);\n const int HW = (int)(y.size(2) * y.size(3));\n const int threads = 256;\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n if ((total % 4 == 0) && (HW % 4 == 0)) {\n int HW4 = HW / 4;\n if (C == 128 && N * C <= 65535) {\n dim3 blocks((HW4 + threads - 1) / threads, N * C);\n brb_vec4_2d_c128_kernel<<<blocks, threads, 0, stream>>>(reinterpret_cast<float4*>(y.data_ptr<float>()), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), HW4);\n } else {\n long total4 = total / 4;\n int blocks = (int)((total4 + threads - 1) / threads);\n brb_vec4_kernel<<<blocks, threads, 0, stream>>>(reinterpret_cast<float4*>(y.data_ptr<float>()), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), total4, C, HW4);\n }\n } else {\n int blocks = (int)((total + threads - 1) / threads);\n brb_scalar_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), total, C, HW);\n }\n return y;\n}\n'''\n\n_ext = load_inline(name='conv_relu_bias_fused_ext_final',\n cpp_sources=[_cpp_src],\n cuda_sources=[_cuda_src],\n functions=['fused_forward'],\n extra_cuda_cflags=['-O3'],\n extra_cflags=['-O3'],\n verbose=False)\n\nclass ModelNew(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, bias_shape):\n super(ModelNew, self).__init__()\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)\n self.bias = nn.Parameter(torch.randn(bias_shape))\n\n def forward(self, x):\n return _ext.fused_forward(x, self.conv.weight, self.conv.bias, self.bias)\n"
}
result ← submit_kernel (ok)
oksubmit_kernel
submit_kernel PASSED: (5 / 5) correctness trials all passed.
Kernel runtime: 2.73 μs
Runtime stats: mean=2.73μs median=?μs std=0.88μs
Numerical precision: max_abs_err=0.00e+00 mean_abs_err=0.00e+00 max_rel_err=0.00e+00
Memory: 2050.1 MB peak (ref: 2530.06 MB, ratio: 0.8103x)
Kernel launches: 4 (ref: 10, fusion ratio: 2.5)
Energy: 377.28 mJ/run (ref: 754.76 mJ, ratio: 2.0005x)
SOL score: 1.0
final submitted kernel (116 lines, 4,638 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
_cpp_src = r'''
#include <torch/extension.h>
#include <optional>
torch::Tensor apply_bias_relu_bias_cuda(torch::Tensor y, torch::Tensor conv_bias, torch::Tensor add_bias);
torch::Tensor fused_forward(torch::Tensor input, torch::Tensor weight, torch::Tensor conv_bias, torch::Tensor add_bias) {
static const int64_t stride_data[2] = {1, 1};
static const int64_t padding_data[2] = {0, 0};
static const int64_t dilation_data[2] = {1, 1};
at::IntArrayRef stride(stride_data, 2), padding(padding_data, 2), dilation(dilation_data, 2);
std::optional<at::Tensor> b = std::nullopt;
auto y = at::conv2d(input, weight, b, stride, padding, dilation, 1);
return apply_bias_relu_bias_cuda(y, conv_bias, add_bias);
}
'''
_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void brb_vec4_2d_c128_kernel(float4* __restrict__ y4, const float* __restrict__ cb, const float* __restrict__ ab,
int HW4) {
int x = blockIdx.x * blockDim.x + threadIdx.x;
int nc = blockIdx.y;
if (x < HW4) {
int c = nc & 127;
long i = (long)nc * HW4 + x;
float4 v = y4[i];
float b0 = cb[c];
float b1 = ab[c];
v.x += b0; v.y += b0; v.z += b0; v.w += b0;
v.x = (v.x > 0.0f ? v.x : 0.0f) + b1;
v.y = (v.y > 0.0f ? v.y : 0.0f) + b1;
v.z = (v.z > 0.0f ? v.z : 0.0f) + b1;
v.w = (v.w > 0.0f ? v.w : 0.0f) + b1;
y4[i] = v;
}
}
__global__ void brb_vec4_kernel(float4* __restrict__ y4, const float* __restrict__ cb, const float* __restrict__ ab,
long total4, int C, int HW4) {
long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
if (i < total4) {
float4 v = y4[i];
int c = (i / HW4) % C;
float b0 = cb[c];
float b1 = ab[c];
v.x += b0; v.y += b0; v.z += b0; v.w += b0;
v.x = (v.x > 0.0f ? v.x : 0.0f) + b1;
v.y = (v.y > 0.0f ? v.y : 0.0f) + b1;
v.z = (v.z > 0.0f ? v.z : 0.0f) + b1;
v.w = (v.w > 0.0f ? v.w : 0.0f) + b1;
y4[i] = v;
}
}
__global__ void brb_scalar_kernel(float* __restrict__ y, const float* __restrict__ cb, const float* __restrict__ ab,
long total, int C, int HW) {
long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
if (idx < total) {
int c = (idx / HW) % C;
float v = y[idx] + cb[c];
y[idx] = (v > 0.0f ? v : 0.0f) + ab[c];
}
}
torch::Tensor apply_bias_relu_bias_cuda(torch::Tensor y, torch::Tensor conv_bias, torch::Tensor add_bias) {
const long total = y.numel();
if (total == 0) return y;
const int N = (int)y.size(0);
const int C = (int)y.size(1);
const int HW = (int)(y.size(2) * y.size(3));
const int threads = 256;
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
if ((total % 4 == 0) && (HW % 4 == 0)) {
int HW4 = HW / 4;
if (C == 128 && N * C <= 65535) {
dim3 blocks((HW4 + threads - 1) / threads, N * C);
brb_vec4_2d_c128_kernel<<<blocks, threads, 0, stream>>>(reinterpret_cast<float4*>(y.data_ptr<float>()), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), HW4);
} else {
long total4 = total / 4;
int blocks = (int)((total4 + threads - 1) / threads);
brb_vec4_kernel<<<blocks, threads, 0, stream>>>(reinterpret_cast<float4*>(y.data_ptr<float>()), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), total4, C, HW4);
}
} else {
int blocks = (int)((total + threads - 1) / threads);
brb_scalar_kernel<<<blocks, threads, 0, stream>>>(y.data_ptr<float>(), conv_bias.data_ptr<float>(), add_bias.data_ptr<float>(), total, C, HW);
}
return y;
}
'''
_ext = load_inline(name='conv_relu_bias_fused_ext_final',
cpp_sources=[_cpp_src],
cuda_sources=[_cuda_src],
functions=['fused_forward'],
extra_cuda_cflags=['-O3'],
extra_cflags=['-O3'],
verbose=False)
class ModelNew(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, bias_shape):
super(ModelNew, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
self.bias = nn.Parameter(torch.randn(bias_shape))
def forward(self, x):
return _ext.fused_forward(x, self.conv.weight, self.conv.bias, self.bias)