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,
    bias: torch.Tensor,
) -> torch.Tensor:
    """
    Applies 3D convolution followed by ReLU, LeakyReLU, GELU, Sigmoid activations and bias addition.

    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)
        bias (torch.Tensor): Bias tensor for addition of shape (out_channels, 1, 1, 1)

    Returns:
        torch.Tensor: Output tensor after applying convolution and activations
    """
    x = F.conv3d(x, conv_weight, bias=conv_bias)
    x = F.relu(x)
    x = F.leaky_relu(x, negative_slope=0.01)
    x = F.gelu(x)
    x = torch.sigmoid(x)
    x = x + bias
    return x


class Model(nn.Module):
    """
    Model that performs a 3D convolution, applies ReLU, LeakyReLU, GELU, Sigmoid activations, and bias in sequence.
    """

    def __init__(self, in_channels, out_channels, kernel_size, bias_shape):
        super(Model, self).__init__()
        conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.bias = nn.Parameter(torch.randn(bias_shape) * 0.02)

        self.conv_weight = nn.Parameter(conv.weight)
        self.conv_bias = nn.Parameter(conv.bias)
        self.bias = self.bias

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


batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 32, 32
kernel_size = 3
bias_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, bias_shape]

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

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

#define WARP_SIZE 32
#define BLOCK_SIZE 256
#define ELEMENTS_PER_THREAD 8

__device__ __forceinline__ float process_element(float val, const float* bias, int bias_idx) {
    val = fmaxf(0.0f, val);
    val = fmaxf(0.01f * val, val);
    const float sqrt_2_over_pi = sqrtf(2.0f / M_PI);
    val = 0.5f * val * (1.0f + tanhf(sqrt_2_over_pi * (val + 0.044715f * powf(val, 3.0f))));
    val = 1.0f / (1.0f + expf(-val));
    return val + __ldg(&bias[bias_idx]);
}

__global__ void apply_activations_and_bias_kernel(
    float* __restrict__ output,
    const float* __restrict__ bias,
    int batch_size,
    int out_channels,
    int depth,
    int height,
    int width
) {
    const int tid = threadIdx.x;
    const int warp_id = tid / WARP_SIZE;
    const int lane_id = tid % WARP_SIZE;
    const int global_warp_id = blockIdx.x * (BLOCK_SIZE / WARP_SIZE) + warp_id;
    
    const int spatial_size = depth * height * width;
    const int total_elements = batch_size * out_channels * spatial_size;
    
    // Calculate base index for this thread
    int base_idx = global_warp_id * (WARP_SIZE * ELEMENTS_PER_THREAD) + lane_id;
    
    // Process elements in chunks of 4 when possible
    #pragma unroll
    for (int i = 0; i < ELEMENTS_PER_THREAD; i += 4) {
        int idx = base_idx + i * WARP_SIZE;
        
        if (idx + 3 * WARP_SIZE < total_elements) {
            // Load 4 elements
            float4 data;
            data.x = output[idx];
            data.y = output[idx + WARP_SIZE];
            data.z = output[idx + 2 * WARP_SIZE];
            data.w = output[idx + 3 * WARP_SIZE];
            
            // Calculate bias indices
            int bias_idx_x = (idx / spatial_size) % out_channels;
            int bias_idx_y = ((idx + WARP_SIZE) / spatial_size) % out_channels;
            int bias_idx_z = ((idx + 2 * WARP_SIZE) / spatial_size) % out_channels;
            int bias_idx_w = ((idx + 3 * WARP_SIZE) / spatial_size) % out_channels;
            
            // Process elements
            data.x = process_element(data.x, bias, bias_idx_x);
            data.y = process_element(data.y, bias, bias_idx_y);
            data.z = process_element(data.z, bias, bias_idx_z);
            data.w = process_element(data.w, bias, bias_idx_w);
            
            // Store results
            output[idx] = data.x;
            output[idx + WARP_SIZE] = data.y;
            output[idx + 2 * WARP_SIZE] = data.z;
            output[idx + 3 * WARP_SIZE] = data.w;
        } else {
            // Handle remaining elements
            for (int j = 0; j < 4; j++) {
                int curr_idx = idx + j * WARP_SIZE;
                if (curr_idx < total_elements) {
                    int bias_idx = (curr_idx / spatial_size) % out_channels;
                    float val = output[curr_idx];
                    output[curr_idx] = process_element(val, bias, bias_idx);
                }
            }
        }
    }
}

torch::Tensor module_fn_cuda(
    torch::Tensor x,
    torch::Tensor conv_weight,
    torch::Tensor conv_bias,
    torch::Tensor bias
) {
    CHECK_INPUT(x);
    CHECK_INPUT(conv_weight);
    CHECK_INPUT(conv_bias);
    CHECK_INPUT(bias);

    auto output = torch::conv3d(x, conv_weight, conv_bias);

    int batch_size = output.size(0);
    int out_channels = output.size(1);
    int depth = output.size(2);
    int height = output.size(3);
    int width = output.size(4);

    const int total_elements = batch_size * out_channels * depth * height * width;
    const int elements_per_block = BLOCK_SIZE * ELEMENTS_PER_THREAD;
    const int num_blocks = (total_elements + elements_per_block - 1) / elements_per_block;
    
    apply_activations_and_bias_kernel<<<num_blocks, BLOCK_SIZE>>>(
        output.data_ptr<float>(),
        bias.data_ptr<float>(),
        batch_size,
        out_channels,
        depth,
        height,
        width
    );

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &module_fn_cuda, "CUDA implementation of module_fn");
}