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,
    bias_flag: bool,
    conv_transpose: torch.Tensor,
    conv_transpose_bias: torch.Tensor,
) -> torch.Tensor:
    """
    Applies a 3D transposed convolution operation followed by softmax and sigmoid.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, in_channels, D, H, W)
        stride (int): Stride of the transposed convolution
        padding (int): Padding of the transposed convolution
        output_padding (int): Additional size added to output shape
        bias_flag (bool): Whether to use bias in conv_transpose
        conv_transpose (torch.Tensor): Transposed convolution weight tensor
        conv_transpose_bias (torch.Tensor): Bias tensor for transposed convolution

    Returns:
        torch.Tensor: Output tensor after applying transposed convolution, softmax and sigmoid
    """
    bias = conv_transpose_bias if bias_flag else None
    x = F.conv_transpose3d(
        x,
        conv_transpose,
        bias=bias,
        stride=stride,
        padding=padding,
        output_padding=output_padding,
    )
    x = F.softmax(x, dim=1)
    x = torch.sigmoid(x)
    return x


class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, applies Softmax and Sigmoid.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride,
        padding,
        output_padding,
        bias,
    ):
        super(Model, self).__init__()
        conv_transpose = nn.ConvTranspose3d(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            output_padding=output_padding,
            bias=bias,
        )
        self.conv_transpose_parameter = nn.Parameter(conv_transpose.weight)
        self.conv_transpose_bias = (
            nn.Parameter(
                conv_transpose.bias
                + torch.randn(
                    conv_transpose.bias.shape,
                    device=conv_transpose.bias.device,
                    dtype=conv_transpose.bias.dtype,
                )
                * 0.02
            )
            if bias
            else None
        )

    def forward(self, x, stride, padding, output_padding, bias, fn=module_fn):
        return fn(
            x,
            stride,
            padding,
            output_padding,
            bias,
            self.conv_transpose_parameter,
            self.conv_transpose_bias,
        )


batch_size = 16
in_channels = 32
out_channels = 64
D, H, W = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1
output_padding = 1
bias = True


def get_inputs():
    return [
        torch.randn(batch_size, in_channels, D, H, W),
        stride,
        padding,
        output_padding,
        bias,
    ]


def get_init_inputs():
    return [
        in_channels,
        out_channels,
        kernel_size,
        stride,
        padding,
        output_padding,
        bias,
    ]

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

// Declare constant memory for read-only parameters
__constant__ int c_channels;
__constant__ int c_depth;
__constant__ int c_height;
__constant__ int c_width;

// CUDA kernel: each thread processes one spatial position across all channels
// using the parameters stored in constant memory
template <typename scalar_t>
__global__ void const_params_fused_softmax_sigmoid_kernel(
    const scalar_t* __restrict__ input,
    scalar_t* __restrict__ output,
    const int batch) {

    // Compute spatial dimensions from constant memory
    int spatial = c_depth * c_height * c_width;
    int total_pixels = batch * spatial;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < total_pixels) {
        // Determine batch and pixel index
        int b = idx / spatial;
        int pixel_idx = idx % spatial;

        // Base offset for this batch and spatial position
        int base = b * c_channels * spatial + pixel_idx;
        int stride = spatial;  // jump to the same spatial location in the next channel

        // Online reduction: compute max and sum of exponentials for softmax
        scalar_t max_val = input[base];
        scalar_t sum_exp = 1.0f;  // corresponds to exp(input[base] - max_val)
        
        // Loop over channels to compute max and sum_exp
        for (int c = 1; c < c_channels; c++) {
            int pos = base + c * stride;
            scalar_t val = input[pos];
            if (val > max_val) {
                sum_exp = sum_exp * exp(max_val - val) + 1.0f;
                max_val = val;
            } else {
                sum_exp += exp(val - max_val);
            }
        }

        // Compute softmax then apply sigmoid for each channel
        for (int c = 0; c < c_channels; c++) {
            int pos = base + c * stride;
            scalar_t softmax_val = exp(input[pos] - max_val) / sum_exp;
            output[pos] = 1.0f / (1.0f + exp(-softmax_val));
        }
    }
}

// Forward function: applies conv_transpose3d and then the fused CUDA kernel
// with parameters for spatial dimensions stored in constant memory
torch::Tensor forward(
    torch::Tensor input,
    int stride,
    int padding,
    int output_padding,
    bool bias_flag,
    torch::Tensor conv_transpose,
    torch::Tensor conv_transpose_bias) {

    // Perform 3D transpose convolution using PyTorch's built-in function
    auto x = torch::conv_transpose3d(
        input,
        conv_transpose,
        bias_flag ? conv_transpose_bias : torch::Tensor(),
        stride,
        padding,
        output_padding
    );

    // Retrieve tensor dimensions
    const int batch = x.size(0);
    const int channels = x.size(1);
    const int depth = x.size(2);
    const int height = x.size(3);
    const int width = x.size(4);

    auto output = torch::empty_like(x);

    // Copy read-only parameters to constant memory
    cudaMemcpyToSymbol(c_channels, &channels, sizeof(int));
    cudaMemcpyToSymbol(c_depth, &depth, sizeof(int));
    cudaMemcpyToSymbol(c_height, &height, sizeof(int));
    cudaMemcpyToSymbol(c_width, &width, sizeof(int));

    int spatial = depth * height * width;
    int total_pixels = batch * spatial;
    int threads = 256;
    int blocks = (total_pixels + threads - 1) / threads;

    AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "const_params_fused_softmax_sigmoid_kernel", ([&] {
        const_params_fused_softmax_sigmoid_kernel<scalar_t><<<blocks, threads>>>(
            x.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            batch);
    }));

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Fused ConvTranspose3d with Softmax and Sigmoid using constant memory for parameters");
}