The AI CUDA Engineer 👷

import torch
import torch.nn as nn
import torch.nn.functional as F


def module_fn(
    x: torch.Tensor,
    stride: int,
    padding: int,
    output_padding: int,
    pool_kernel_size: int,
    pool_stride: int,
    pool_padding: int,
    conv_transpose: torch.Tensor,
    conv_transpose_bias: torch.Tensor,
    subtract: torch.Tensor,
) -> torch.Tensor:
    """
    Applies sequence of operations:
        - ConvTranspose3d
        - MaxPool3d
        - Softmax
        - Subtract
        - Swish
        - Max

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, in_channels, depth, height, width)
        stride (int): Stride for conv transpose
        padding (int): Padding for conv transpose
        output_padding (int): Output padding for conv transpose
        pool_kernel_size (int): Kernel size for max pooling
        pool_stride (int): Stride for max pooling
        pool_padding (int): Padding for max pooling
        conv_transpose (torch.Tensor): Weight tensor for transposed convolution
        conv_transpose_bias (torch.Tensor): Bias tensor for transposed convolution
        subtract (torch.Tensor): Subtraction parameter tensor
    """
    x = F.conv_transpose3d(
        x,
        conv_transpose,
        bias=conv_transpose_bias,
        stride=stride,
        padding=padding,
        output_padding=output_padding,
    )
    x = F.max_pool3d(
        x, kernel_size=pool_kernel_size, stride=pool_stride, padding=pool_padding
    )
    x = F.softmax(x, dim=1)
    x = x - subtract.view(1, -1, 1, 1, 1)
    x = torch.sigmoid(x) * x  # Swish
    x = torch.max(x, dim=1)[0]
    return x


class Model(nn.Module):
    """
    A model that performs a sequence of operations:
        - ConvTranspose3d
        - MaxPool3d
        - Softmax
        - Subtract
        - Swish
        - Max
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride,
        padding,
        output_padding,
        pool_kernel_size,
        pool_stride,
        pool_padding,
    ):
        super(Model, self).__init__()
        conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size)
        self.conv_transpose_parameter = conv_transpose.weight
        self.conv_transpose_bias = conv_transpose.bias
        self.subtract_parameter = nn.Parameter(torch.randn(out_channels) * 0.02)

    def forward(
        self,
        x,
        stride,
        padding,
        output_padding,
        pool_kernel_size,
        pool_stride,
        pool_padding,
        fn=module_fn,
    ):
        return fn(
            x,
            stride,
            padding,
            output_padding,
            pool_kernel_size,
            pool_stride,
            pool_padding,
            self.conv_transpose_parameter,
            self.conv_transpose_bias,
            self.subtract_parameter,
        )


batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1
output_padding = 1
pool_kernel_size = 2
pool_stride = 2
pool_padding = 0


def get_inputs():
    return [
        torch.randn(batch_size, in_channels, depth, height, width),
        stride,
        padding,
        output_padding,
        pool_kernel_size,
        pool_stride,
        pool_padding,
    ]


def get_init_inputs():
    return [
        in_channels,
        out_channels,
        kernel_size,
        stride,
        padding,
        output_padding,
        pool_kernel_size,
        pool_stride,
        pool_padding,
    ]

#include <torch/extension.h>
#include <pybind11/pybind11.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>
#include <float.h>

namespace py = pybind11;

// Fused kernel that performs softmax (across the channel dimension), subtraction, swish activation, and final channel-wise max reduction,
// using warp shuffle intrinsics to minimize warp divergence and ensure uniform control flow.

// Each block processes one spatial location (n, d, h, w) from an input tensor of shape [N, C, D, H, W].
// Dynamic shared memory is used to store the per-channel exponentials (for softmax) and to facilitate warp-level reductions.

__global__ void fused_kernel(const float* __restrict__ input,
                              float* output,
                              const float* __restrict__ subtract,
                              int N, int C, int D, int H, int W) {
    // Compute spatial indices from blockIdx.x
    int pos = blockIdx.x;  // pos in [0, N*D*H*W)
    int spatial = D * H * W;
    int n = pos / spatial;
    int rem = pos % spatial;
    int d = rem / (H * W);
    int rem2 = rem % (H * W);
    int h = rem2 / W;
    int w = rem2 % W;

    // Allocate dynamic shared memory. Layout:
    // [0, C) -> to hold per-channel computed exp(x) values
    // [C, C + blockDim.x) -> workspace for warp-level reductions
    extern __shared__ float sdata[];
    float* sh_exp = sdata;               // size: C floats
    float* warp_arr = sdata + C;           // size: blockDim.x floats

    // Phase 1: Compute exponentials for softmax and accumulate partial sum
    float partial_sum = 0.0f;
    for (int c = threadIdx.x; c < C; c += blockDim.x) {
        int index = (((n * C + c) * D + d) * H + h) * W + w;
        float val = __ldg(&input[index]);
        float exp_val = expf(val);
        sh_exp[c] = exp_val;
        partial_sum += exp_val;
    }

    // Use warp shuffle to reduce partial_sum across threads within each warp
    unsigned int mask = 0xffffffff;
    float sum = partial_sum;
    for (int offset = warpSize/2; offset > 0; offset /= 2) {
        sum += __shfl_down_sync(mask, sum, offset);
    }
    
    int lane = threadIdx.x & (warpSize - 1);
    int warp_id = threadIdx.x / warpSize;
    if (lane == 0) {
        warp_arr[warp_id] = sum;
    }
    __syncthreads();

    // Reduce the per-warp sums to compute the total sum
    int num_warps = (blockDim.x + warpSize - 1) / warpSize;
    float total_sum = 0.0f;
    if (threadIdx.x < num_warps) {
        total_sum = warp_arr[threadIdx.x];
    }
    if (threadIdx.x < warpSize) {
        for (int offset = warpSize/2; offset > 0; offset /= 2) {
            total_sum += __shfl_down_sync(mask, total_sum, offset);
        }
    }
    if (threadIdx.x == 0) {
        warp_arr[0] = total_sum;
    }
    __syncthreads();
    total_sum = warp_arr[0];

    // Phase 2: For each channel, compute fused operation:
    // softmax: soft = exp_val / total_sum, then subtract subtract_tensor value, and apply swish: y * sigmoid(y).
    // Also, compute a local maximum across channels.
    float local_max = -FLT_MAX;
    for (int c = threadIdx.x; c < C; c += blockDim.x) {
        float soft = sh_exp[c] / total_sum;
        float y = soft - __ldg(&subtract[c]);
        float fused_val = y * (1.0f / (1.0f + expf(-y))); // swish activation
        local_max = fmaxf(local_max, fused_val);
    }

    // Warp-level reduction to compute maximum
    float max_val = local_max;
    for (int offset = warpSize/2; offset > 0; offset /= 2) {
        max_val = fmaxf(max_val, __shfl_down_sync(mask, max_val, offset));
    }
    if (lane == 0) {
        warp_arr[warp_id] = max_val;
    }
    __syncthreads();
    float block_max = -FLT_MAX;
    if (threadIdx.x < num_warps) {
        block_max = warp_arr[threadIdx.x];
    }
    if (threadIdx.x < warpSize) {
        for (int offset = warpSize/2; offset > 0; offset /= 2) {
            block_max = fmaxf(block_max, __shfl_down_sync(mask, block_max, offset));
        }
    }
    if (threadIdx.x == 0) {
        output[pos] = block_max;
    }
}

// The forward function first computes the transposed convolution and max pooling using ATen's functions,
// then launches the fused CUDA kernel to perform softmax along the channel dimension, subtract, swish activation, and finally channel-wise max reduction.

torch::Tensor forward(
    torch::Tensor x,
    int64_t stride,
    int64_t padding,
    int64_t output_padding,
    int64_t pool_kernel_size,
    int64_t pool_stride,
    int64_t pool_padding,
    torch::Tensor conv_transpose_weight,
    torch::Tensor conv_transpose_bias,
    torch::Tensor subtract_tensor
) {
    // Transposed convolution
    auto conv_out = at::conv_transpose3d(
        x,
        conv_transpose_weight,
        conv_transpose_bias,
        {stride, stride, stride},
        {padding, padding, padding},
        {output_padding, output_padding, output_padding},
        1,
        {1, 1, 1}
    );

    // MaxPool
    auto pool_out = at::max_pool3d(
        conv_out,
        {pool_kernel_size, pool_kernel_size, pool_kernel_size},
        {pool_stride, pool_stride, pool_stride},
        {pool_padding, pool_padding, pool_padding}
    );

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

    auto output = torch::empty({N, D, H, W}, pool_out.options());

    // Each block processes one spatial location (n, d, h, w)
    int spatial = D * H * W;
    int grid_size = N * spatial;
    int block_size = 128;  // chosen to balance occupancy
    // Dynamic shared memory: C floats for storing exp results + block_size floats for warp reduction workspace
    size_t shm_size = (C + block_size) * sizeof(float);

    const float* pool_ptr = pool_out.data_ptr<float>();
    float* output_ptr = output.data_ptr<float>();
    const float* subtract_ptr = subtract_tensor.data_ptr<float>();

    fused_kernel<<<grid_size, block_size, shm_size>>>(pool_ptr, output_ptr, subtract_ptr, N, C, D, H, W);

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "CUDA forward pass for fused operations with minimized warp divergence");
}