The AI CUDA Engineer 👷

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


def module_fn(
    x: torch.Tensor,
    conv_weight: torch.Tensor,
    conv_bias: torch.Tensor,
    sum_tensor: torch.Tensor,
) -> torch.Tensor:
    """
    Applies 3D convolution, LeakyReLU, tensor addition, clamping and GELU activation.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, in_channels, depth, height, width)
        conv_weight (torch.Tensor): 3D convolution weight tensor of shape
            (out_channels, in_channels, kernel_size, kernel_size, kernel_size)
        conv_bias (torch.Tensor): Bias tensor for 3D convolution of shape (out_channels)
        sum_tensor (torch.Tensor): Tensor to add of shape (out_channels, 1, 1, 1)

    Returns:
        torch.Tensor: Output tensor after applying convolution, LeakyReLU, addition,
            clamping and GELU activation
    """
    x = F.conv3d(x, conv_weight, bias=conv_bias)
    x = F.leaky_relu(x, negative_slope=0.2)
    x = x + sum_tensor
    x = torch.clamp(x, min=-1.0, max=1.0)
    x = F.gelu(x)
    return x


class Model(nn.Module):
    """
    Model that performs a 3D convolution, applies LeakyReLU, sums with a tensor, clamps, and applies GELU activation.
    """

    def __init__(self, in_channels, out_channels, kernel_size, sum_tensor_shape):
        super(Model, self).__init__()
        conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.conv_weight = conv.weight
        self.conv_bias = conv.bias
        self.sum_tensor = nn.Parameter(torch.randn(sum_tensor_shape) * 0.02)

    def forward(self, x, fn=module_fn):
        return fn(x, self.conv_weight, self.conv_bias, self.sum_tensor)


batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 32, 32
kernel_size = 3
sum_tensor_shape = (out_channels, 1, 1, 1)


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


def get_init_inputs():
    return [in_channels, out_channels, kernel_size, sum_tensor_shape]

#include <torch/extension.h>
#include <ATen/ATen.h>
#include <math.h>

// Modular device functions for improved readability and maintainability
__device__ __forceinline__ float leaky_relu(float x) {
    // Apply LeakyReLU with negative slope 0.2
    return (x > 0.0f) ? x : 0.2f * x;
}

__device__ __forceinline__ float add_channel_bias(float x, int c, const float* sum_tensor) {
    // Add per-channel bias using __ldg for read-only cache
    return x + __ldg(&sum_tensor[c]);
}

__device__ __forceinline__ float clamp_val(float x) {
    // Clamp x between -1 and 1
    return fmaxf(fminf(x, 1.0f), -1.0f);
}

__device__ __forceinline__ float gelu_activation(float x) {
    // Apply GELU activation function
    float cube = x * x * x;
    float inner = 0.7978845608f * (x + 0.044715f * cube);
    return 0.5f * x * (1.0f + tanhf(inner));
}

// Composite function that applies the activation sequence
__device__ __forceinline__ float process_element(float x, int c, const float* sum_tensor) {
    x = leaky_relu(x);
    x = add_channel_bias(x, c, sum_tensor);
    x = clamp_val(x);
    return gelu_activation(x);
}

// Optimized kernel using modular device functions
__global__ void optimized_kernel(
    const float* __restrict__ input,
    const float* __restrict__ sum_tensor,
    float* __restrict__ output,
    int64_t num_elements,
    int64_t channels,
    int64_t depth,
    int64_t height,
    int64_t width) {

    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if(idx < num_elements) {
        // Compute channel index from linear index
        int c = (idx / (width * height * depth)) % channels;

        // Load input using read-only cache
        float x = __ldg(&input[idx]);
        
        // Process the element using modular functions
        output[idx] = process_element(x, c, sum_tensor);
    }
}

// Kernel launcher
void launch_optimized(torch::Tensor& input, torch::Tensor& sum_tensor) {
    int64_t num_elements = input.numel();
    const int threads = 256;
    const int blocks = (num_elements + threads - 1) / threads;

    optimized_kernel<<<blocks, threads>>>(
        input.data_ptr<float>(),
        sum_tensor.data_ptr<float>(),
        input.data_ptr<float>(),
        num_elements,
        input.size(1),   // channels
        input.size(2),   // depth
        input.size(3),   // height
        input.size(4));  // width

    cudaDeviceSynchronize();
}

// Forward function: performs 3D convolution and then applies the custom CUDA kernel
torch::Tensor forward(
    torch::Tensor x,
    torch::Tensor conv_weight,
    torch::Tensor conv_bias,
    torch::Tensor sum_tensor) {

    // Validate that all tensors are on CUDA and are float32
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(conv_weight.is_cuda(), "conv_weight must be a CUDA tensor");
    TORCH_CHECK(conv_bias.is_cuda(), "conv_bias must be a CUDA tensor");
    TORCH_CHECK(sum_tensor.is_cuda(), "sum_tensor must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be in float32");

    // Perform 3D convolution using cuDNN
    auto x_conv = at::conv3d(x, conv_weight, conv_bias);
    auto output = x_conv.contiguous();

    // Apply the custom elementwise operations using our optimized kernel
    launch_optimized(output, sum_tensor);
    
    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Optimized forward function with modular device functions (CUDA)");
}