2.57. 2D Convolution with ReLU and HardSwish Activation
Performs a 2D convolution on the input, applies the rectified linear unit (ReLU) activation to filter negative values, and then calculates the hard swish activation by multiplying the result element-wise with a clamped linear transformation of the data ((x + 3) / 6).
Progress
Step 3 of 3
Reward
3.95x
Optimization Timeline
Step 3 Details
Reasoning
We can further optimize our fused convolution kernel by specializing the common case when kernel_size==3 and in_channels==3 (e.g. RGB images). In this case, we unroll completely the convolution loops over the kernel and input channels to reduce loop overhead and improve instruction-level parallelism. This specialized branch will compute the convolution for the 3×3 kernel for exactly three input channels by writing out each multiplication and addition explicitly. For other cases, we fall back on our previous unrolled loop for kernel_size==3 or the general triple‐nested loop for arbitrary kernel sizes. We then apply bias addition, ReLU, and HardSwish activation as before. This improvement genuinely helps by reducing loop overhead and enabling more compiler optimizations (such as FMA scheduling) for the common case, which is not a workaround but a genuine micro-optimization. Below is the updated Python module (which remains unchanged) and the updated CUDA kernel code.
#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
// Fused kernel computes convolution, adds bias, applies ReLU and then HardSwish.
// HardSwish: hswish(x) = ReLU(x) * min((ReLU(x)+3)/6, 1)
// Each thread computes one output element.
__global__ void fused_conv_relu_hardswish_kernel(
const float* __restrict__ input, // [batch, in_channels, H, W]
const float* __restrict__ weight, // [out_channels, in_channels, k, k]
const float* __restrict__ bias, // [out_channels]
float* __restrict__ output, // [batch, out_channels, out_H, out_W]
int batch, int in_channels,
int in_H, int in_W,
int out_channels,
int kernel_size)
{
// Compute output dimensions.
int out_H = in_H - kernel_size + 1;
int out_W = in_W - kernel_size + 1;
int index = blockIdx.x * blockDim.x + threadIdx.x;
int total = batch * out_channels * out_H * out_W;
if (index >= total)
return;
// Decode index into (b, oc, oh, ow)
int ow = index % out_W;
int tmp = index / out_W;
int oh = tmp % out_H;
tmp = tmp / out_H;
int oc = tmp % out_channels;
int b = tmp / out_channels;
float sum = 0.0f;
// Check for the specialized case: kernel_size == 3 and in_channels == 3.
if (kernel_size == 3 && in_channels == 3) {
// Fully unroll for 3x3 convolution and 3 channels.
// Loop over each input channel explicitly.
// For each input channel, compute base indices.
int base_input0 = ((b * 3 + 0) * in_H + oh) * in_W + ow;
int base_input1 = ((b * 3 + 1) * in_H + oh) * in_W + ow;
int base_input2 = ((b * 3 + 2) * in_H + oh) * in_W + ow;
// Weight base indices for each channel: each channel has 9 elements.
int base_weight0 = ((oc * 3 + 0) * 9);
int base_weight1 = ((oc * 3 + 1) * 9);
int base_weight2 = ((oc * 3 + 2) * 9);
// For channel 0:
sum += input[base_input0 + 0] * weight[base_weight0 + 0];
sum += input[base_input0 + 1] * weight[base_weight0 + 1];
sum += input[base_input0 + 2] * weight[base_weight0 + 2];
sum += input[base_input0 + in_W] * weight[base_weight0 + 3];
sum += input[base_input0 + in_W + 1] * weight[base_weight0 + 4];
sum += input[base_input0 + in_W + 2] * weight[base_weight0 + 5];
sum += input[base_input0 + 2*in_W] * weight[base_weight0 + 6];
sum += input[base_input0 + 2*in_W + 1] * weight[base_weight0 + 7];
sum += input[base_input0 + 2*in_W + 2] * weight[base_weight0 + 8];
// For channel 1:
sum += input[base_input1 + 0] * weight[base_weight1 + 0];
sum += input[base_input1 + 1] * weight[base_weight1 + 1];
sum += input[base_input1 + 2] * weight[base_weight1 + 2];
sum += input[base_input1 + in_W] * weight[base_weight1 + 3];
sum += input[base_input1 + in_W + 1] * weight[base_weight1 + 4];
sum += input[base_input1 + in_W + 2] * weight[base_weight1 + 5];
sum += input[base_input1 + 2*in_W] * weight[base_weight1 + 6];
sum += input[base_input1 + 2*in_W + 1] * weight[base_weight1 + 7];
sum += input[base_input1 + 2*in_W + 2] * weight[base_weight1 + 8];
// For channel 2:
sum += input[base_input2 + 0] * weight[base_weight2 + 0];
sum += input[base_input2 + 1] * weight[base_weight2 + 1];
sum += input[base_input2 + 2] * weight[base_weight2 + 2];
sum += input[base_input2 + in_W] * weight[base_weight2 + 3];
sum += input[base_input2 + in_W + 1] * weight[base_weight2 + 4];
sum += input[base_input2 + in_W + 2] * weight[base_weight2 + 5];
sum += input[base_input2 + 2*in_W] * weight[base_weight2 + 6];
sum += input[base_input2 + 2*in_W + 1] * weight[base_weight2 + 7];
sum += input[base_input2 + 2*in_W + 2] * weight[base_weight2 + 8];
} else if (kernel_size == 3) {
// For kernel_size 3 but arbitrary in_channels.
for (int ic = 0; ic < in_channels; ic++) {
int input_base = ((b * in_channels + ic) * in_H + oh) * in_W + ow;
int weight_base = ((oc * in_channels + ic) * 9); // 3*3 = 9
sum += input[input_base + 0] * weight[weight_base + 0];
sum += input[input_base + 1] * weight[weight_base + 1];
sum += input[input_base + 2] * weight[weight_base + 2];
sum += input[input_base + in_W] * weight[weight_base + 3];
sum += input[input_base + in_W + 1] * weight[weight_base + 4];
sum += input[input_base + in_W + 2] * weight[weight_base + 5];
sum += input[input_base + 2*in_W] * weight[weight_base + 6];
sum += input[input_base + 2*in_W + 1] * weight[weight_base + 7];
sum += input[input_base + 2*in_W + 2] * weight[weight_base + 8];
}
} else {
// General case for any kernel size.
for (int ic = 0; ic < in_channels; ic++) {
for (int kh = 0; kh < kernel_size; kh++) {
for (int kw = 0; kw < kernel_size; kw++) {
int in_h = oh + kh;
int in_w = ow + kw;
int input_idx = ((b * in_channels + ic) * in_H + in_h) * in_W + in_w;
int weight_idx = ((oc * in_channels + ic) * kernel_size + kh) * kernel_size + kw;
sum += input[input_idx] * weight[weight_idx];
}
}
}
}
// Add bias.
sum += bias[oc];
// Apply ReLU.
float relu_val = (sum > 0.0f ? sum : 0.0f);
// Compute HardSwish: hswish(x) = x * min((x+3)/6, 1)
float multiplier = (relu_val + 3.0f) / 6.0f;
if (multiplier > 1.0f) multiplier = 1.0f;
float result = relu_val * multiplier;
output[index] = result;
}
torch::Tensor forward(torch::Tensor input, torch::Tensor weight, torch::Tensor bias, int kernel_size) {
c10::cuda::CUDAGuard device_guard(input.device());
int batch = input.size(0);
int in_channels = input.size(1);
int in_H = input.size(2);
int in_W = input.size(3);
int out_channels = weight.size(0);
int out_H = in_H - kernel_size + 1;
int out_W = in_W - kernel_size + 1;
auto output = torch::empty({batch, out_channels, out_H, out_W}, input.options());
int total = batch * out_channels * out_H * out_W;
int threads = 256;
int blocks = (total + threads - 1) / threads;
fused_conv_relu_hardswish_kernel<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
input.data_ptr<float>(),
weight.data_ptr<float>(),
bias.data_ptr<float>(),
output.data_ptr<float>(),
batch, in_channels,
in_H, in_W,
out_channels,
kernel_size
);
return output;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("forward", &forward, "Fused Conv2d, ReLU, and HardSwish activation (CUDA)");
}