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,
    scaling_factor: float,
) -> torch.Tensor:
    """
    Applies linear transformation, Swish activation, and scaling.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, in_features)
        weight (torch.Tensor): Weight matrix of shape (out_features, in_features)
        bias (torch.Tensor): Bias vector of shape (out_features)
        scaling_factor (float): Factor to scale the output by

    Returns:
        torch.Tensor: Output tensor of shape (batch_size, out_features)
    """
    x = F.linear(x, weight, bias)
    x = x * torch.sigmoid(x)  # Swish activation
    x = x * scaling_factor
    return x


class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication, applies Swish activation, and scales the result.
    """

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

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


batch_size = 128
in_features = 1024
out_features = 512
scaling_factor = 2.0


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


def get_init_inputs():
    return [in_features, out_features, scaling_factor]

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

// Kernel using grid-stride loops in 2D for efficient thread mapping
__global__ void swish_scaling_kernel_grid_stride(const float* __restrict__ input,
                                                   float* output,
                                                   float scaling_factor,
                                                   int rows,
                                                   int cols) {
    // Calculate initial row and column for this thread
    int init_row = blockIdx.y * blockDim.y + threadIdx.y;
    int init_col = blockIdx.x * blockDim.x + threadIdx.x;

    // Use grid-stride loops to cover the full matrix
    for (int row = init_row; row < rows; row += gridDim.y * blockDim.y) {
        for (int col = init_col; col < cols; col += gridDim.x * blockDim.x) {
            int idx = row * cols + col;
            float x = input[idx];
            float sigmoid = 1.0f / (1.0f + expf(-x));
            output[idx] = x * sigmoid * scaling_factor;
        }
    }
}

// The forward function performing matrix multiplication and invoking the CUDA kernel
torch::Tensor forward(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    double scaling_factor) {

    // Ensure tensors are contiguous and on CUDA
    x = x.contiguous();
    weight = weight.contiguous();
    bias = bias.contiguous();

    TORCH_CHECK(x.is_cuda(), "Input tensor 'x' must be a CUDA tensor.");
    TORCH_CHECK(weight.is_cuda(), "Weight tensor must be a CUDA tensor.");
    TORCH_CHECK(bias.is_cuda(), "Bias tensor must be a CUDA tensor.");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "Input tensor 'x' must be of type torch.float32.");
    TORCH_CHECK(weight.scalar_type() == at::kFloat, "Weight tensor must be of type torch.float32.");
    TORCH_CHECK(bias.scalar_type() == at::kFloat, "Bias tensor must be of type torch.float32.");

    // Compute the linear transformation: y = bias + x @ weight.T
    auto y = at::addmm(bias, x, weight.t());
    auto output = at::empty_like(y);

    // Determine matrix dimensions
    int rows = (y.dim() == 1) ? 1 : y.size(0);
    int cols = (y.dim() == 1) ? y.size(0) : y.size(1);

    // Choose block dimensions (16x16 is a balanced choice)
    dim3 threads(16, 16);
    dim3 blocks((cols + threads.x - 1) / threads.x,
                (rows + threads.y - 1) / threads.y);

    // Launch the kernel with grid-stride loops for robust mapping
    swish_scaling_kernel_grid_stride<<<blocks, threads>>>(
        y.data_ptr<float>(),
        output.data_ptr<float>(),
        static_cast<float>(scaling_factor),
        rows,
        cols);

    // Check for any kernel launch errors
    cudaError_t err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess, "CUDA kernel failed: ", cudaGetErrorString(err));

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Matrix Multiplication with Swish Activation and Scaling using Grid-Stride Loops");
}