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 coalesced_maxpool_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    const int N, const int C, const int D, const int H, const int W,
    const int outD, const int outH, const int outW,
    const int items_per_block
) {
    // Reorganized memory layout for coalesced access
    const int stride_n = C * D * H * W;
    const int stride_c = D * H * W;
    const int stride_d = H * W;
    const int stride_h = W;
    
    // New thread mapping: prioritize width dimension for coalescing
    const int tid = threadIdx.x;
    const int local_w = tid % 4;          // Fastest dimension [0-3]
    const int local_h = (tid / 4) % 4;    // Medium dimension [0-3]
    const int local_d = tid / 16;         // Slowest dimension [0-3]

    const int total_elements = N * C * outD * outH * outW;
    const int block_start = blockIdx.x * items_per_block;

    // Process multiple elements per block
    #pragma unroll 4
    for (int item_idx = 0; item_idx < items_per_block; item_idx++) {
        const int global_idx = block_start + item_idx;
        if (global_idx >= total_elements) break;

        // Convert linear index to 5D output coordinates
        const int n = global_idx / (C * outD * outH * outW);
        int rem = global_idx % (C * outD * outH * outW);
        const int c = rem / (outD * outH * outW);
        rem %= outD * outH * outW;
        const int out_d = rem / (outH * outW);
        rem %= outH * outW;
        const int out_h = rem / outW;
        const int out_w = rem % outW;

        // Input window starting positions
        const int d_start = out_d * 4;
        const int h_start = out_h * 4;
        const int w_start = out_w * 4;

        // Global input coordinates with coalesced access pattern
        const int d = d_start + local_d;
        const int h = h_start + local_h;
        const int w = w_start + local_w;

        // Coalesced memory access with boundary checks
        float val = -FLT_MAX;
        if (d < D && h < H && w < W) {
            const int input_idx = n * stride_n + c * stride_c + d * stride_d + h * stride_h + w;
            val = __ldg(&input[input_idx]);
        }

        // Warp-level reduction with optimized shuffle pattern
        #pragma unroll
        for (int offset = 16; offset > 0; offset /= 2) {
            val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
        }

        // Efficient shared memory reduction
        __shared__ float smem[2];
        if (tid % 32 == 0) smem[tid/32] = val;
        __syncthreads();

        // Final reduction in first warp
        if (tid < 32) {
            float tmp = tid < 2 ? smem[tid] : -FLT_MAX;
            #pragma unroll
            for (int offset = 1; offset > 0; offset >>= 1) {
                tmp = fmaxf(tmp, __shfl_down_sync(0xffffffff, tmp, offset));
            }
            if (tid == 0) {
                output[global_idx] = tmp;
            }
        }
        __syncthreads();
    }
}

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

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

    const int outD = D / 4;
    const int outH = H / 4;
    const int outW = W / 4;

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

    const int items_per_block = 8;
    const int total_elements = N * C * outD * outH * outW;
    const int num_blocks = (total_elements + items_per_block - 1) / items_per_block;
    
    coalesced_maxpool_kernel<<<num_blocks, 64>>>(
        softmax_output.data_ptr<float>(),
        output.data_ptr<float>(),
        N, C, D, H, W,
        outD, outH, outW,
        items_per_block
    );

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Coalesced memory access CUDA forward");
}