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>
#include <algorithm>

// Tile dimensions for spatial tiling
#define TILE_W 16
#define TILE_H 16
#define KERNEL_SIZE 3  // Assuming a 3x3 kernel for simplicity

//---------------------------------------------------------------------
// Fused convolution and ReLU kernel
// Each thread computes one output pixel for a given sample and output channel.
// Grid organization: 
//   grid.x: covers output width in tiles
//   grid.y: covers output height in tiles
//   grid.z: covers (N * K) where N is the number of samples and K the number of output channels
// Within a block, blockDim.x = TILE_W and blockDim.y = TILE_H.
// The output position is computed as:
//    out_x = blockIdx.x * TILE_W + threadIdx.x
//    out_y = blockIdx.y * TILE_H + threadIdx.y
// Sample and channel indices are recovered from blockIdx.z as:
//    n = blockIdx.z / K
//    k = blockIdx.z % K
//---------------------------------------------------------------------

__global__ void fused_conv2d_relu_kernel(const float* __restrict__ input,
                                           const float* __restrict__ weight,
                                           const float* __restrict__ bias,
                                           float* __restrict__ output,
                                           int N, int C, int H, int W,
                                           int K, int pad, int stride) {
    // Recover sample and output channel from gridIdx.z
    int n = blockIdx.z / K;
    int k = blockIdx.z % K;

    // Compute output spatial coordinates
    int out_x = blockIdx.x * TILE_W + threadIdx.x;
    int out_y = blockIdx.y * TILE_H + threadIdx.y;

    // Check bounds (assuming output spatial dimensions are same as input HxW for stride=1 and pad=1)
    if (out_x < W && out_y < H) {
        // Initialize with bias value for this output channel
        float sum = bias[k];

        // Perform convolution sum over input channels and kernel window
        for (int c = 0; c < C; c++) {
            for (int i = 0; i < KERNEL_SIZE; i++) {
                for (int j = 0; j < KERNEL_SIZE; j++) {
                    // Compute input coordinates with stride and padding
                    int in_y = out_y * stride - pad + i;
                    int in_x = out_x * stride - pad + j;
                    
                    // Check if within input bounds
                    if (in_y >= 0 && in_y < H && in_x >= 0 && in_x < W) {
                        float in_val = input[n * C * H * W + c * H * W + in_y * W + in_x];
                        float wt = weight[k * C * KERNEL_SIZE * KERNEL_SIZE + c * KERNEL_SIZE * KERNEL_SIZE + i * KERNEL_SIZE + j];
                        sum += in_val * wt;
                    }
                }
            }
        }

        // Fused ReLU activation
        sum = fmaxf(sum, 0.0f);
        
        // Write the result
        output[n * K * H * W + k * H * W + out_y * W + out_x] = sum;
    }
}

//---------------------------------------------------------------------
// Custom convolution function that invokes the fused kernel for the first conv layer
//---------------------------------------------------------------------

torch::Tensor custom_conv2d_fused(torch::Tensor input, torch::Tensor weight, torch::Tensor bias, int stride = 1, int pad = 1) {
    const int N = input.size(0);
    const int C = input.size(1);
    const int H = input.size(2);
    const int W = input.size(3);
    const int K = weight.size(0);

    // Allocate output tensor. Here we assume output spatial dimensions equal H x W (for stride=1, pad=1)
    auto output = torch::empty({N, K, H, W}, input.options());

    // Define block dimensions
    dim3 block(TILE_W, TILE_H, 1);

    // Define grid dimensions
    int grid_x = (W + TILE_W - 1) / TILE_W; // number of tiles along width
    int grid_y = (H + TILE_H - 1) / TILE_H; // number of tiles along height
    // Combine N and K into grid.z so that each block processes one (n, k) pair
    dim3 grid(grid_x, grid_y, N * K);

    fused_conv2d_relu_kernel<<<grid, block>>>(
        input.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        output.data_ptr<float>(),
        N, C, H, W, K, pad, stride
    );

    return output;
}

//---------------------------------------------------------------------
// VGG16 forward pass using the fused convolution kernel for Block 1,
// and standard torch::conv2d (cuDNN) for subsequent 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 fused custom conv2d + ReLU for the first layer
    current = custom_conv2d_fused(current, conv_weights[0], conv_biases[0], /*stride=*/1, /*pad=*/1);
    // Second convolution in Block 1 using cuDNN optimized conv2d
    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 (Fused Conv2D + ReLU, CUDA)");
}