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>

// This fused kernel performs two successive max poolings (each with a window of 2x2x2 and stride 2)
// in a single kernel launch by directly computing the maximum over a 4x4x4 window from the softmax output.
// The final output shape is [N, C, D/4, H/4, W/4].

__global__ void max_pool3d_fused_kernel(const float* __restrict__ input, float* __restrict__ output,
                                          const int N, const int C, const int D, const int H, const int W) {
    // The input is assumed to have shape [N, C, D, H, W] from softmax (D, H, W are divisible by 4).
    // The fused pooling downsamples by a factor of 4 in each spatial dimension.
    int D_out = D / 4;
    int H_out = H / 4;
    int W_out = W / 4;

    // Combine the sample, channel, and depth index (for the output) into gridDim.z
    // Each block in grid.z corresponds to one slice: (n, c, d_out).
    int idx = blockIdx.z;  // Range: 0 to (N * C * D_out - 1)
    int d_out = idx % D_out;
    int nc = idx / D_out;
    int c = nc % C;
    int n = nc / C;

    // Determine the output height and width indices
    int h_out_idx = blockIdx.y * blockDim.y + threadIdx.y;
    int w_out_idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (w_out_idx < W_out && h_out_idx < H_out) {
        // Each output element corresponds to a 4x4x4 window in the input starting here:
        int d_in_start = d_out * 4;
        int h_in_start = h_out_idx * 4;
        int w_in_start = w_out_idx * 4;

        float max_val = -FLT_MAX;

        // Calculate the starting index for (n, c, d_in_start, h_in_start, w_in_start)
        int base_index = n * (C * D * H * W) + c * (D * H * W) + d_in_start * (H * W) + h_in_start * W + w_in_start;

        // Iterate over the 4x4x4 window
        for (int dd = 0; dd < 4; dd++) {
            for (int hh = 0; hh < 4; hh++) {
                for (int ww = 0; ww < 4; ww++) {
                    int index = base_index + dd * (H * W) + hh * W + ww;
                    float val = input[index];
                    if (val > max_val) {
                        max_val = val;
                    }
                }
            }
        }

        // Compute the output index for [n, c, d_out, h_out_idx, w_out_idx]
        int out_index = n * (C * D_out * H_out * W_out) + c * (D_out * H_out * W_out) +
                        d_out * (H_out * W_out) + h_out_idx * W_out + w_out_idx;
        output[out_index] = max_val;
    }
}

// Forward function that applies 3D convolution, softmax, and then the fused max pooling in one kernel call

torch::Tensor forward(
    torch::Tensor x,
    torch::Tensor conv_weight,
    torch::Tensor conv_bias
) {
    // Ensure the input tensors are contiguous
    x = x.contiguous();
    conv_weight = conv_weight.contiguous();
    conv_bias = conv_bias.contiguous();

    // Apply 3D convolution
    auto conv_output = at::conv3d(x, conv_weight, conv_bias, {1, 1, 1}, {0, 0, 0});

    // Apply softmax over the channel dimension
    auto softmax_output = at::softmax(conv_output, 1);

    // Get dimensions of the softmax output: assumed shape [N, C, D, H, W]
    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);

    // For the fused pooling, the spatial dimensions must be divisible by 4
    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());

    // Define block dimensions: using 16x16 threads for the (w, h) plane and grid.z for (n, c, d_out)
    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 = output.data_ptr<float>();

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

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Fused max pooling CUDA forward function (combining two pooling stages)");
}