The AI CUDA Engineer 👷

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


def module_fn(
    x: torch.Tensor,
    stride: int,
    padding: int,
    output_padding: int,
    pool_kernel_size: int,
    pool_stride: int,
    pool_padding: int,
    conv_transpose: torch.Tensor,
    conv_transpose_bias: torch.Tensor,
    subtract: torch.Tensor,
) -> torch.Tensor:
    """
    Applies sequence of operations:
        - ConvTranspose3d
        - MaxPool3d
        - Softmax
        - Subtract
        - Swish
        - Max

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, in_channels, depth, height, width)
        stride (int): Stride for conv transpose
        padding (int): Padding for conv transpose
        output_padding (int): Output padding for conv transpose
        pool_kernel_size (int): Kernel size for max pooling
        pool_stride (int): Stride for max pooling
        pool_padding (int): Padding for max pooling
        conv_transpose (torch.Tensor): Weight tensor for transposed convolution
        conv_transpose_bias (torch.Tensor): Bias tensor for transposed convolution
        subtract (torch.Tensor): Subtraction parameter tensor
    """
    x = F.conv_transpose3d(
        x,
        conv_transpose,
        bias=conv_transpose_bias,
        stride=stride,
        padding=padding,
        output_padding=output_padding,
    )
    x = F.max_pool3d(
        x, kernel_size=pool_kernel_size, stride=pool_stride, padding=pool_padding
    )
    x = F.softmax(x, dim=1)
    x = x - subtract.view(1, -1, 1, 1, 1)
    x = torch.sigmoid(x) * x  # Swish
    x = torch.max(x, dim=1)[0]
    return x


class Model(nn.Module):
    """
    A model that performs a sequence of operations:
        - ConvTranspose3d
        - MaxPool3d
        - Softmax
        - Subtract
        - Swish
        - Max
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride,
        padding,
        output_padding,
        pool_kernel_size,
        pool_stride,
        pool_padding,
    ):
        super(Model, self).__init__()
        conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size)
        self.conv_transpose_parameter = conv_transpose.weight
        self.conv_transpose_bias = conv_transpose.bias
        self.subtract_parameter = nn.Parameter(torch.randn(out_channels) * 0.02)

    def forward(
        self,
        x,
        stride,
        padding,
        output_padding,
        pool_kernel_size,
        pool_stride,
        pool_padding,
        fn=module_fn,
    ):
        return fn(
            x,
            stride,
            padding,
            output_padding,
            pool_kernel_size,
            pool_stride,
            pool_padding,
            self.conv_transpose_parameter,
            self.conv_transpose_bias,
            self.subtract_parameter,
        )


batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1
output_padding = 1
pool_kernel_size = 2
pool_stride = 2
pool_padding = 0


def get_inputs():
    return [
        torch.randn(batch_size, in_channels, depth, height, width),
        stride,
        padding,
        output_padding,
        pool_kernel_size,
        pool_stride,
        pool_padding,
    ]


def get_init_inputs():
    return [
        in_channels,
        out_channels,
        kernel_size,
        stride,
        padding,
        output_padding,
        pool_kernel_size,
        pool_stride,
        pool_padding,
    ]

#include <torch/extension.h>
#include <pybind11/pybind11.h>

namespace py = pybind11;

// Declare constant memory for kernel parameters
__constant__ int STRIDE[3];
__constant__ int PADDING[3];
__constant__ int OUTPUT_PADDING[3];
__constant__ int POOL_PARAMS[9];  // kernel_size, stride, padding for each dimension

// Declare constant memory for subtract tensor
__constant__ float SUBTRACT_TENSOR[1024];  // Assuming max channel size is 1024

__device__ __forceinline__ float swish(float x) {
    return x * (1.0f / (1.0f + expf(-x)));
}

torch::Tensor forward(
    torch::Tensor x,
    int64_t stride,
    int64_t padding,
    int64_t output_padding,
    int64_t pool_kernel_size,
    int64_t pool_stride,
    int64_t pool_padding,
    torch::Tensor conv_transpose_weight,
    torch::Tensor conv_transpose_bias,
    torch::Tensor subtract_tensor
) {
    // Copy constant parameters to device memory
    int h_stride[3] = {stride, stride, stride};
    int h_padding[3] = {padding, padding, padding};
    int h_output_padding[3] = {output_padding, output_padding, output_padding};
    int h_pool_params[9] = {
        pool_kernel_size, pool_kernel_size, pool_kernel_size,
        pool_stride, pool_stride, pool_stride,
        pool_padding, pool_padding, pool_padding
    };
    
    cudaMemcpyToSymbol(STRIDE, h_stride, sizeof(int) * 3);
    cudaMemcpyToSymbol(PADDING, h_padding, sizeof(int) * 3);
    cudaMemcpyToSymbol(OUTPUT_PADDING, h_output_padding, sizeof(int) * 3);
    cudaMemcpyToSymbol(POOL_PARAMS, h_pool_params, sizeof(int) * 9);

    // Copy subtract tensor to constant memory
    cudaMemcpyToSymbol(SUBTRACT_TENSOR, subtract_tensor.data_ptr<float>(), subtract_tensor.numel() * sizeof(float));

    // Transposed convolution
    auto out = at::conv_transpose3d(
        x,
        conv_transpose_weight,
        conv_transpose_bias,
        {stride, stride, stride},
        {padding, padding, padding},
        {output_padding, output_padding, output_padding},
        1,
        {1, 1, 1}
    );

    // MaxPool
    out = at::max_pool3d(
        out,
        {pool_kernel_size, pool_kernel_size, pool_kernel_size},
        {pool_stride, pool_stride, pool_stride},
        {pool_padding, pool_padding, pool_padding}
    );

    // Softmax along channel dimension
    out = at::softmax(out, 1, /*dtype=*/c10::nullopt);

    // Subtract using constant memory
    auto sub_view = subtract_tensor.view({1, -1, 1, 1, 1});
    out = out - sub_view;

    // Swish
    out = out * at::sigmoid(out);

    // Max over channel dimension
    out = std::get<0>(out.max(/*dim=*/1, /*keepdim=*/false));

    return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "CUDA forward pass for the module");
}