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,
    conv_transpose: torch.Tensor,
    conv_transpose_bias: torch.Tensor,
    bias: torch.Tensor,
) -> torch.Tensor:
    """
    Applies a 3D transposed convolution followed by bias addition and residual operations.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, in_channels, depth, height, width)
        stride (int): Stride of the transposed convolution
        padding (int): Padding of the transposed convolution
        output_padding (int): Additional size added to output shape
        conv_transpose (torch.Tensor): Transposed convolution weight tensor
        conv_transpose_bias (torch.Tensor): Bias tensor for transposed convolution
        bias (torch.Tensor): Bias tensor for addition

    Returns:
        torch.Tensor: Output tensor after applying operations
    """
    x = F.conv_transpose3d(
        x,
        conv_transpose,
        bias=conv_transpose_bias,
        stride=stride,
        padding=padding,
        output_padding=output_padding,
    )
    original_x = x.clone().detach()
    x = x + bias
    x = x + original_x
    x = x * original_x
    x = x + original_x
    return x


class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, followed by a sum,
    a residual add, a multiplication, and another residual add.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride,
        padding,
        output_padding,
        bias_shape,
    ):
        super(Model, self).__init__()
        conv_transpose = nn.ConvTranspose3d(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            output_padding=output_padding,
        )
        self.conv_transpose_parameter = conv_transpose.weight
        self.conv_transpose_bias = nn.Parameter(
            conv_transpose.bias + torch.ones_like(conv_transpose.bias) * 0.02
        )  # make sure its nonzero
        self.bias_parameter = nn.Parameter(torch.randn(bias_shape) * 0.02)

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


batch_size = 16
in_channels = 32
out_channels = 64
depth, height, width = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1
output_padding = 1
bias_shape = (out_channels, 1, 1, 1)


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


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

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

// Declare constant memory for element bias
__constant__ float const_element_bias[1024];

__global__ void constant_memory_kernel(
    const float* __restrict__ conv_output,
    float* __restrict__ output,
    int num_elements,
    int channels,
    int spatial_size
) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;

    // Process elements in groups of 4 (128-bit aligned)
    int num_vec = num_elements / 4;

    for (int i = tid; i < num_vec; i += stride) {
        int base_idx = i * 4;
        float4 in_val = __ldg(reinterpret_cast<const float4*>(conv_output) + i);
        float4 result;

        int c0 = (base_idx / spatial_size) % channels;
        int c1 = ((base_idx + 1) / spatial_size) % channels;
        int c2 = ((base_idx + 2) / spatial_size) % channels;
        int c3 = ((base_idx + 3) / spatial_size) % channels;

        float b0 = const_element_bias[c0];
        float b1 = const_element_bias[c1];
        float b2 = const_element_bias[c2];
        float b3 = const_element_bias[c3];

        result.x = in_val.x * (2.0f * in_val.x + b0 + 1.0f);
        result.y = in_val.y * (2.0f * in_val.y + b1 + 1.0f);
        result.z = in_val.z * (2.0f * in_val.z + b2 + 1.0f);
        result.w = in_val.w * (2.0f * in_val.w + b3 + 1.0f);

        reinterpret_cast<float4*>(output)[i] = result;
    }

    // Handle remaining elements
    int remaining_start = num_vec * 4;
    for (int i = tid + remaining_start; i < num_elements; i += stride) {
        int c = (i / spatial_size) % channels;
        float orig = __ldg(conv_output + i);
        float b = const_element_bias[c];
        output[i] = orig * (2.0f * orig + b + 1.0f);
    }
}

void copy_to_constant(const torch::Tensor& bias, int channels) {
    // Copy the bias data to constant memory
    cudaMemcpyToSymbol(const_element_bias, bias.data_ptr<float>(), channels * sizeof(float));
}

torch::Tensor forward(
    torch::Tensor x,
    int stride,
    int padding,
    int output_padding,
    torch::Tensor conv_transpose,
    torch::Tensor conv_transpose_bias,
    torch::Tensor bias
) {
    auto conv_result = torch::conv_transpose3d(
        x,
        conv_transpose,
        conv_transpose_bias,
        stride,
        padding,
        output_padding
    );

    auto sizes = conv_result.sizes();
    const int channels = sizes[1];
    const int spatial_size = sizes[2] * sizes[3] * sizes[4];
    const int num_elements = conv_result.numel();

    // Ensure the number of channels fits within the constant memory limit
    if (channels > 1024) {
        throw std::runtime_error("Number of channels exceeds constant memory limit.");
    }

    // Copy bias to constant memory
    copy_to_constant(bias, channels);

    auto output = torch::empty_like(conv_result);

    const int threads_per_block = 512;
    int max_blocks = (num_elements + threads_per_block - 1) / threads_per_block;
    int blocks = (max_blocks < 1024) ? max_blocks : 1024;

    constant_memory_kernel<<<blocks, threads_per_block>>>(
        conv_result.data_ptr<float>(),
        output.data_ptr<float>(),
        num_elements,
        channels,
        spatial_size
    );

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Fused ConvTranspose3D with Constant Memory Optimization");
}