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 WARP_SIZE 32
#define BLOCK_SIZE 256
#define TILE_SIZE 16

__constant__ float c_bias[MAX_OUT_FEATURES];
__constant__ float c_subtract[MAX_OUT_FEATURES];

//------------------------------------------------------------------------------
// Modular device functions for core operations
namespace device {

// Efficient dot product computation
__device__ __forceinline__ float dot_product(
    const float* __restrict__ a,
    const float* __restrict__ b,
    const int size
) {
    float sum = 0.0f;
    #pragma unroll
    for (int i = 0; i < size; ++i) {
        sum += a[i] * b[i];
    }
    return sum;
}

// Warp-level reduction
__device__ __forceinline__ float warp_reduce_sum(float val) {
    #pragma unroll
    for (int offset = WARP_SIZE/2; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

// GELU activation
__device__ __forceinline__ float gelu(float x) {
    const float kAlpha = 0.044715f;
    const float kBeta = 0.7978845608f;
    float inner = kBeta * (x + kAlpha * x * x * x);
    return x * 0.5f * (1.0f + tanhf(inner));
}

// Efficient block-level reduction
__device__ __forceinline__ float block_reduce_sum(
    float val,
    float* shared_mem
) {
    const int tid = threadIdx.x;
    const int lane_id = tid % WARP_SIZE;
    const int warp_id = tid / WARP_SIZE;
    const int warps_per_block = BLOCK_SIZE / WARP_SIZE;

    // Warp-level reduction
    val = warp_reduce_sum(val);

    // Write reduced warp results to shared memory
    if (lane_id == 0) {
        shared_mem[warp_id] = val;
    }
    __syncthreads();

    // Final reduction across warps
    if (tid < warps_per_block) {
        val = (tid < (BLOCK_SIZE / WARP_SIZE)) ? shared_mem[tid] : 0.0f;
        val = warp_reduce_sum(val);
    }
    
    return val;
}

} // namespace device

//------------------------------------------------------------------------------
// Optimized GEMM kernel with tiling and fused operations
__global__ void fused_gemm_subtract_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    float* __restrict__ out,
    const int batch_size,
    const int in_features,
    const int out_features
) {
    __shared__ float x_shared[TILE_SIZE][TILE_SIZE + 1];  // +1 for bank conflicts
    __shared__ float w_shared[TILE_SIZE][TILE_SIZE + 1];

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

    float acc = 0.0f;

    for (int t = 0; t < (in_features + TILE_SIZE - 1) / TILE_SIZE; ++t) {
        // Collaborative loading of tiles into shared memory
        if (row < batch_size && t * TILE_SIZE + ty < in_features) {
            x_shared[tx][ty] = x[row * in_features + t * TILE_SIZE + ty];
        } else {
            x_shared[tx][ty] = 0.0f;
        }
        
        if (col < out_features && t * TILE_SIZE + tx < in_features) {
            w_shared[ty][tx] = weight[col * in_features + t * TILE_SIZE + tx];
        } else {
            w_shared[ty][tx] = 0.0f;
        }
        
        __syncthreads();

        // Compute partial dot products
        acc += device::dot_product(x_shared[tx], w_shared[ty], TILE_SIZE);
        
        __syncthreads();
    }

    // Write result with fused bias and subtract
    if (row < batch_size && col < out_features) {
        out[row * out_features + col] = acc + c_bias[col] - c_subtract[col];
    }
}

//------------------------------------------------------------------------------
// Optimized average pooling kernel with fused logsumexp
__global__ void fused_avgpool_logsumexp_kernel(
    const float* __restrict__ x,
    float* __restrict__ out,
    const int batch_size,
    const int out_features
) {
    extern __shared__ float shared_mem[];
    
    const int tid = threadIdx.x;
    const int bid = blockIdx.x;
    
    float sum = 0.0f;
    for (int j = tid; j < out_features; j += blockDim.x) {
        sum += x[bid * out_features + j];
    }
    
    sum = device::block_reduce_sum(sum, shared_mem);
    
    if (tid == 0) {
        float avg = sum / static_cast<float>(out_features);
        out[bid] = logf(expf(avg));
    }
}

//------------------------------------------------------------------------------
// Optimized GELU and residual add kernel
__global__ void fused_gelu_residual_kernel(
    const float* __restrict__ scalar,
    const float* __restrict__ original_x,
    float* __restrict__ out,
    const int batch_size,
    const int in_features
) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = batch_size * in_features;
    
    if (idx < total) {
        const int i = idx / in_features;
        out[idx] = original_x[idx] + device::gelu(scalar[i]);
    }
}

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() && weight.is_cuda() && bias.is_cuda() && subtract.is_cuda(),
                "All tensors must be CUDA tensors");

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

    auto x_ = x.contiguous();
    auto w_ = weight.contiguous();
    auto b_ = bias.contiguous();
    auto s_ = subtract.contiguous();

    cudaMemcpyToSymbol(c_bias, b_.data_ptr<float>(), out_features * sizeof(float));
    cudaMemcpyToSymbol(c_subtract, s_.data_ptr<float>(), out_features * sizeof(float));

    auto original_x = x_.clone();
    auto gemm_out = torch::empty({batch_size, out_features}, x.options());
    auto pool_out = torch::empty({batch_size}, x.options());

    // Launch optimized GEMM kernel
    dim3 blockGemm(TILE_SIZE, TILE_SIZE);
    dim3 gridGemm(
        (batch_size + TILE_SIZE - 1) / TILE_SIZE,
        (out_features + TILE_SIZE - 1) / TILE_SIZE
    );
    fused_gemm_subtract_kernel<<<gridGemm, blockGemm>>>(
        x_.data_ptr<float>(),
        w_.data_ptr<float>(),
        gemm_out.data_ptr<float>(),
        batch_size,
        in_features,
        out_features
    );

    // Launch fused avgpool and logsumexp kernel
    const int threads = BLOCK_SIZE;
    const int shared_mem_size = (threads / WARP_SIZE) * sizeof(float);
    fused_avgpool_logsumexp_kernel<<<batch_size, threads, shared_mem_size>>>(
        gemm_out.data_ptr<float>(),
        pool_out.data_ptr<float>(),
        batch_size,
        out_features
    );

    // Launch fused GELU and residual add kernel
    auto out = torch::empty({batch_size, in_features}, x.options());
    const int total = batch_size * in_features;
    const int blocks = (total + threads - 1) / threads;
    fused_gelu_residual_kernel<<<blocks, threads>>>(
        pool_out.data_ptr<float>(),
        original_x.data_ptr<float>(),
        out.data_ptr<float>(),
        batch_size,
        in_features
    );

    return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward_cuda, 
          "Modular CUDA implementation with optimized device functions");
}