The AI CUDA Engineer 👷

import torch
import torch.nn as nn
import torch.nn.functional as F


def module_fn(
    x: torch.Tensor,
    conv1_weight: nn.Parameter,
    conv1_bias: nn.Parameter,
    conv2_weight: nn.Parameter,
    conv2_bias: nn.Parameter,
    fc1_weight: nn.Parameter,
    fc1_bias: nn.Parameter,
    fc2_weight: nn.Parameter,
    fc2_bias: nn.Parameter,
    fc3_weight: nn.Parameter,
    fc3_bias: nn.Parameter,
) -> torch.Tensor:
    """
    Implements a LeNet-5 architecture with ReLU activation.

    Args:
        x (torch.Tensor): The input tensor, shape (batch_size, 1, 32, 32)
        conv1_weight (nn.Parameter): Parameters for first conv layer
        conv1_bias (nn.Parameter): Parameters for first conv layer
        conv2_weight (nn.Parameter): Parameters for second conv layer
        conv2_bias (nn.Parameter): Parameters for second conv layer
        fc1_weight (nn.Parameter): Parameters for first FC layer
        fc1_bias (nn.Parameter): Parameters for first FC layer
        fc2_weight (nn.Parameter): Parameters for second FC layer
        fc3_weight (nn.Parameter): Parameters for third FC layer
        fc3_bias (nn.Parameter): Parameters for third FC layer

    Returns:
        torch.Tensor: The output tensor, shape (batch_size, num_classes)
    """
    # First convolutional layer with ReLU activation and max pooling
    x = F.conv2d(x, conv1_weight, conv1_bias, stride=1)
    x = F.relu(x)
    x = F.max_pool2d(x, kernel_size=2, stride=2)

    # Second convolutional layer with ReLU activation and max pooling
    x = F.conv2d(x, conv2_weight, conv2_bias, stride=1)
    x = F.relu(x)
    x = F.max_pool2d(x, kernel_size=2, stride=2)

    # Flatten the output for the fully connected layers
    x = x.view(-1, 16 * 5 * 5)

    # First fully connected layer with ReLU activation
    x = F.linear(x, fc1_weight, fc1_bias)
    x = F.relu(x)

    # Second fully connected layer with ReLU activation
    x = F.linear(x, fc2_weight, fc2_bias)
    x = F.relu(x)

    # Final fully connected layer
    x = F.linear(x, fc3_weight, fc3_bias)

    return x


class Model(nn.Module):
    def __init__(self, num_classes):
        """
        LeNet-5 architecture implementation in PyTorch.

        :param num_classes: The number of output classes.
        """
        super(Model, self).__init__()

        # Extract parameters from convolutional layers
        conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1)
        self.conv1_weight = nn.Parameter(conv1.weight.data.clone())
        self.conv1_bias = nn.Parameter(conv1.bias.data.clone())

        conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1)
        self.conv2_weight = nn.Parameter(conv2.weight.data.clone())
        self.conv2_bias = nn.Parameter(conv2.bias.data.clone())

        # Extract parameters from fully connected layers
        fc1 = nn.Linear(in_features=16 * 5 * 5, out_features=120)
        self.fc1_weight = nn.Parameter(fc1.weight.data.clone())
        self.fc1_bias = nn.Parameter(fc1.bias.data.clone())

        fc2 = nn.Linear(in_features=120, out_features=84)
        self.fc2_weight = nn.Parameter(fc2.weight.data.clone())
        self.fc2_bias = nn.Parameter(fc2.bias.data.clone())

        fc3 = nn.Linear(in_features=84, out_features=num_classes)
        self.fc3_weight = nn.Parameter(fc3.weight.data.clone())
        self.fc3_bias = nn.Parameter(fc3.bias.data.clone())

    def forward(self, x, fn=module_fn):
        return fn(
            x,
            self.conv1_weight,
            self.conv1_bias,
            self.conv2_weight,
            self.conv2_bias,
            self.fc1_weight,
            self.fc1_bias,
            self.fc2_weight,
            self.fc2_bias,
            self.fc3_weight,
            self.fc3_bias,
        )


# Test code for the LeNet-5 model
batch_size = 1
num_classes = 10


def get_inputs():
    return [torch.randn(batch_size, 1, 32, 32)]


def get_init_inputs():
    return [num_classes]

#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cublas_v2.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/CUDAUtils.h>

// ReLU kernel using fixed iteration loops with predicated operations
// The boundary check is implemented via a ternary operator to help the compiler generate predicated instructions,
// minimizing warp divergence for threads in the last (partial) iteration.
__global__ void relu_kernel(float* input, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;
    int iterations = (size + stride - 1) / stride;
    for (int i = 0; i < iterations; i++) {
        int pos = idx + i * stride;
        // Use a ternary operator to load a valid value; out‐of‐bounds threads load 0.0f
        float val = (pos < size) ? input[pos] : 0.0f;
        float res = fmaxf(0.0f, val);
        if (pos < size) {
            input[pos] = res;
        }
    }
}

// Optimized 2x2 max pooling kernel with uniform control flow
// Assumes pooling window is 2x2 with stride 2. The inner max calculation uses branchless fmaxf calls.
__global__ void max_pool2d_kernel(const float* __restrict__ input, float* __restrict__ output,
    int batch_size, int channels, int height, int width, int stride) {
    // Compute output dimensions (assumes height and width are divisible by 2)
    int out_h = height / 2;
    int out_w = width / 2;
    int total = batch_size * channels * out_h * out_w;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride_total = gridDim.x * blockDim.x;
    for (int i = idx; i < total; i += stride_total) {
         int temp = i;
         int b = temp / (channels * out_h * out_w);
         temp = temp % (channels * out_h * out_w);
         int c = temp / (out_h * out_w);
         temp = temp % (out_h * out_w);
         int ph = temp / out_w;
         int pw = temp % out_w;
         int h_offset = ph * stride;
         int w_offset = pw * stride;
         int base = ((b * channels + c) * height);
         const float* row0 = input + (base + h_offset) * width + w_offset;
         const float* row1 = row0 + width;
         float m0 = fmaxf(row0[0], row0[1]);
         float m1 = fmaxf(row1[0], row1[1]);
         output[i] = fmaxf(m0, m1);
    }
}

// Flatten kernel to copy data with uniform loop execution
__global__ void flatten_kernel(const float* input, float* output, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = gridDim.x * blockDim.x;
    int iterations = (size + stride - 1) / stride;
    for (int i = 0; i < iterations; i++) {
         int pos = idx + i * stride;
         if (pos < size) {
             output[pos] = input[pos];
         }
    }
}

// Forward function implementing a LeNet-5 architecture using custom kernels
// The conv and fc layers use PyTorch operators to ensure numerical correctness, while custom CUDA kernels
// (relu_kernel and max_pool2d_kernel) are used to reduce warp divergence in elementwise and pooling operations.
// Note: The max pooling kernel is specialized for 2x2 pooling (stride = 2).

torch::Tensor forward(
    torch::Tensor x,
    torch::Tensor conv1_weight, torch::Tensor conv1_bias,
    torch::Tensor conv2_weight, torch::Tensor conv2_bias,
    torch::Tensor fc1_weight, torch::Tensor fc1_bias,
    torch::Tensor fc2_weight, torch::Tensor fc2_bias,
    torch::Tensor fc3_weight, torch::Tensor fc3_bias
) {
    // Move inputs to CUDA and ensure contiguous layout
    x = x.to(torch::kCUDA);
    conv1_weight = conv1_weight.to(torch::kCUDA);
    conv1_bias = conv1_bias.to(torch::kCUDA);
    conv2_weight = conv2_weight.to(torch::kCUDA);
    conv2_bias = conv2_bias.to(torch::kCUDA);
    fc1_weight = fc1_weight.to(torch::kCUDA);
    fc1_bias = fc1_bias.to(torch::kCUDA);
    fc2_weight = fc2_weight.to(torch::kCUDA);
    fc2_bias = fc2_bias.to(torch::kCUDA);
    fc3_weight = fc3_weight.to(torch::kCUDA);
    fc3_bias = fc3_bias.to(torch::kCUDA);

    const int block_size = 256;
    const int num_blocks = 32;

    // First convolutional layer followed by ReLU activation
    auto conv1 = torch::conv2d(x, conv1_weight, conv1_bias, {1, 1});
    relu_kernel<<<num_blocks, block_size>>>(conv1.data_ptr<float>(), conv1.numel());

    // Custom 2x2 max pooling: allocate output tensor and apply kernel (stride = 2)
    auto pool1 = torch::empty({conv1.size(0), conv1.size(1), conv1.size(2) / 2, conv1.size(3) / 2}, conv1.options());
    max_pool2d_kernel<<<num_blocks, block_size>>>(conv1.data_ptr<float>(), pool1.data_ptr<float>(),
         conv1.size(0), conv1.size(1), conv1.size(2), conv1.size(3), 2);

    // Second convolutional layer followed by ReLU activation
    auto conv2 = torch::conv2d(pool1, conv2_weight, conv2_bias, {1, 1});
    relu_kernel<<<num_blocks, block_size>>>(conv2.data_ptr<float>(), conv2.numel());
    auto pool2 = torch::empty({conv2.size(0), conv2.size(1), conv2.size(2) / 2, conv2.size(3) / 2}, conv2.options());
    max_pool2d_kernel<<<num_blocks, block_size>>>(conv2.data_ptr<float>(), pool2.data_ptr<float>(),
         conv2.size(0), conv2.size(1), conv2.size(2), conv2.size(3), 2);

    // Flatten the pooled output
    auto flat = pool2.view({pool2.size(0), -1});

    // Fully connected layers use torch::linear to preserve high precision;
    // custom ReLU is applied between fc layers to reduce divergent branching in elementwise operations.
    auto fc1 = torch::linear(flat, fc1_weight, fc1_bias);
    relu_kernel<<<num_blocks, block_size>>>(fc1.data_ptr<float>(), fc1.numel());
    auto fc2 = torch::linear(fc1, fc2_weight, fc2_bias);
    relu_kernel<<<num_blocks, block_size>>>(fc2.data_ptr<float>(), fc2.numel());
    auto fc3 = torch::linear(fc2, fc3_weight, fc3_bias);

    return fc3;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "LeNet-5 forward pass");
}