The AI CUDA Engineer 👷

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


def module_fn(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
    subtract: torch.Tensor,
) -> torch.Tensor:
    """
    Performs a series of operations: Gemm, Subtract, GlobalAvgPool, LogSumExp, GELU, and ResidualAdd.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, in_features)
        weight (torch.Tensor): Weight matrix for linear layer of shape (out_features, in_features)
        bias (torch.Tensor): Bias vector for linear layer of shape (out_features)
        subtract (torch.Tensor): Vector to subtract of shape (out_features)

    Returns:
        torch.Tensor: Output tensor after applying all operations
    """
    original_x = x.clone().detach()

    # Gemm
    x = F.linear(x, weight, bias)

    # Subtract
    x = x - subtract

    # GlobalAvgPool
    x = torch.mean(x, dim=1, keepdim=True)

    # LogSumExp
    x = torch.logsumexp(x, dim=1, keepdim=True)

    # GELU
    x = F.gelu(x)

    # ResidualAdd
    x = x + original_x

    return x


class Model(nn.Module):
    """
    Model that performs a series of operations: Gemm, Subtract, GlobalAvgPool, LogSumExp, GELU, and ResidualAdd.
    """

    def __init__(self, in_features, out_features):
        super(Model, self).__init__()
        gemm = nn.Linear(in_features, out_features)
        self.weight = nn.Parameter(gemm.weight)
        self.bias = nn.Parameter(gemm.bias)
        self.subtract = nn.Parameter(torch.randn(out_features) * 0.02)

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


batch_size = 128
in_features = 1024
out_features = 512


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


def get_init_inputs():
    return [in_features, out_features]

#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

#define MAX_OUT_FEATURES 4096
#define TILE_DIM 16
#define BLOCK_ROWS 16

// Constant memory for bias and subtract vectors
__constant__ float c_bias[MAX_OUT_FEATURES];
__constant__ float c_subtract[MAX_OUT_FEATURES];

//---------------------------------------------------------------------------
// Optimized GEMM kernel with coalesced memory access using tiling
// Computes: out[r, c] = dot(x[r, :], weight[c, :]) + c_bias[c] - c_subtract[c]
// x: [batch_size x in_features]
// weight: [out_features x in_features]
// out: [batch_size x out_features]
//---------------------------------------------------------------------------
__global__ void coalesced_gemm_subtract_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    float* __restrict__ out,
    int batch_size,
    int in_features,
    int out_features
) {
    __shared__ float tile_x[TILE_DIM][TILE_DIM+1];
    __shared__ float tile_w[TILE_DIM][TILE_DIM+1];

    int bx = blockIdx.x;
    int by = blockIdx.y;
    int tx = threadIdx.x;
    int ty = threadIdx.y;

    int row = by * TILE_DIM + ty;
    int col = bx * TILE_DIM + tx;

    float sum = 0.0f;

    // Loop over tiles
    for (int t = 0; t < (in_features + TILE_DIM - 1) / TILE_DIM; t++) {
        // Load tile from x
        if (row < batch_size && (t * TILE_DIM + tx) < in_features)
            tile_x[ty][tx] = x[row * in_features + t * TILE_DIM + tx];
        else
            tile_x[ty][tx] = 0.0f;
        
        // Load tile from weight
        if (col < out_features && (t * TILE_DIM + ty) < in_features)
            tile_w[ty][tx] = weight[col * in_features + t * TILE_DIM + ty];
        else
            tile_w[ty][tx] = 0.0f;

        __syncthreads();

        #pragma unroll
        for (int k = 0; k < TILE_DIM; k++) {
            sum += tile_x[ty][k] * tile_w[k][tx];
        }

        __syncthreads();
    }

    if (row < batch_size && col < out_features) {
        out[row * out_features + col] = sum + c_bias[col] - c_subtract[col];
    }
}

//---------------------------------------------------------------------------
// Fused kernel: for each batch row, perform average pooling (reduction) over the GEMM output,
// apply the GELU activation on the computed average, and then add the result to each element
// of the corresponding original input row (residual addition).
// gemm_out: [batch_size x out_features]
// original_x: [batch_size x in_features]
// out: [batch_size x in_features]
//---------------------------------------------------------------------------
__global__ void fused_pool_gelu_residual_kernel(
    const float* __restrict__ gemm_out,
    const float* __restrict__ original_x,
    float* __restrict__ out,
    int batch_size,
    int out_features,
    int in_features
) {
    // One block per batch row
    int row = blockIdx.x;
    
    // Shared memory for reduction
    extern __shared__ float sdata[];
    float local_sum = 0.0f;
    
    // Each thread reduces part of the GEMM output row
    for (int i = threadIdx.x; i < out_features; i += blockDim.x) {
        local_sum += gemm_out[row * out_features + i];
    }
    sdata[threadIdx.x] = local_sum;
    __syncthreads();

    // Reduction in shared memory
    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s)
            sdata[threadIdx.x] += sdata[threadIdx.x + s];
        __syncthreads();
    }

    // Thread 0 computes the average and applies GELU activation
    float avg = sdata[0] / static_cast<float>(out_features);
    float gelu = avg * 0.5f * (1.0f + tanhf(0.7978845608f * (avg + 0.044715f * avg * avg * avg)));
    __syncthreads();

    // Residual addition: add the activated scalar to each element of original_x row
    for (int j = threadIdx.x; j < in_features; j += blockDim.x) {
        int idx = row * in_features + j;
        out[idx] = original_x[idx] + gelu;
    }
}

//---------------------------------------------------------------------------
// Forward function that launches the fused pipeline.
// Sequence of operations:
// 1. GEMM with bias and subtract using coalesced memory accesses
// 2. Fused kernel that performs average pooling over GEMM output, applies GELU activation,
//    and performs residual addition with the original input.
//---------------------------------------------------------------------------

torch::Tensor forward_cuda(
    const torch::Tensor& x,
    const torch::Tensor& weight,
    const torch::Tensor& bias,
    const torch::Tensor& subtract
) {
    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");
    TORCH_CHECK(subtract.is_cuda(), "subtract must be a CUDA tensor");
    TORCH_CHECK(x.dim() == 2, "x must be 2D (batch_size x in_features)");
    TORCH_CHECK(weight.dim() == 2, "weight must be 2D (out_features x in_features)");
    TORCH_CHECK(bias.dim() == 1, "bias must be 1D (out_features)");
    TORCH_CHECK(subtract.dim() == 1, "subtract must be 1D (out_features)");

    int batch_size = x.size(0);
    int in_features = x.size(1);
    int out_features = weight.size(0);

    TORCH_CHECK(weight.size(1) == in_features, "Mismatch between weight and x dimensions");
    TORCH_CHECK(bias.size(0) == out_features, "bias dimension must match weight output features");
    TORCH_CHECK(subtract.size(0) == out_features, "subtract dimension must match weight output features");
    TORCH_CHECK(out_features <= MAX_OUT_FEATURES, "out_features exceeds maximum allowed for constant memory");

    auto x_contig = x.contiguous();
    auto weight_contig = weight.contiguous();
    auto bias_contig = bias.contiguous();
    auto subtract_contig = subtract.contiguous();

    // Copy bias and subtract vectors to constant memory
    cudaMemcpyToSymbol(c_bias, bias_contig.data_ptr<float>(), out_features * sizeof(float));
    cudaMemcpyToSymbol(c_subtract, subtract_contig.data_ptr<float>(), out_features * sizeof(float));

    // Clone original input for residual addition
    auto original_x = x_contig.clone();

    // Allocate intermediate tensor for GEMM output and final output tensor
    auto gemm_out = torch::empty({batch_size, out_features}, x.options());
    auto out_tensor = torch::empty({batch_size, in_features}, x.options());

    // Launch the optimized GEMM kernel
    dim3 threadsGemm(TILE_DIM, BLOCK_ROWS);
    dim3 blocksGemm((out_features + TILE_DIM - 1) / TILE_DIM, (batch_size + TILE_DIM - 1) / TILE_DIM);
    coalesced_gemm_subtract_kernel<<<blocksGemm, threadsGemm>>>(
        x_contig.data_ptr<float>(),
        weight_contig.data_ptr<float>(),
        gemm_out.data_ptr<float>(),
        batch_size,
        in_features,
        out_features
    );

    // Launch the fused pooling, GELU, and residual addition kernel
    // One block per batch row; using 256 threads per block and dynamic shared memory for reduction
    int threadsFused = 256;
    fused_pool_gelu_residual_kernel<<<batch_size, threadsFused, threadsFused * sizeof(float)>>>(
        gemm_out.data_ptr<float>(),
        original_x.data_ptr<float>(),
        out_tensor.data_ptr<float>(),
        batch_size,
        out_features,
        in_features
    );

    return out_tensor;
}

// PyBind11 interface
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward_cuda, "Fused GEMM, pooling, GELU, and residual add CUDA kernel");
}