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,
    conv_weight: torch.Tensor,
    conv_bias: torch.Tensor,
) -> torch.Tensor:
    """
    Applies 3D convolution followed by Mish and Tanh activations.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, in_channels, D, H, W)
        stride (int): Stride of the convolution
        padding (int): Padding of the convolution
        conv_weight (torch.Tensor): Convolution weight tensor of shape
            (out_channels, in_channels, kernel_size, kernel_size, kernel_size)
        conv_bias (torch.Tensor): Bias tensor for convolution of shape (out_channels)

    Returns:
        torch.Tensor: Output tensor after applying convolution, Mish and Tanh activations
    """
    x = F.conv3d(x, conv_weight, bias=conv_bias, stride=stride, padding=padding)
    x = F.mish(x)
    x = torch.tanh(x)
    return x


class Model(nn.Module):
    """
    Model that performs a 3D convolution, applies Mish activation, and then applies Tanh activation.
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super(Model, self).__init__()
        conv = nn.Conv3d(
            in_channels, out_channels, kernel_size, stride=stride, padding=padding
        )
        self.conv_weight = nn.Parameter(conv.weight)
        self.conv_bias = nn.Parameter(
            conv.bias
            + torch.randn(
                conv.bias.shape, device=conv.bias.device, dtype=conv.bias.dtype
            )
            * 0.02
        )

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


batch_size = 16
in_channels = 3
out_channels = 16
D, H, W = 16, 32, 32
kernel_size = 3
stride = 1
padding = 0


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


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

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

#define BLOCK_SIZE 256
#define ELEMENTS_PER_THREAD 4
#define SHARED_MEM_SIZE (BLOCK_SIZE * ELEMENTS_PER_THREAD)

__device__ __forceinline__ float fused_mish_tanh_activation(float x) {
    float softplus = logf(1.0f + expf(x));
    float mish = x * tanhf(softplus);
    return tanhf(mish);
}

__global__ void manual_loop_unroll_kernel(
    float* __restrict__ output,
    const float* __restrict__ input,
    const int total_elements
) {
    __shared__ float shared_data[SHARED_MEM_SIZE];
    
    const int tid = threadIdx.x;
    const int global_idx = blockIdx.x * SHARED_MEM_SIZE + tid;
    
    // Load data into shared memory in coalesced manner
    #pragma unroll
    for (int i = 0; i < ELEMENTS_PER_THREAD; ++i) {
        const int idx = global_idx + i * BLOCK_SIZE;
        if (idx < total_elements) {
            shared_data[tid + i * BLOCK_SIZE] = input[idx];
        }
    }
    
    __syncthreads(); // Synchronize to ensure shared memory is fully populated

    // Manually unroll the loop for processing the data
    for (int i = 0; i < ELEMENTS_PER_THREAD; i += 4) {
        const int idx_0 = global_idx + i * BLOCK_SIZE;
        const int idx_1 = global_idx + (i + 1) * BLOCK_SIZE;
        const int idx_2 = global_idx + (i + 2) * BLOCK_SIZE;
        const int idx_3 = global_idx + (i + 3) * BLOCK_SIZE;
        if (idx_0 < total_elements) {
            shared_data[tid + i * BLOCK_SIZE] = fused_mish_tanh_activation(shared_data[tid + i * BLOCK_SIZE]);
            output[idx_0] = shared_data[tid + i * BLOCK_SIZE];
        }
        if (idx_1 < total_elements) {
            shared_data[tid + (i + 1) * BLOCK_SIZE] = fused_mish_tanh_activation(shared_data[tid + (i + 1) * BLOCK_SIZE]);
            output[idx_1] = shared_data[tid + (i + 1) * BLOCK_SIZE];
        }
        if (idx_2 < total_elements) {
            shared_data[tid + (i + 2) * BLOCK_SIZE] = fused_mish_tanh_activation(shared_data[tid + (i + 2) * BLOCK_SIZE]);
            output[idx_2] = shared_data[tid + (i + 2) * BLOCK_SIZE];
        }
        if (idx_3 < total_elements) {
            shared_data[tid + (i + 3) * BLOCK_SIZE] = fused_mish_tanh_activation(shared_data[tid + (i + 3) * BLOCK_SIZE]);
            output[idx_3] = shared_data[tid + (i + 3) * BLOCK_SIZE];
        }
    }
}

torch::Tensor module_fn_forward(
    torch::Tensor x,
    int64_t stride,
    int64_t padding,
    torch::Tensor conv_weight,
    torch::Tensor conv_bias
) {
    TORCH_CHECK(x.is_cuda(), "Input tensor x must be a CUDA tensor");
    TORCH_CHECK(conv_weight.is_cuda(), "Convolution weight must be a CUDA tensor");
    TORCH_CHECK(conv_bias.is_cuda(), "Convolution bias must be a CUDA tensor");

    auto x_conv = at::conv3d(
        x,
        conv_weight,
        conv_bias,
        {stride, stride, stride},
        {padding, padding, padding}
    );

    auto output = torch::empty_like(x_conv);
    const int total_elements = x_conv.numel();
    const int num_blocks = (total_elements + SHARED_MEM_SIZE - 1) / SHARED_MEM_SIZE;
    
    manual_loop_unroll_kernel<<<num_blocks, BLOCK_SIZE>>>(
        output.data_ptr<float>(),
        x_conv.data_ptr<float>(),
        total_elements
    );

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &module_fn_forward, "Manual loop unroll Mish-Tanh activation with shared memory (CUDA)");
}