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,
    conv_transpose: torch.Tensor,
    conv_transpose_bias: torch.Tensor,
    multiplier: float,
) -> torch.Tensor:
    """
    Applies transposed convolution, scalar multiplication, and multiple global average pooling operations.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width)
        stride (int): Stride of the transposed convolution
        padding (int): Padding of the transposed convolution
        output_padding (int): Additional size added to output shape
        conv_transpose (torch.Tensor): Transposed convolution weight tensor
        conv_transpose_bias (torch.Tensor): Bias tensor for transposed convolution
        multiplier (float): Scalar multiplier value

    Returns:
        torch.Tensor: Scalar output after applying operations
    """
    x = F.conv_transpose2d(
        x,
        conv_transpose,
        bias=conv_transpose_bias,
        stride=stride,
        padding=padding,
        output_padding=output_padding,
    )
    x = x * multiplier
    x = torch.mean(x, dim=[2, 3], keepdim=True)
    x = torch.mean(x, dim=[2, 3], keepdim=True)
    x = torch.mean(x)
    return x


class Model(nn.Module):
    """
    Model that performs a transposed convolution, multiplies by a scalar, applies global average pooling,
    another global average pooling, and then calculates the mean.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        stride,
        padding,
        output_padding,
        multiplier,
    ):
        super(Model, self).__init__()
        conv = nn.ConvTranspose2d(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            output_padding=output_padding,
        )
        self.conv_transpose_parameter = nn.Parameter(conv.weight)
        self.conv_transpose_bias = nn.Parameter(
            conv.bias
            + torch.randn(
                conv.bias.shape, device=conv.bias.device, dtype=conv.bias.dtype
            )
            * 0.02
        )
        self.multiplier = multiplier

    def forward(self, x, stride, padding, output_padding, fn=module_fn):
        return fn(
            x,
            stride,
            padding,
            output_padding,
            self.conv_transpose_parameter,
            self.conv_transpose_bias,
            self.multiplier,
        )


batch_size = 128
in_channels = 3
out_channels = 16
height, width = 32, 32
kernel_size = 3
stride = 2
padding = 1
output_padding = 1
multiplier = 0.5


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


def get_init_inputs():
    return [
        in_channels,
        out_channels,
        kernel_size,
        stride,
        padding,
        output_padding,
        multiplier,
    ]

#include <torch/extension.h>
#include <ATen/ATen.h>
#include <cuda_runtime.h>

// Define tile dimensions for better memory access patterns
#define TILE_DIM 32
#define BLOCK_ROWS 8

template<int BLOCK_SIZE>
__global__ void tiled_reduction_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    const int N,
    const int C,
    const int H,
    const int W,
    const float multiplier
) {
    __shared__ float tile[TILE_DIM][TILE_DIM + 1];  // +1 to avoid bank conflicts
    __shared__ float row_sums[TILE_DIM];
    
    const int tid = threadIdx.x;
    const int bid = blockIdx.x;
    const int batch_idx = bid / C;
    const int channel_idx = bid % C;
    const int spatial_size = H * W;
    
    // Calculate base input pointer for this block
    const float* block_input = input + (batch_idx * C * spatial_size) + (channel_idx * spatial_size);
    
    float thread_sum = 0.0f;
    
    // Process input in tiles
    #pragma unroll 2
    for (int by = 0; by < H; by += TILE_DIM) {
        #pragma unroll 2
        for (int bx = 0; bx < W; bx += TILE_DIM) {
            // Load tile into shared memory
            #pragma unroll
            for (int y = 0; y < TILE_DIM; y += BLOCK_ROWS) {
                int gy = by + y + (tid / TILE_DIM);
                int gx = bx + (tid % TILE_DIM);
                
                if (gy < H && gx < W) {
                    tile[y + (tid / TILE_DIM)][tid % TILE_DIM] = 
                        block_input[gy * W + gx] * multiplier;
                } else {
                    tile[y + (tid / TILE_DIM)][tid % TILE_DIM] = 0.0f;
                }
            }
            __syncthreads();
            
            // Reduce tile
            if (tid < TILE_DIM) {
                float sum = 0.0f;
                #pragma unroll
                for (int i = 0; i < TILE_DIM; i++) {
                    sum += tile[tid][i];
                }
                row_sums[tid] = sum;
            }
            __syncthreads();
            
            // Final reduction within the tile
            if (tid < (TILE_DIM/2)) {
                row_sums[tid] += row_sums[tid + (TILE_DIM/2)];
            }
            __syncthreads();
            
            if (tid < (TILE_DIM/4)) {
                row_sums[tid] += row_sums[tid + (TILE_DIM/4)];
            }
            __syncthreads();
            
            if (tid < (TILE_DIM/8)) {
                row_sums[tid] += row_sums[tid + (TILE_DIM/8)];
            }
            __syncthreads();
            
            if (tid < (TILE_DIM/16)) {
                row_sums[tid] += row_sums[tid + (TILE_DIM/16)];
            }
            __syncthreads();
            
            if (tid == 0) {
                thread_sum += (row_sums[0] + row_sums[1]);
            }
            __syncthreads();
        }
    }
    
    // Write final result
    if (tid == 0) {
        output[bid] = thread_sum / spatial_size;
    }
}

at::Tensor module_fn(
    at::Tensor x,
    int64_t stride,
    int64_t padding,
    int64_t output_padding,
    at::Tensor conv_transpose,
    at::Tensor conv_transpose_bias,
    double multiplier
) {
    // Apply transposed convolution
    at::Tensor y = at::conv_transpose2d(
        x,
        conv_transpose,
        conv_transpose_bias,
        {stride, stride},
        {padding, padding},
        {output_padding, output_padding},
        1,
        {1, 1}
    );
    
    // Prepare output tensor
    auto options = torch::TensorOptions().device(y.device()).dtype(y.dtype());
    auto dims = y.sizes();
    at::Tensor output = torch::zeros({dims[0], dims[1]}, options);
    
    // Launch kernel with optimized configuration
    constexpr int BLOCK_SIZE = TILE_DIM * BLOCK_ROWS;  // Optimized for tile-based processing
    const int blocks = dims[0] * dims[1];
    const int shared_mem_size = (TILE_DIM * (TILE_DIM + 1) + TILE_DIM) * sizeof(float);
    
    tiled_reduction_kernel<BLOCK_SIZE><<<blocks, BLOCK_SIZE, shared_mem_size>>>(
        y.data_ptr<float>(),
        output.data_ptr<float>(),
        dims[0], dims[1], dims[2], dims[3],
        static_cast<float>(multiplier)
    );
    
    // Compute final mean
    return output.mean();
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &module_fn, "Module function");
}