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>

// Kernel that performs both max pooling operations in a single pass with balanced workload distribution
__global__ void balanced_maxpool3d_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    const int N, const int C, const int D, const int H, const int W
) {
    // Final output dimensions after both pooling operations
    const int D_out = D / 4;
    const int H_out = H / 4;
    const int W_out = W / 4;

    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) {
        // Map to input coordinates
        int d_in_base = d_out * 4;
        int h_in_base = h_out * 4;
        int w_in_base = w_out * 4;

        float max_val = -FLT_MAX;
        
        // Single pass over 4x4x4 input region
        #pragma unroll
        for (int dz = 0; dz < 4; ++dz) {
            #pragma unroll
            for (int dy = 0; dy < 4; ++dy) {
                #pragma unroll
                for (int dx = 0; dx < 4; ++dx) {
                    int d_in = d_in_base + dz;
                    int h_in = h_in_base + dy;
                    int w_in = w_in_base + dx;
                    
                    int input_idx = n * (C * D * H * W) + 
                                  c * (D * H * W) + 
                                  d_in * (H * W) + 
                                  h_in * W + w_in;
                    
                    float val = input[input_idx];
                    max_val = max(max_val, val);
                }
            }
        }

        // Write final output
        int output_idx = 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_idx] = max_val;
    }
}

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();

    // Use ATen ops for conv3d and softmax as they're highly optimized
    auto conv_output = at::conv3d(x, conv_weight, conv_bias, {1, 1, 1}, {0, 0, 0});
    auto softmax_output = at::softmax(conv_output, 1);

    // Get dimensions
    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);

    // Final output dimensions
    int D_out = D / 4;
    int H_out = H / 4;
    int W_out = W / 4;

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

    // Launch kernel with balanced block size
    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);

    balanced_maxpool3d_kernel<<<grid, block>>>(
        softmax_output.data_ptr<float>(),
        output.data_ptr<float>(),
        N, C, D, H, W
    );

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Balanced max pooling CUDA forward function");
}