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>
#include <cuda_runtime.h>

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

// Kernel: Each thread processes one element.
// Note: Atomic operations are not required here because each thread writes to a unique element.
__global__ void kernel_atomic_minimal(
    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) return;

    // Calculate channel index assuming input tensor has shape [N, C, D, H, W] in contiguous layout
    int64_t whd = width * height * depth;
    int64_t c = (idx / whd) % channels;

    // Load input using read-only cache
    float x = __ldg(&input[idx]);

    // Apply LeakyReLU
    float y = (x > 0.0f) ? x : 0.2f * x;

    // Add per-channel bias from sum_tensor (read-only). No atomic operation is needed because there's no write race.
    y += __ldg(&sum_tensor[c]);

    // Clamp the result between -1 and 1
    y = fmaxf(fminf(y, 1.0f), -1.0f);

    // Apply GELU activation
    y = gelu_activation(y);

    // Write final result. Since each thread writes to a unique location, no atomics are used.
    output[idx] = y;
}

// Kernel launcher
void launch_kernel_atomic_minimal(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;
    
    kernel_atomic_minimal<<<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 then applies the element-wise operations
torch::Tensor forward(
    torch::Tensor x,
    torch::Tensor conv_weight,
    torch::Tensor conv_bias,
    torch::Tensor sum_tensor) {
    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 float32");

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

    // Launch the kernel to apply LeakyReLU, addition, clamp, and GELU elementwise
    launch_kernel_atomic_minimal(x_conv, sum_tensor);

    return x_conv;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Forward pass with minimal atomic operations (CUDA)");
}