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,
) -> torch.Tensor:
    """Applies 3D convolution, softmax activation, and two max pooling operations.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, in_channels, depth, height, width)
        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, softmax and max pooling,
            with shape (batch_size, out_channels, depth', height', width') where:
            depth' = ((depth - kernel_size + 1) // 4)
            height' = ((height - kernel_size + 1) // 4)
            width' = ((width - kernel_size + 1) // 4)
            The //4 comes from two max pooling operations with kernel_size=2
    """
    x = F.conv3d(x, conv_weight, conv_bias, stride=1, padding=0)
    x = F.softmax(x, dim=1)
    x = F.max_pool3d(x, kernel_size=2)
    x = F.max_pool3d(x, kernel_size=2)
    return x


class Model(nn.Module):
    """
    Model that performs a 3D convolution, applies Softmax, and performs two max pooling operations.
    """

    def __init__(self, in_channels, out_channels, kernel_size, pool_kernel_size):
        super(Model, self).__init__()
        conv = nn.Conv3d(in_channels, out_channels, kernel_size, padding=1)
        self.conv_weight = nn.Parameter(conv.weight)
        self.conv_bias = nn.Parameter(conv.bias)

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


batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 32, 32
kernel_size = 3
pool_kernel_size = 2


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


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

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

__global__ void max_pool3d_coalesced_kernel(const float* __restrict__ input, float* __restrict__ output,
                                             const int N, const int C, const int D, const int H, const int W) {
    int D_out = D / 2;
    int H_out = H / 2;
    int W_out = W / 2;

    int nc_d = blockIdx.z;  
    int d_out = nc_d % D_out;
    int nc = nc_d / D_out;
    int c = nc % C;
    int n = nc / C;

    int h_out = blockIdx.y * blockDim.y + threadIdx.y;
    int w_out = blockIdx.x * blockDim.x + threadIdx.x;

    if (w_out < W_out && h_out < H_out) {
        int d_in = d_out * 2;
        int h_in = h_out * 2;
        int w_in = w_out * 2;
        int input_base = n * (C * D * H * W) + c * (D * H * W) + d_in * (H * W) + h_in * W + w_in;

        float max_val = -FLT_MAX;
        #pragma unroll
        for (int dz = 0; dz < 2; ++dz) {
            #pragma unroll
            for (int dy = 0; dy < 2; ++dy) {
                #pragma unroll
                for (int dx = 0; dx < 2; ++dx) {
                    int index = input_base + dz * (H * W) + dy * W + dx;
                    max_val = fmaxf(max_val, input[index]);
                }
            }
        }

        int output_index = n * (C * D_out * H_out * W_out) +
                           c * (D_out * H_out * W_out) +
                           d_out * (H_out * W_out) +
                           h_out * W_out +
                           w_out;
        output[output_index] = max_val;
    }
}

// Optimized forward function
torch::Tensor forward(
    torch::Tensor x,
    torch::Tensor conv_weight,
    torch::Tensor conv_bias
) {
    x = x.contiguous();
    conv_weight = conv_weight.contiguous();
    conv_bias = conv_bias.contiguous();

    auto conv_output = at::conv3d(x, conv_weight, conv_bias, {1, 1, 1}, {0, 0, 0});
    auto softmax_output = at::softmax(conv_output, 1);

    int N = softmax_output.size(0);
    int C = softmax_output.size(1);
    int D = softmax_output.size(2);
    int H = softmax_output.size(3);
    int W = softmax_output.size(4);

    int D_out = D / 2;
    int H_out = H / 2;
    int W_out = W / 2;

    auto pool1 = at::empty({N, C, D_out, H_out, W_out}, softmax_output.options());

    dim3 block(16, 16, 1);
    dim3 grid((W_out + block.x - 1) / block.x,
              (H_out + block.y - 1) / block.y,
              N * C * D_out);

    const float* input_ptr = softmax_output.data_ptr<float>();
    float* output_ptr = pool1.data_ptr<float>();

    max_pool3d_coalesced_kernel<<<grid, block>>>(input_ptr, output_ptr, N, C, D, H, W);

    int D1 = D_out;
    int H1 = H_out;
    int W1 = W_out;
    int D_out2 = D1 / 2;
    int H_out2 = H1 / 2;
    int W_out2 = W1 / 2;

    auto pool2 = at::empty({N, C, D_out2, H_out2, W_out2}, softmax_output.options());

    dim3 grid2((W_out2 + block.x - 1) / block.x,
               (H_out2 + block.y - 1) / block.y,
               N * C * D_out2);

    const float* input_ptr2 = pool1.data_ptr<float>();
    float* output_ptr2 = pool2.data_ptr<float>();

    max_pool3d_coalesced_kernel<<<grid2, block>>>(input_ptr2, output_ptr2, N, C, D1, H1, W1);
    cudaDeviceSynchronize();

    return pool2;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Optimized CUDA double max pooling with coalesced memory accesses");
}