GPUCodeForces/S1/2/groupnorm_cuda.py

# groupnorm_cuda.py
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

from groupnorm_torch import NUM_GROUPS, C, EPS, W


assert W % 4 == 0, "Width (W) must be a multiple of 4 for float4 vectorization"

class ModelNew(nn.Module):

    def __init__(self, num_groups, num_channels, eps):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(num_channels))
        self.bias = nn.Parameter(torch.zeros(num_channels))
        self.num_groups = num_groups
        self.eps = eps
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor groupnorm_forward_cuda(
            torch::Tensor input, torch::Tensor weight, torch::Tensor bias,
            int num_groups, float eps);
        """
        cuda_source = """
        #include <cuda_runtime.h>
        #include <float.h>

        #define BLOCK_SIZE 512
        #define WARP_SIZE 32


        __device__ __forceinline__ float warp_reduce_sum(float val) {
            for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2)
                val += __shfl_down_sync(0xffffffff, val, offset);
            return val;
        }


        __global__ void groupnorm_fused_vectorized_kernel(
            const float* __restrict__ input, const float* __restrict__ weight,
            const float* __restrict__ bias, float* __restrict__ output,
            int N, int C, int H, int W, int num_groups, float eps
        ) {
            __shared__ float s_warp_sums[BLOCK_SIZE / WARP_SIZE];
            __shared__ float s_warp_sum_sqs[BLOCK_SIZE / WARP_SIZE];

            int group_idx = blockIdx.x;
            if (group_idx >= N * num_groups) return;

            int sample_idx = group_idx / num_groups;
            int group_in_sample = group_idx % num_groups;

            int C_per_group = C / num_groups;
            int group_size = C_per_group * H * W;
            

            float thread_sum = 0.0f;
            float thread_sum_sq = 0.0f;
            
            for (int i = threadIdx.x; i < group_size / 4; i += BLOCK_SIZE) {
                int c_local = i / (H * (W / 4));
                int remainder = i % (H * (W / 4));
                int h_local = remainder / (W / 4);
                int w4_local = remainder % (W / 4);

                int c_global = group_in_sample * C_per_group + c_local;
                int offset = (sample_idx * C + c_global) * H * W + h_local * W + w4_local * 4;
                
                float4 val4 = *reinterpret_cast<const float4*>(input + offset);
                
                thread_sum += val4.x + val4.y + val4.z + val4.w;
                thread_sum_sq += val4.x * val4.x + val4.y * val4.y + val4.z * val4.z + val4.w * val4.w;
            }
            
            float warp_sum = warp_reduce_sum(thread_sum);
            float warp_sum_sq = warp_reduce_sum(thread_sum_sq);

            int warp_id = threadIdx.x / WARP_SIZE;
            int lane_id = threadIdx.x % WARP_SIZE;
            if (lane_id == 0) { s_warp_sums[warp_id] = warp_sum; s_warp_sum_sqs[warp_id] = warp_sum_sq; }
            __syncthreads();

            warp_sum = (threadIdx.x < BLOCK_SIZE / WARP_SIZE) ? s_warp_sums[lane_id] : 0.0f;
            warp_sum_sq = (threadIdx.x < BLOCK_SIZE / WARP_SIZE) ? s_warp_sum_sqs[lane_id] : 0.0f;
            if (warp_id == 0) { warp_sum = warp_reduce_sum(warp_sum); warp_sum_sq = warp_reduce_sum(warp_sum_sq); }

            if (threadIdx.x == 0) {
                s_warp_sums[0] = warp_sum / group_size; // mean
                float variance = (warp_sum_sq / group_size) - (s_warp_sums[0] * s_warp_sums[0]);
                s_warp_sum_sqs[0] = rsqrtf(variance + eps); // inv_stddev
            }
            __syncthreads();
            
            float mean = s_warp_sums[0];
            float inv_stddev = s_warp_sum_sqs[0];

            for (int i = threadIdx.x; i < group_size / 4; i += BLOCK_SIZE) {
                int c_local = i / (H * (W / 4));
                int remainder = i % (H * (W / 4));
                int h_local = remainder / (W / 4);
                int w4_local = remainder % (W / 4);
                
                int c_global = group_in_sample * C_per_group + c_local;
                int offset = (sample_idx * C + c_global) * H * W + h_local * W + w4_local * 4;

                float4 val4 = *reinterpret_cast<const float4*>(input + offset);
                float w0 = weight[c_global], b0 = bias[c_global];

                val4.x = (val4.x - mean) * inv_stddev * w0 + b0;
                val4.y = (val4.y - mean) * inv_stddev * w0 + b0;
                val4.z = (val4.z - mean) * inv_stddev * w0 + b0;
                val4.w = (val4.w - mean) * inv_stddev * w0 + b0;

                *reinterpret_cast<float4*>(output + offset) = val4;
            }
        }

        torch::Tensor groupnorm_forward_cuda(
            torch::Tensor input, torch::Tensor weight, torch::Tensor bias,
            int num_groups, float eps) {

            const int N = input.size(0);
            const int C = input.size(1);
            const int H = input.size(2);
            const int W = input.size(3);
            auto output = torch::empty_like(input);

            const int blocks = N * num_groups;
            const int threads = BLOCK_SIZE;
            size_t shared_mem_size = (threads / WARP_SIZE) * 2 * sizeof(float);

            groupnorm_fused_vectorized_kernel<<<blocks, threads, shared_mem_size>>>(
                input.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
                output.data_ptr<float>(), N, C, H, W, num_groups, eps);
            return output;
        }
        """
        self.groupnorm_op = load_inline(
            name="groupnorm_fused_vectorized_op",
            cpp_sources=cpp_source, cuda_sources=cuda_source,
            functions=["groupnorm_forward_cuda"],
            extra_cuda_cflags=["-O3", "--use_fast_math"], verbose=True
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.groupnorm_op.groupnorm_forward_cuda(
            x.contiguous(), self.weight, self.bias, self.num_groups, self.eps
        )