GPUCodeForces/S1/10/batchnorm1d_cuda.py

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

batchnorm_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>

__global__ void batchnorm_forward_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    float* __restrict__ y,
    int batch,
    int features,
    float eps
) {
    int feature = blockIdx.x;
    if (feature >= features) return;
    int tid = threadIdx.x;

    extern __shared__ float shared[];
    float* shm_sum = shared;
    float* shm_sq = shared + blockDim.x;

    float sum = 0.0f;
    float sum_sq = 0.0f;

    for (int row = tid; row < batch; row += blockDim.x) {
        float v = x[row * features + feature];
        sum += v;
        sum_sq += v * v;
    }

    shm_sum[tid] = sum;
    shm_sq[tid] = sum_sq;
    __syncthreads();

    for (int offset = blockDim.x >> 1; offset > 0; offset >>= 1) {
        if (tid < offset) {
            shm_sum[tid] += shm_sum[tid + offset];
            shm_sq[tid] += shm_sq[tid + offset];
        }
        __syncthreads();
    }

    __shared__ float s_mean;
    __shared__ float s_inv_std;

    if (tid == 0) {
        float mean = shm_sum[0] / batch;
        float var = shm_sq[0] / batch - mean * mean;
        var = var > 0.f ? var : 0.f;
        s_mean = mean;
        s_inv_std = rsqrtf(var + eps);
    }
    __syncthreads();

    float mean = s_mean;
    float inv_std = s_inv_std;
    float g = gamma[feature];
    float b = beta[feature];

    for (int row = tid; row < batch; row += blockDim.x) {
        float v = x[row * features + feature];
        float norm = (v - mean) * inv_std;
        y[row * features + feature] = norm * g + b;
    }
}

torch::Tensor batchnorm_cuda_forward(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    double eps
) {
    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(x.dtype() == torch::kFloat32, "only float32 tensors are supported");
    TORCH_CHECK(weight.dtype() == torch::kFloat32, "weight must be float32");
    TORCH_CHECK(bias.dtype() == torch::kFloat32, "bias must be float32");
    TORCH_CHECK(x.dim() == 2, "input must be 2D [batch, features]");
    TORCH_CHECK(weight.dim() == 1, "weight must be 1D");
    TORCH_CHECK(bias.dim() == 1, "bias must be 1D");
    TORCH_CHECK(x.size(1) == weight.size(0), "feature size mismatch between input and weight");
    TORCH_CHECK(weight.size(0) == bias.size(0), "weight and bias must have the same length");

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

    int batch = x_contig.size(0);
    int features = x_contig.size(1);

    auto y = torch::empty_like(x_contig);

    int threads = 256;
    if (batch < threads) {
        threads = 1;
        while (threads < batch) threads <<= 1;
        if (threads < 32) threads = 32;
    }
    size_t shared_mem = threads * 2 * sizeof(float);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    batchnorm_forward_kernel<<<features, threads, shared_mem, stream>>>(
        x_contig.data_ptr<float>(),
        weight_contig.data_ptr<float>(),
        bias_contig.data_ptr<float>(),
        y.data_ptr<float>(),
        batch,
        features,
        static_cast<float>(eps)
    );
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
"""

batchnorm_cpp_source = r"""
torch::Tensor batchnorm_cuda_forward(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    double eps
);
"""

batchnorm_cuda = load_inline(
    name="batchnorm_cuda_ext",
    cpp_sources=batchnorm_cpp_source,
    cuda_sources=batchnorm_source,
    functions=["batchnorm_cuda_forward"],
    verbose=True
)

class ModelNew(nn.Module):
    """
    Model performing matrix multiplication followed by custom CUDA BatchNorm and ReLU.
    """
    def __init__(self, mat_weight: torch.Tensor, bn_weight: torch.Tensor, bn_bias: torch.Tensor, eps: float = 1e-5):
        super().__init__()
        if mat_weight.dim() != 2:
            raise ValueError("mat_weight must be a 2D tensor [input_dim, output_dim].")
        if bn_weight.dim() != 1 or bn_bias.dim() != 1:
            raise ValueError("BatchNorm weight and bias must be 1D.")
        if bn_weight.size(0) != mat_weight.size(1):
            raise ValueError("BatchNorm parameter size must match output_dim.")
        if bn_weight.size(0) != bn_bias.size(0):
            raise ValueError("BatchNorm weight and bias must share shape.")
        self.weight = nn.Parameter(mat_weight.clone())
        self.bn_weight = nn.Parameter(bn_weight.clone())
        self.bn_bias = nn.Parameter(bn_bias.clone())
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not x.is_cuda:
            raise ValueError("Input must be a CUDA tensor.")
        if not self.weight.is_cuda:
            raise ValueError("Model weight must be on CUDA.")
        if not self.bn_weight.is_cuda or not self.bn_bias.is_cuda:
            raise ValueError("BatchNorm parameters must be on CUDA.")
        x = torch.matmul(x, self.weight)
        x = batchnorm_cuda.batchnorm_cuda_forward(x, self.bn_weight, self.bn_bias, self.eps)
        return torch.relu(x)