The AI CUDA Engineer 👷

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


def module_fn(
    x: torch.Tensor,
    conv_weights: nn.ParameterList,
    conv_biases: nn.ParameterList,
    fc_weights: nn.ParameterList,
    fc_biases: nn.ParameterList,
    is_training: bool,
) -> torch.Tensor:
    """
    Implements the VGG16 module.

    Args:
        x (torch.Tensor): Input tensor, shape (batch_size, in_channels, height, width)
        conv_weights (nn.ParameterList): List of convolutional weights
        conv_biases (nn.ParameterList): List of convolutional biases
        fc_weights (nn.ParameterList): List of fully connected weights
        fc_biases (nn.ParameterList): List of fully connected biases
        is_training (bool): Whether in training mode

    Returns:
        torch.Tensor: Output tensor, shape (batch_size, num_classes)
    """
    # Block 1
    x = F.conv2d(x, conv_weights[0], conv_biases[0], padding=1)
    x = F.relu(x)
    x = F.conv2d(x, conv_weights[1], conv_biases[1], padding=1)
    x = F.relu(x)
    x = F.max_pool2d(x, kernel_size=2, stride=2)

    # Block 2
    x = F.conv2d(x, conv_weights[2], conv_biases[2], padding=1)
    x = F.relu(x)
    x = F.conv2d(x, conv_weights[3], conv_biases[3], padding=1)
    x = F.relu(x)
    x = F.max_pool2d(x, kernel_size=2, stride=2)

    # Block 3
    x = F.conv2d(x, conv_weights[4], conv_biases[4], padding=1)
    x = F.relu(x)
    x = F.conv2d(x, conv_weights[5], conv_biases[5], padding=1)
    x = F.relu(x)
    x = F.conv2d(x, conv_weights[6], conv_biases[6], padding=1)
    x = F.relu(x)
    x = F.max_pool2d(x, kernel_size=2, stride=2)

    # Block 4
    x = F.conv2d(x, conv_weights[7], conv_biases[7], padding=1)
    x = F.relu(x)
    x = F.conv2d(x, conv_weights[8], conv_biases[8], padding=1)
    x = F.relu(x)
    x = F.conv2d(x, conv_weights[9], conv_biases[9], padding=1)
    x = F.relu(x)
    x = F.max_pool2d(x, kernel_size=2, stride=2)

    # Block 5
    x = F.conv2d(x, conv_weights[10], conv_biases[10], padding=1)
    x = F.relu(x)
    x = F.conv2d(x, conv_weights[11], conv_biases[11], padding=1)
    x = F.relu(x)
    x = F.conv2d(x, conv_weights[12], conv_biases[12], padding=1)
    x = F.relu(x)
    x = F.max_pool2d(x, kernel_size=2, stride=2)

    # Classifier
    x = torch.flatten(x, 1)
    x = F.linear(x, fc_weights[0], fc_biases[0])
    x = F.relu(x)
    x = F.dropout(x, p=0.0, training=is_training)
    x = F.linear(x, fc_weights[1], fc_biases[1])
    x = F.relu(x)
    x = F.dropout(x, p=0.0, training=is_training)
    x = F.linear(x, fc_weights[2], fc_biases[2])

    return x


class Model(nn.Module):
    def __init__(self, num_classes=1000):
        super(Model, self).__init__()

        # Extract convolutional parameters
        self.conv_weights = nn.ParameterList()
        self.conv_biases = nn.ParameterList()

        # Block 1
        conv = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.conv_weights.append(nn.Parameter(conv.weight.data.clone()))
        self.conv_biases.append(nn.Parameter(conv.bias.data.clone()))

        conv = nn.Conv2d(64, 64, kernel_size=3, padding=1)
        self.conv_weights.append(nn.Parameter(conv.weight.data.clone()))
        self.conv_biases.append(nn.Parameter(conv.bias.data.clone()))

        # Block 2
        conv = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.conv_weights.append(nn.Parameter(conv.weight.data.clone()))
        self.conv_biases.append(nn.Parameter(conv.bias.data.clone()))

        conv = nn.Conv2d(128, 128, kernel_size=3, padding=1)
        self.conv_weights.append(nn.Parameter(conv.weight.data.clone()))
        self.conv_biases.append(nn.Parameter(conv.bias.data.clone()))

        # Block 3
        conv = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.conv_weights.append(nn.Parameter(conv.weight.data.clone()))
        self.conv_biases.append(nn.Parameter(conv.bias.data.clone()))

        conv = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.conv_weights.append(nn.Parameter(conv.weight.data.clone()))
        self.conv_biases.append(nn.Parameter(conv.bias.data.clone()))

        conv = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.conv_weights.append(nn.Parameter(conv.weight.data.clone()))
        self.conv_biases.append(nn.Parameter(conv.bias.data.clone()))

        # Block 4
        conv = nn.Conv2d(256, 512, kernel_size=3, padding=1)
        self.conv_weights.append(nn.Parameter(conv.weight.data.clone()))
        self.conv_biases.append(nn.Parameter(conv.bias.data.clone()))

        conv = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv_weights.append(nn.Parameter(conv.weight.data.clone()))
        self.conv_biases.append(nn.Parameter(conv.bias.data.clone()))

        conv = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv_weights.append(nn.Parameter(conv.weight.data.clone()))
        self.conv_biases.append(nn.Parameter(conv.bias.data.clone()))

        # Block 5
        conv = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv_weights.append(nn.Parameter(conv.weight.data.clone()))
        self.conv_biases.append(nn.Parameter(conv.bias.data.clone()))

        conv = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv_weights.append(nn.Parameter(conv.weight.data.clone()))
        self.conv_biases.append(nn.Parameter(conv.bias.data.clone()))

        conv = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv_weights.append(nn.Parameter(conv.weight.data.clone()))
        self.conv_biases.append(nn.Parameter(conv.bias.data.clone()))

        # Extract fully connected parameters
        self.fc_weights = nn.ParameterList()
        self.fc_biases = nn.ParameterList()

        fc = nn.Linear(512 * 7 * 7, 4096)
        self.fc_weights.append(nn.Parameter(fc.weight.data.clone()))
        self.fc_biases.append(nn.Parameter(fc.bias.data.clone()))

        fc = nn.Linear(4096, 4096)
        self.fc_weights.append(nn.Parameter(fc.weight.data.clone()))
        self.fc_biases.append(nn.Parameter(fc.bias.data.clone()))

        fc = nn.Linear(4096, num_classes)
        self.fc_weights.append(nn.Parameter(fc.weight.data.clone()))
        self.fc_biases.append(nn.Parameter(fc.bias.data.clone()))

    def forward(self, x, fn=module_fn):
        return fn(
            x,
            self.conv_weights,
            self.conv_biases,
            self.fc_weights,
            self.fc_biases,
            self.training,
        )


# Test code
batch_size = 10
num_classes = 1000


def get_inputs():
    return [torch.randn(batch_size, 3, 224, 224)]


def get_init_inputs():
    return [num_classes]

#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

// This kernel is specialized for VGG16's first convolution layer where input has 3 channels and kernel size is 3.
// It manually unrolls the loops over the 3x3 kernel and the small channel dimension (C==3) to reduce loop overhead.

__global__ void conv2d_first_layer_kernel(
    const float* __restrict__ input,   // shape: (N, 3, H, W)
    const float* __restrict__ weight,  // shape: (K, 3, 3, 3) -> K x 27
    const float* __restrict__ bias,    // shape: (K)
    float* __restrict__ output,        // shape: (N, K, H, W)
    int N, int H, int W, int K
) {
    // Compute spatial coordinates
    int w = blockIdx.x * blockDim.x + threadIdx.x;
    int h = blockIdx.y * blockDim.y + threadIdx.y;
    // Combine batch and output channel in blockIdx.z
    int nk = blockIdx.z;  // n * K + k
    int n = nk / K;
    int k = nk % K;

    if (h >= H || w >= W) return;

    // Padding = 1 assumed
    float sum = bias[k];

    // For each input channel (C=3) and 3x3 kernel, completely unroll the loops.
    // Compute base indexes
    int in_base = n * 3 * H * W;  // starting offset for input of batch n
    int weight_base = k * 27;     // each filter has 27 elements

    // Unroll for channel 0
    {
        int c = 0;
        int offset = in_base + c * H * W;
        int r = h - 1; // starting row index in input
        int cidx;
        float v0, v1, v2;

        // position (0,0)
        cidx = w - 1;
        v0 = ((r >= 0 && r < H && cidx >= 0 && cidx < W) ? input[offset + r * W + cidx] : 0.0f);
        // position (0,1)
        cidx = w;
        v1 = ((r >= 0 && r < H && cidx >= 0 && cidx < W) ? input[offset + r * W + cidx] : 0.0f);
        // position (0,2)
        cidx = w + 1;
        v2 = ((r >= 0 && r < H && cidx < W) ? input[offset + r * W + cidx] : 0.0f);
        sum += v0 * weight[weight_base + 0];
        sum += v1 * weight[weight_base + 1];
        sum += v2 * weight[weight_base + 2];
    }

    // Unroll for channel 1
    {
        int c = 1;
        int offset = in_base + c * H * W;
        int weight_base_c = weight_base + 9; // next 9 weights for channel 1
        int r = h - 1;
        int cidx;
        float v0, v1, v2;

        cidx = w - 1;
        v0 = ((r >= 0 && r < H && cidx >= 0 && cidx < W) ? input[offset + r * W + cidx] : 0.0f);
        cidx = w;
        v1 = ((r >= 0 && r < H && cidx >= 0 && cidx < W) ? input[offset + r * W + cidx] : 0.0f);
        cidx = w + 1;
        v2 = ((r >= 0 && r < H && cidx < W) ? input[offset + r * W + cidx] : 0.0f);
        sum += v0 * weight[weight_base_c + 0];
        sum += v1 * weight[weight_base_c + 1];
        sum += v2 * weight[weight_base_c + 2];
    }

    // Unroll for channel 2
    {
        int c = 2;
        int offset = in_base + c * H * W;
        int weight_base_c = weight_base + 18; // next 9 weights for channel 2
        int r = h - 1;
        int cidx;
        float v0, v1, v2;

        cidx = w - 1;
        v0 = ((r >= 0 && r < H && cidx >= 0 && cidx < W) ? input[offset + r * W + cidx] : 0.0f);
        cidx = w;
        v1 = ((r >= 0 && r < H && cidx >= 0 && cidx < W) ? input[offset + r * W + cidx] : 0.0f);
        cidx = w + 1;
        v2 = ((r >= 0 && r < H && cidx < W) ? input[offset + r * W + cidx] : 0.0f);
        sum += v0 * weight[weight_base_c + 0];
        sum += v1 * weight[weight_base_c + 1];
        sum += v2 * weight[weight_base_c + 2];
    }

    output[n * K * H * W + k * H * W + h * W + w] = sum;
}

// Host function to launch the specialized first-layer convolution kernel
torch::Tensor custom_conv2d_first_layer(torch::Tensor input, torch::Tensor weight, torch::Tensor bias) {
    int N = input.size(0);
    int H = input.size(2);
    int W = input.size(3);
    int K = weight.size(0); // output channels

    auto output = torch::empty({N, K, H, W}, input.options());

    dim3 block(16, 16);
    dim3 grid((W + block.x - 1) / block.x, (H + block.y - 1) / block.y, N * K);

    conv2d_first_layer_kernel<<<grid, block>>>(
        input.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        output.data_ptr<float>(),
        N, H, W, K
    );

    return output;
}

// VGG16 forward pass using the specialized unrolled conv kernel for the first layer
// and standard torch calls for the remaining layers.

torch::Tensor vgg16_forward_cuda(
    torch::Tensor x,
    std::vector<torch::Tensor> conv_weights,
    std::vector<torch::Tensor> conv_biases,
    std::vector<torch::Tensor> fc_weights,
    std::vector<torch::Tensor> fc_biases,
    bool is_training
) {
    auto current = x;

    // Block 1 - use the custom unrolled convolution for the first layer (assumes input channels == 3 and kernel size == 3, padding=1)
    current = custom_conv2d_first_layer(current, conv_weights[0], conv_biases[0]);
    current = torch::relu(current);
    current = torch::conv2d(current, conv_weights[1], conv_biases[1], /*stride=*/1, /*padding=*/1);
    current = torch::relu(current);
    current = torch::max_pool2d(current, /*kernel_size=*/2, /*stride=*/2);

    // Block 2
    current = torch::conv2d(current, conv_weights[2], conv_biases[2], /*stride=*/1, /*padding=*/1);
    current = torch::relu(current);
    current = torch::conv2d(current, conv_weights[3], conv_biases[3], /*stride=*/1, /*padding=*/1);
    current = torch::relu(current);
    current = torch::max_pool2d(current, /*kernel_size=*/2, /*stride=*/2);

    // Block 3
    current = torch::conv2d(current, conv_weights[4], conv_biases[4], /*stride=*/1, /*padding=*/1);
    current = torch::relu(current);
    current = torch::conv2d(current, conv_weights[5], conv_biases[5], /*stride=*/1, /*padding=*/1);
    current = torch::relu(current);
    current = torch::conv2d(current, conv_weights[6], conv_biases[6], /*stride=*/1, /*padding=*/1);
    current = torch::relu(current);
    current = torch::max_pool2d(current, /*kernel_size=*/2, /*stride=*/2);

    // Block 4
    current = torch::conv2d(current, conv_weights[7], conv_biases[7], /*stride=*/1, /*padding=*/1);
    current = torch::relu(current);
    current = torch::conv2d(current, conv_weights[8], conv_biases[8], /*stride=*/1, /*padding=*/1);
    current = torch::relu(current);
    current = torch::conv2d(current, conv_weights[9], conv_biases[9], /*stride=*/1, /*padding=*/1);
    current = torch::relu(current);
    current = torch::max_pool2d(current, /*kernel_size=*/2, /*stride=*/2);

    // Block 5
    current = torch::conv2d(current, conv_weights[10], conv_biases[10], /*stride=*/1, /*padding=*/1);
    current = torch::relu(current);
    current = torch::conv2d(current, conv_weights[11], conv_biases[11], /*stride=*/1, /*padding=*/1);
    current = torch::relu(current);
    current = torch::conv2d(current, conv_weights[12], conv_biases[12], /*stride=*/1, /*padding=*/1);
    current = torch::relu(current);
    current = torch::max_pool2d(current, /*kernel_size=*/2, /*stride=*/2);

    // Classifier
    current = current.flatten(1);
    current = torch::linear(current, fc_weights[0], fc_biases[0]);
    current = torch::relu(current);
    if (is_training) {
        current = torch::dropout(current, /*p=*/0.0, /*train=*/true);
    }
    current = torch::linear(current, fc_weights[1], fc_biases[1]);
    current = torch::relu(current);
    if (is_training) {
        current = torch::dropout(current, /*p=*/0.0, /*train=*/true);
    }
    current = torch::linear(current, fc_weights[2], fc_biases[2]);

    return current;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &vgg16_forward_cuda, "VGG16 forward (CUDA)");
}