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>

#define WARP_SIZE 32
#define BLOCK_SIZE 256
#define TILE_SIZE 16

// Kernel for 2D convolution with atomic operations optimization
__global__ void conv2d_atomic_kernel(const float* input, const float* weight, const float* bias,
                                      float* output, int N, int C, int H, int W,
                                      int K, int P, int stride) {
    extern __shared__ float shared_mem[];
    float* tile = shared_mem;

    int n = blockIdx.z;
    int k = blockIdx.y;
    int h = blockIdx.x * TILE_SIZE + threadIdx.y;
    int w = threadIdx.x;

    float sum = 0.0f;
    if (k < K && h < H && w < W) {
        sum = bias[k];

        #pragma unroll
        for (int c = 0; c < C; ++c) {
            #pragma unroll
            for (int kh = 0; kh < 3; ++kh) {
                #pragma unroll
                for (int kw = 0; kw < 3; ++kw) {
                    int ih = h - P + kh;
                    int iw = w - P + kw;
                    if (ih >= 0 && ih < H && iw >= 0 && iw < W) {
                        float in_val = input[n * C * H * W + c * H * W + ih * W + iw];
                        float weight_val = weight[k * C * 3 * 3 + c * 3 * 3 + kh * 3 + kw];
                        sum += in_val * weight_val;
                    }
                }
            }
        }
        tile[threadIdx.y * TILE_SIZE + threadIdx.x] = sum;
    }
    __syncthreads();

    if (threadIdx.y == 0 && w < W) {
        // Use atomic operations only once per output element computation
        atomicAdd(&output[n * K * H * W + k * H * W + h * W + w], sum);
    }
}

// Custom convolution function using the optimized atomic kernel
torch::Tensor custom_conv2d_atomic(torch::Tensor input, torch::Tensor weight, torch::Tensor bias) {
    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);
    const int kernel_size = weight.size(2);
    const int P = kernel_size / 2;

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

    dim3 block(TILE_SIZE, BLOCK_SIZE / TILE_SIZE);
    dim3 grid((W + TILE_SIZE - 1) / TILE_SIZE, (K + TILE_SIZE - 1) / TILE_SIZE, N);

    size_t shared_mem_size = TILE_SIZE * TILE_SIZE * sizeof(float);
    conv2d_atomic_kernel<<<grid, block, shared_mem_size>>>(
        input.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        output.data_ptr<float>(),
        N, C, H, W, K, P, 1
    );

    return output;
}

// VGG16 forward pass using custom convolution
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 custom conv2d for first layer
    current = custom_conv2d_atomic(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);

    // Rest of the network using optimized cudnn calls
    // 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 with aligned memory access
    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)");
}