The AI CUDA Engineer 👷

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


def module_fn(
    x: torch.Tensor,
    pool_kernel_size: int,
    scale_factor: float,
    weight: torch.Tensor,
    bias: torch.Tensor,
) -> torch.Tensor:
    """
    Implements Matmul_AvgPool_GELU_Scale_Max pattern using functional operations.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, in_features)
        pool_kernel_size (int): Kernel size for average pooling
        scale_factor (float): Scale factor to multiply features by
        weight (torch.Tensor): Weight matrix for linear layer
        bias (torch.Tensor): Bias vector for linear layer

    Returns:
        torch.Tensor: Output tensor of shape (batch_size,)
    """
    x = F.linear(x, weight, bias)
    x = F.avg_pool1d(x.unsqueeze(1), kernel_size=pool_kernel_size).squeeze(1)
    x = F.gelu(x)
    x = x * scale_factor
    x = torch.max(x, dim=1).values
    return x


class Model(nn.Module):
    """
    A model implementing the pattern "Matmul_AvgPool_GELU_Scale_Max".
    """

    def __init__(self, in_features, out_features, pool_kernel_size, scale_factor):
        super(Model, self).__init__()
        gemm = nn.Linear(in_features, out_features)
        self.weight = gemm.weight
        self.bias = gemm.bias

    def forward(self, x, pool_kernel_size, scale_factor, fn=module_fn):
        return fn(x, pool_kernel_size, scale_factor, self.weight, self.bias)


batch_size = 128
in_features = 512
out_features = 256
pool_kernel_size = 4
scale_factor = 2.0


def get_inputs():
    return [torch.randn(batch_size, in_features), pool_kernel_size, scale_factor]


def get_init_inputs():
    return [in_features, out_features, pool_kernel_size, scale_factor]

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

#ifndef TILE_SIZE
#define TILE_SIZE 16
#endif

__global__ void FusedMatMulBiasKernel(const float* __restrict__ A,
                                      const float* __restrict__ B,
                                      const float* __restrict__ bias,
                                      float* __restrict__ C,
                                      int M, int N, int K) {
    __shared__ float Asub[TILE_SIZE][TILE_SIZE];
    __shared__ float Bsub[TILE_SIZE][TILE_SIZE];
    __shared__ float shared_bias[TILE_SIZE];  // Cache bias values

    int row = blockIdx.y * TILE_SIZE + threadIdx.y;
    int col = blockIdx.x * TILE_SIZE + threadIdx.x;

    // Preload bias into shared memory
    if (threadIdx.y == 0 && (blockIdx.x * TILE_SIZE + threadIdx.x) < N) {
        shared_bias[threadIdx.x] = bias[blockIdx.x * TILE_SIZE + threadIdx.x];
    }
    __syncthreads();

    float sum = 0.0f;

    for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; t++) {
        int tiled_k = t * TILE_SIZE;
        
        // Coalesced loading of A and B tiles
        if (row < M && (tiled_k + threadIdx.x) < K)
            Asub[threadIdx.y][threadIdx.x] = A[row * K + tiled_k + threadIdx.x];
        else
            Asub[threadIdx.y][threadIdx.x] = 0.0f;

        if (col < N && (tiled_k + threadIdx.y) < K)
            Bsub[threadIdx.y][threadIdx.x] = B[col * K + tiled_k + threadIdx.y];
        else
            Bsub[threadIdx.y][threadIdx.x] = 0.0f;

        __syncthreads();

        #pragma unroll
        for (int i = 0; i < TILE_SIZE; i++) {
            sum += Asub[threadIdx.y][i] * Bsub[i][threadIdx.x];
        }
        __syncthreads();
    }

    if (row < M && col < N) {
        C[row * N + col] = sum + shared_bias[threadIdx.x];
    }
}

__global__ void FusedPoolActMaxKernel(const float* __restrict__ linear_output,
                                      float* __restrict__ output,
                                      int M, int N,
                                      int pool_kernel_size,
                                      int output_length,
                                      float scale_factor) {
    extern __shared__ float shared_mem[];
    float* pool_buffer = shared_mem;  // Used for storing pooling results
    float* reduction_buffer = &shared_mem[blockDim.x];  // Used for max reduction

    int row = blockIdx.x;
    int tid = threadIdx.x;
    
    float local_max = -FLT_MAX;
    
    // Process multiple pooling windows per thread
    for (int bin = tid; bin < output_length; bin += blockDim.x) {
        int start = bin * pool_kernel_size;
        float sum = 0.0f;
        int count = 0;

        // Load pooling window elements into shared memory
        #pragma unroll
        for (int j = 0; j < pool_kernel_size; j++) {
            int col = start + j;
            if (col < N) {
                pool_buffer[tid] = linear_output[row * N + col];
                sum += pool_buffer[tid];
                count++;
            }
        }
        
        // Compute average and GELU in registers
        float avg = (count > 0) ? sum / count : 0.0f;
        float gelu = 0.5f * avg * (1.0f + erff(avg * 0.70710678f));
        gelu *= scale_factor;
        
        local_max = fmaxf(local_max, gelu);
    }

    // Store local max in reduction buffer
    reduction_buffer[tid] = local_max;
    __syncthreads();

    // Parallel reduction in shared memory
    for (unsigned int s = blockDim.x/2; s > 32; s >>= 1) {
        if (tid < s) {
            reduction_buffer[tid] = fmaxf(reduction_buffer[tid], reduction_buffer[tid + s]);
        }
        __syncthreads();
    }

    // Warp-level reduction (no sync needed within a warp)
    if (tid < 32) {
        volatile float* smem = reduction_buffer;
        if (blockDim.x >= 64) smem[tid] = fmaxf(smem[tid], smem[tid + 32]);
        if (blockDim.x >= 32) smem[tid] = fmaxf(smem[tid], smem[tid + 16]);
        if (blockDim.x >= 16) smem[tid] = fmaxf(smem[tid], smem[tid + 8]);
        if (blockDim.x >= 8) smem[tid] = fmaxf(smem[tid], smem[tid + 4]);
        if (blockDim.x >= 4) smem[tid] = fmaxf(smem[tid], smem[tid + 2]);
        if (blockDim.x >= 2) smem[tid] = fmaxf(smem[tid], smem[tid + 1]);
    }

    if (tid == 0) {
        output[row] = reduction_buffer[0];
    }
}

torch::Tensor forward(
    torch::Tensor x,
    int pool_kernel_size,
    float scale_factor,
    torch::Tensor weight,
    torch::Tensor bias) {
    
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
    TORCH_CHECK(bias.is_cuda(), "bias must be a CUDA tensor");

    x = x.contiguous();
    weight = weight.contiguous();
    bias = bias.contiguous();

    int M = x.size(0);
    int K = x.size(1);
    int N = weight.size(0);

    auto options = torch::TensorOptions().dtype(x.dtype()).device(x.device());
    auto linear_output = torch::empty({M, N}, options);

    dim3 blockDim(TILE_SIZE, TILE_SIZE);
    dim3 gridDim((N + TILE_SIZE - 1) / TILE_SIZE, (M + TILE_SIZE - 1) / TILE_SIZE);
    
    FusedMatMulBiasKernel<<<gridDim, blockDim>>>(
        x.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        linear_output.data_ptr<float>(),
        M, N, K);

    int output_length = (N + pool_kernel_size - 1) / pool_kernel_size;
    auto output = torch::empty({M}, options);

    int threads = 256;
    // Shared memory size = space for pooling buffer + reduction buffer
    size_t sharedMemSize = threads * sizeof(float) * 2;
    
    FusedPoolActMaxKernel<<<M, threads, sharedMemSize>>>(
        linear_output.data_ptr<float>(),
        output.data_ptr<float>(),
        M, N, pool_kernel_size, output_length, scale_factor);

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Shared Memory Optimized Forward");
}