reference PyTorch source (the prompt the agent sees)
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Model that performs a convolution, subtraction, tanh activation, subtraction and average pooling.
"""
def __init__(self, in_channels, out_channels, kernel_size, subtract1_value, subtract2_value, kernel_size_pool):
super(Model, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
self.subtract1_value = subtract1_value
self.subtract2_value = subtract2_value
self.avgpool = nn.AvgPool2d(kernel_size_pool)
def forward(self, x):
x = self.conv(x)
x = x - self.subtract1_value
x = torch.tanh(x)
x = x - self.subtract2_value
x = self.avgpool(x)
return x
batch_size = 128
in_channels = 64
out_channels = 128
height, width = 128, 128
kernel_size = 3
subtract1_value = 0.5
subtract2_value = 0.2
kernel_size_pool = 2
def get_inputs():
return [torch.rand(batch_size, in_channels, height, width)]
def get_init_inputs():
return [in_channels, out_channels, kernel_size, subtract1_value, subtract2_value, kernel_size_pool]final submitted kernel (126 lines, 5,179 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cpp_source = r'''
#include <torch/extension.h>
#include <ATen/ops/conv2d.h>
void launch_post(torch::Tensor conv, torch::Tensor out, double subtract1, double subtract2, int64_t pool_k);
torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
double subtract1, double subtract2, int64_t pool_k) {
auto conv = at::conv2d(x, weight, c10::optional<at::Tensor>(bias),
at::IntArrayRef({1,1}), at::IntArrayRef({0,0}),
at::IntArrayRef({1,1}), 1);
auto out = torch::empty({conv.size(0), conv.size(1), conv.size(2) / pool_k, conv.size(3) / pool_k}, conv.options());
launch_post(conv, out, subtract1, subtract2, pool_k);
return out;
}
'''
cuda_source = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#define C0 128
#define HC 126
#define WC 126
#define HO 63
#define WO 63
__global__ void post_pool2_fixed_kernel(const float* __restrict__ y, float* __restrict__ out,
int total, float s1, float s2) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total) return;
int ow = idx % WO;
int tmp = idx / WO;
int oh = tmp % HO;
tmp /= HO;
int c = tmp & 127;
int n = tmp >> 7;
long long base = (((long long)n * C0 + c) * HC + (oh << 1)) * WC + (ow << 1);
float a = __tanhf(y[base] - s1);
float b = __tanhf(y[base + 1] - s1);
float c0 = __tanhf(y[base + WC] - s1);
float d = __tanhf(y[base + WC + 1] - s1);
out[idx] = 0.25f * (a + b + c0 + d) - s2;
}
__global__ void post_pool2_flat_kernel(const float* __restrict__ y, float* __restrict__ out,
int total, int C, int Hc, int Wc, float s1, float s2) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total) return;
int Wo = Wc >> 1;
int Ho = Hc >> 1;
int ow = idx % Wo;
int tmp = idx / Wo;
int oh = tmp % Ho;
tmp /= Ho;
int c = tmp % C;
int n = tmp / C;
long long base = (((long long)n * C + c) * Hc + (oh << 1)) * Wc + (ow << 1);
float a = __tanhf(y[base] - s1);
float b = __tanhf(y[base + 1] - s1);
float c0 = __tanhf(y[base + Wc] - s1);
float d = __tanhf(y[base + Wc + 1] - s1);
out[idx] = 0.25f * (a + b + c0 + d) - s2;
}
__global__ void post_pool_generic_kernel(const float* __restrict__ y, float* __restrict__ out,
int total, int C, int Hc, int Wc, int pk, float s1, float s2) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total) return;
int Wo = Wc / pk;
int Ho = Hc / pk;
int ow = idx % Wo;
int tmp = idx / Wo;
int oh = tmp % Ho;
tmp /= Ho;
int c = tmp % C;
int n = tmp / C;
float sum = 0.0f;
long long base0 = (((long long)n * C + c) * Hc + oh * pk) * Wc + ow * pk;
for (int i = 0; i < pk; ++i) {
long long base = base0 + (long long)i * Wc;
for (int j = 0; j < pk; ++j) sum += __tanhf(y[base + j] - s1);
}
out[idx] = sum / (float)(pk * pk) - s2;
}
void launch_post(torch::Tensor conv, torch::Tensor out, double subtract1, double subtract2, int64_t pool_k) {
int64_t N = conv.size(0), C = conv.size(1), Hc = conv.size(2), Wc = conv.size(3);
int64_t total64 = N * C * (Hc / pool_k) * (Wc / pool_k);
int threads = 256;
int blocks = (int)((total64 + threads - 1) / threads);
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
if (pool_k == 2 && C == 128 && Hc == 126 && Wc == 126) {
post_pool2_fixed_kernel<<<blocks, threads, 0, stream>>>(conv.data_ptr<float>(), out.data_ptr<float>(), (int)total64, (float)subtract1, (float)subtract2);
} else if (pool_k == 2) {
post_pool2_flat_kernel<<<blocks, threads, 0, stream>>>(conv.data_ptr<float>(), out.data_ptr<float>(), (int)total64, (int)C, (int)Hc, (int)Wc, (float)subtract1, (float)subtract2);
} else {
post_pool_generic_kernel<<<blocks, threads, 0, stream>>>(conv.data_ptr<float>(), out.data_ptr<float>(), (int)total64, (int)C, (int)Hc, (int)Wc, (int)pool_k, (float)subtract1, (float)subtract2);
}
}
'''
_ext = load_inline(name='fused_conv_tanh_pool_final',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['fused_forward'],
extra_cflags=['-O3'],
extra_cuda_cflags=['-O3'],
verbose=False)
class ModelNew(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, subtract1_value, subtract2_value, kernel_size_pool):
super(ModelNew, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
self.subtract1_value = subtract1_value
self.subtract2_value = subtract2_value
self.kernel_size_pool = kernel_size_pool
def forward(self, x):
return _ext.fused_forward(x, self.conv.weight, self.conv.bias,
float(self.subtract1_value), float(self.subtract2_value), int(self.kernel_size_pool))