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>

// Fused kernel to perform two sequential max pooling operations as one 4x4x4 pooling
// It fuses the two max pooling steps (each with kernel size 2 and stride 2) into one kernel
// that performs a 4x4x4 reduction using shared memory and warp-level primitives.

__global__ void fused_maxpool_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    int N, int C, int D, int H, int W,
    int outD, int outH, int outW
) {
    // Grid mapping:
    //   blockIdx.x -> output width index (outW)
    //   blockIdx.y -> output height index (outH)
    //   blockIdx.z -> combined index for (N, C, outD): index = n * (C*outD) + c * outD + d_out
    
    int out_w = blockIdx.x;
    int out_h = blockIdx.y;
    int index = blockIdx.z;
    int n = index / (C * outD);
    int rem = index % (C * outD);
    int c = rem / outD;
    int out_d = rem % outD;

    // Compute starting indices in the input tensor for a 4x4x4 pooling window
    int d_start = out_d * 4;
    int h_start = out_h * 4;
    int w_start = out_w * 4;

    // Each block uses 64 threads to reduce a 4x4x4 window (64 elements)
    int tid = threadIdx.x;

    // Map thread id [0,63] to 3D coordinates in 4x4x4 window
    int local_d = tid / 16;        // [0, 3]
    int local_h = (tid % 16) / 4;    // [0, 3]
    int local_w = tid % 4;           // [0, 3]

    int d = d_start + local_d;
    int h = h_start + local_h;
    int w = w_start + local_w;

    float val = -FLT_MAX;
    if (d < D && h < H && w < W) {
        int input_idx = n * (C * D * H * W) + c * (D * H * W) + d * (H * W) + h * W + w;
        val = input[input_idx];
    }

    // Perform warp-level reduction using __shfl_down_sync
    unsigned int mask = 0xffffffff;
    #pragma unroll
    for (int offset = 16; offset > 0; offset /= 2) {
        float other = __shfl_down_sync(mask, val, offset);
        val = fmaxf(val, other);
    }

    // Each warp's leader writes its result to shared memory
    __shared__ float shared[2];
    int warp_id = tid / 32;  // two warps per block
    if ((tid & 31) == 0) {
        shared[warp_id] = val;
    }
    __syncthreads();

    // Final reduction across warps in the first warp
    if (tid < 32) {
        if (tid < 2) {
            float v = shared[tid];
            #pragma unroll
            for (int offset = 1; offset < 2; offset *= 2) {
                float other = __shfl_down_sync(mask, v, offset);
                v = fmaxf(v, other);
            }
            if (tid == 0) {
                int out_idx = n * (C * outD * outH * outW) + c * (outD * outH * outW) + out_d * (outH * outW) + out_h * outW + out_w;
                output[out_idx] = v;
            }
        }
    }
}

// Forward function implementing 6_Conv3d_Softmax_MaxPool_MaxPool with fused double max pooling
// This function applies 3D convolution, softmax activation over channels, and then fuses two
// consecutive max pooling operations into one optimized CUDA kernel.

torch::Tensor forward(
    torch::Tensor x,
    torch::Tensor conv_weight,
    torch::Tensor conv_bias
) {
    // Ensure 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, /*dim=*/1);

    // Get dimensions (assuming softmax_output is [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);

    // Compute output dimensions after two pooling operations fused as a 4x4x4 pooling
    int outD = D / 4;
    int outH = H / 4;
    int outW = W / 4;

    auto options = softmax_output.options();
    auto output = torch::empty({N, C, outD, outH, outW}, options);

    // Launch the fused max pooling kernel
    // Grid: (outW, outH, N * C * outD) so that each block computes one output element
    dim3 grid(outW, outH, N * C * outD);
    int threads = 64;  // one thread per element in the 4x4x4 window

    fused_maxpool_kernel<<<grid, threads>>>(
        softmax_output.data_ptr<float>(),
        output.data_ptr<float>(),
        N, C, D, H, W,
        outD, outH, outW
    );
    cudaDeviceSynchronize();

    return output;
}


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