Kernel Details - fused_strided_groupnorm_base

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


def module_fn(
    x: torch.Tensor,
    conv_weight: torch.Tensor,
    conv_bias: torch.Tensor,
    bias: torch.Tensor,
    scale: torch.Tensor,
    group_norm_weight: torch.Tensor,
    group_norm_bias: torch.Tensor,
    num_groups: int,
) -> torch.Tensor:
    """
    Applies convolution, bias addition, scaling, sigmoid activation and group normalization.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width)
        conv_weight (torch.Tensor): Convolution weight tensor
        conv_bias (torch.Tensor): Convolution bias tensor
        bias (torch.Tensor): Bias tensor for addition
        scale (torch.Tensor): Scale tensor for multiplication
        group_norm_weight (torch.Tensor): Group norm weight tensor
        group_norm_bias (torch.Tensor): Group norm bias tensor
        num_groups (int): Number of groups for group normalization

    Returns:
        torch.Tensor: Output tensor after applying convolution, bias, scale, sigmoid and group norm
    """
    x = F.conv2d(x, conv_weight, bias=conv_bias)
    x = x + bias
    x = x * scale
    x = torch.sigmoid(x)
    x = F.group_norm(x, num_groups, group_norm_weight, group_norm_bias)
    return x


class Model(nn.Module):
    """
    Model that performs a convolution, adds a bias term, scales, applies sigmoid, and performs group normalization.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size,
        num_groups,
        bias_shape,
        scale_shape,
    ):
        super(Model, self).__init__()
        conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.conv_weight = conv.weight
        self.conv_bias = nn.Parameter(
            conv.bias + torch.ones_like(conv.bias) * 0.02
        )  # make sure its nonzero
        self.bias = nn.Parameter(torch.randn(bias_shape) * 0.02)
        self.scale = nn.Parameter(torch.randn(scale_shape) * 0.02)
        group_norm = nn.GroupNorm(num_groups, out_channels)
        self.group_norm_weight = group_norm.weight
        self.group_norm_bias = nn.Parameter(
            group_norm.bias + torch.ones_like(group_norm.bias) * 0.02
        )  # make sure its nonzero

    def forward(self, x, num_groups, fn=module_fn):
        return fn(
            x,
            self.conv_weight,
            self.conv_bias,
            self.bias,
            self.scale,
            self.group_norm_weight,
            self.group_norm_bias,
            num_groups,
        )


batch_size = 128
in_channels = 3
out_channels = 16
height, width = 32, 32
kernel_size = 3
num_groups = 8
bias_shape = (out_channels, 1, 1)
scale_shape = (out_channels, 1, 1)


def get_inputs():
    return [torch.randn(batch_size, in_channels, height, width), num_groups]


def get_init_inputs():
    return [in_channels, out_channels, kernel_size, num_groups, bias_shape, scale_shape]

import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a convolution, adds a bias term, scales, applies sigmoid, and performs group normalization.
    """
    def __init__(self, in_channels, out_channels, kernel_size, num_groups, bias_shape, scale_shape):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.conv.bias = nn.Parameter(self.conv.bias + torch.ones_like(self.conv.bias) * 0.02)
        self.bias = nn.Parameter(torch.randn(bias_shape)*0.02) 
        self.scale = nn.Parameter(torch.randn(scale_shape)*0.02)
        self.group_norm = nn.GroupNorm(num_groups, out_channels)
        self.group_norm.bias = nn.Parameter(self.group_norm.bias + torch.ones_like(self.group_norm.bias) * 0.02)

    def forward(self, x):
        x = self.conv(x)
        x = x + self.bias
        x = x * self.scale
        x = torch.sigmoid(x)
        x = self.group_norm(x)
        return x

batch_size = 128
in_channels = 3
out_channels = 16
height, width = 32, 32
kernel_size = 3
num_groups = 8
bias_shape = (out_channels, 1, 1)
scale_shape = (out_channels, 1, 1)

def get_inputs():
    return [torch.randn(batch_size, in_channels, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, num_groups, bias_shape, scale_shape]

Download Evaluation Download PyTorch Download CUDA Download Profiles

Kernel Information

Operation Name	21_Conv2d_Add_Scale_Sigmoid_GroupNorm
Level ID	2
Task ID	21
Kernel Name	fused_strided_groupnorm_base_base
CUDA Speedup (Native)	1.666x
CUDA Speedup (Compile)	1.427x
CUDA Runtime	0.045 ms
PyTorch Runtime (Native)	0.075 ms
PyTorch Runtime (Compile)	0.064 ms
Correct	True
Max Diff (vs. Reference)	0.005000
Model	bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0
Temperature	0.50

View Experiment Progress Details

Related Kernels (Level 2, Task 21 • 21_Conv2d_Add_Scale_Sigmoid_GroupNorm)

Rank	Kernel Name	Runtime (ms)	Speedup Native	Speedup Compile
🥇	shared_memory_coalesced_access_kernel_base	0.04	1.79	1.53
🥈	fused_elem_groupnorm_reduced_sync_base_base	0.04	1.74	1.49
🥈	fused_optimized_warp_base	0.04	1.74	1.49
🥈	optimized_memory_access_kernel_base	0.04	1.74	1.49
🥈	atomic_optimized_kernel_base_base	0.04	1.74	1.49
🥈	optimized_memory_access_kernel_base	0.04	1.74	1.49
🥈	optimized_fused_kernel_base	0.04	1.74	1.49
8	fused_warp_reduce_groupnorm_base	0.04	1.70	1.46
8	fused_warp_groupnorm_base	0.04	1.70	1.46
8	shared_memory_reuse_kernel_base	0.04	1.70	1.46
8	fused_lockfree_groupnorm_base_base	0.04	1.70	1.46
8	fused_stride_kernel_base	0.04	1.70	1.46
8	fused_elem_groupnorm_no_atomic_base	0.04	1.70	1.46
8	fused_elem_groupnorm_min_sync_base_base	0.04	1.70	1.46
8	unrolled_fused_kernel_base_base	0.04	1.70	1.46
16	optimized_modular_kernel_base	0.04	1.67	1.43
16	fused_strided_groupnorm_base_base	0.04	1.67	1.43
18	fused_sigmoid_groupnorm_base	0.05	1.63	1.40
19	fused_sigmoid_groupnorm_base	0.05	1.60	1.37
19	block_size_optimized_kernel_base_base	0.05	1.60	1.37

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

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor.")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous.")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

__global__ void fused_strided_groupnorm_kernel(
    const float* __restrict__ x,
    float* __restrict__ y,
    const float* __restrict__ bias,
    const float* __restrict__ scale,
    const float* __restrict__ gn_weight,
    const float* __restrict__ gn_bias,
    int N, int C, int H, int W,
    int num_groups,
    bool bias_broadcast,
    bool scale_broadcast,
    float eps) {

    extern __shared__ float shared_mem[];
    float* shared_sum = shared_mem;
    float* shared_sum_sq = &shared_mem[blockDim.x];

    const int group_idx = blockIdx.x % num_groups;
    const int sample_idx = blockIdx.x / num_groups;
    const int channels_per_group = C / num_groups;
    const int group_size = channels_per_group * H * W;
    const int sample_offset = sample_idx * C * H * W;
    const int group_offset = group_idx * channels_per_group;

    // Initialize thread accumulators
    float local_sum = 0.0f;
    float local_sum_sq = 0.0f;

    // Stride loop for first pass
    const int tid = threadIdx.x;
    const int stride = blockDim.x;
    
    #pragma unroll 4
    for (int i = tid; i < group_size; i += stride) {
        const int c_local = i / (H * W);
        const int hw = i % (H * W);
        const int c = group_offset + c_local;
        const int idx = sample_offset + c * H * W + hw;

        const float x_val = __ldg(&x[idx]);
        const float b_val = bias_broadcast ? __ldg(&bias[0]) : __ldg(&bias[c]);
        const float s_val = scale_broadcast ? __ldg(&scale[0]) : __ldg(&scale[c]);
        
        const float val = 1.0f / (1.0f + expf(-(x_val + b_val) * s_val));
        y[idx] = val;  // Store intermediate result
        
        local_sum += val;
        local_sum_sq += val * val;
    }

    // Parallel reduction in shared memory
    shared_sum[tid] = local_sum;
    shared_sum_sq[tid] = local_sum_sq;
    __syncthreads();

    #pragma unroll
    for (int s = blockDim.x/2; s > 0; s >>= 1) {
        if (tid < s) {
            shared_sum[tid] += shared_sum[tid + s];
            shared_sum_sq[tid] += shared_sum_sq[tid + s];
        }
        __syncthreads();
    }

    // Compute statistics
    const float mean = shared_sum[0] / group_size;
    const float variance = (shared_sum_sq[0] / group_size) - (mean * mean);
    const float inv_std = rsqrtf(variance + eps);

    __syncthreads();

    // Second pass with stride loop
    #pragma unroll 4
    for (int i = tid; i < group_size; i += stride) {
        const int c_local = i / (H * W);
        const int hw = i % (H * W);
        const int c = group_offset + c_local;
        const int idx = sample_offset + c * H * W + hw;

        const float val = y[idx];
        const float normalized = (val - mean) * inv_std;
        const float gamma = __ldg(&gn_weight[c]);
        const float beta = __ldg(&gn_bias[c]);
        
        y[idx] = gamma * normalized + beta;
    }
}

void fused_strided_groupnorm_cuda(
    at::Tensor x,
    at::Tensor bias,
    at::Tensor scale,
    at::Tensor y,
    at::Tensor gn_weight,
    at::Tensor gn_bias,
    int64_t num_groups,
    bool bias_broadcast,
    bool scale_broadcast,
    float eps) {

    const int N = x.size(0);
    const int C = x.size(1);
    const int H = x.size(2);
    const int W = x.size(3);

    const int threads = 512;
    const int blocks = N * num_groups;
    const int shared_mem_size = 2 * threads * sizeof(float);

    fused_strided_groupnorm_kernel<<<blocks, threads, shared_mem_size>>>(
        x.data_ptr<float>(),
        y.data_ptr<float>(),
        bias.data_ptr<float>(),
        scale.data_ptr<float>(),
        gn_weight.data_ptr<float>(),
        gn_bias.data_ptr<float>(),
        N, C, H, W,
        num_groups,
        bias_broadcast,
        scale_broadcast,
        eps);

    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        TORCH_CHECK(false, "CUDA kernel failed : ", cudaGetErrorString(err));
    }
}

at::Tensor module_fn_forward(
    at::Tensor x,
    at::Tensor conv_weight,
    at::Tensor conv_bias,
    at::Tensor bias,
    at::Tensor scale,
    at::Tensor gn_weight,
    at::Tensor gn_bias,
    int64_t num_groups) {

    CHECK_INPUT(x);
    CHECK_INPUT(conv_weight);
    if (conv_bias.defined()) CHECK_INPUT(conv_bias);
    CHECK_INPUT(bias);
    CHECK_INPUT(scale);
    CHECK_INPUT(gn_weight);
    CHECK_INPUT(gn_bias);

    x = at::conv2d(x, conv_weight, conv_bias);
    at::Tensor y = at::empty_like(x);

    fused_strided_groupnorm_cuda(
        x, bias, scale, y, gn_weight, gn_bias,
        num_groups,
        bias.numel() == 1,
        scale.numel() == 1,
        1e-5);

    return y;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &module_fn_forward, "Fused strided groupnorm forward (CUDA)");
}

Performance Metrics

Metric	Value	Unit	Variance	Samples
Executed Ipc Active	2.488	inst/cycle	0.000	5
Executed Ipc Elapsed	2.074	inst/cycle	0.001	5
Issue Slots Busy	62.486	%	0.005	5
Issued Ipc Active	2.498	inst/cycle	0.000	5
SM Busy	62.486	%	0.005	5
Memory Throughput	405284155711.290	byte/second	24008741076730912768.000	5
Mem Busy	23.310	%	0.091	5
Max Bandwidth	23.508	%	0.087	5
L1/TEX Hit Rate	65.308	%	0.001	5
L2 Hit Rate	69.906	%	0.107	5
Mem Pipes Busy	22.118	%	0.081	5
Warp Cycles Per Issued Instruction	21.982	cycle	0.042	5
Warp Cycles Per Executed Instruction	22.066	cycle	0.043	5
Avg. Active Threads Per Warp	31.560		0.000	5
Avg. Not Predicated Off Threads Per Warp	27.670		0.000	5
Max Active Clusters	0.000	cluster	0.000	5
Max Cluster Size	8.000	block	0.000	5
Overall GPU Occupancy	0.000	%	0.000	5
Cluster Occupancy	0.000	%	0.000	5
Block Limit SM	32.000	block	0.000	5
Block Limit Registers	4.000	block	0.000	5
Block Limit Shared Mem	12.000	block	0.000	5
Block Limit Warps	4.000	block	0.000	5
Theoretical Active Warps per SM	64.000	warp	0.000	5
Theoretical Occupancy	100.000	%	0.000	5
Achieved Occupancy	85.554	%	0.008	5
Achieved Active Warps Per SM	54.752	warp	0.004	5

Analysis Rules

Rule	Description
INF HighPipeUtilization	ALU is the highest-utilized pipeline (44.6%) based on active cycles, taking into account the rates of its different instructions. It executes integer and logic operations. It is well-utilized, but should not be a bottleneck.
WRN Occupancy	This kernel's theoretical occupancy is not impacted by any block limit. The difference between calculated theoretical (100.0%) and measured achieved occupancy (85.4%) can be the result of warp scheduling overheads or workload imbalances during the kernel execution. Load imbalances can occur between warps within a block as well as across blocks of the same kernel. See the CUDA Best Practices Guide (https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html#occupancy) for more details on optimizing occupancy.

Rule

Description

INF HighPipeUtilization

ALU is the highest-utilized pipeline (44.6%) based on active cycles, taking into account the rates of its different instructions. It executes integer and logic operations. It is well-utilized, but should not be a bottleneck.

WRN Occupancy

This kernel's theoretical occupancy is not impacted by any block limit. The difference between calculated theoretical (100.0%) and measured achieved occupancy (85.4%) can be the result of warp scheduling overheads or workload imbalances during the kernel execution. Load imbalances can occur between warps within a block as well as across blocks of the same kernel. See the CUDA Best Practices Guide (https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html#occupancy) for more details on optimizing occupancy.

Operation / Metric	Value	Unit
aten::fill_
CPU Time	149991.47	μs
Device Time	992538.44	μs
Self CPU Time	37517.74	μs
Self Device Time	992459.28	μs
CPU Memory Usage	0	B
Device Memory Usage	0	B
Self CPU Memory Usage	0	B
Self Device Memory Usage	0	B
aten::zero_
CPU Time	180059.05	μs
Device Time	992538.44	μs
Self CPU Time	30088.30	μs
Self Device Time	0.00	μs
CPU Memory Usage	0	B
Device Memory Usage	0	B
Self CPU Memory Usage	0	B
Self Device Memory Usage	0	B
aten::conv2d
CPU Time	995965.16	μs
Device Time	355344.08	μs
Self CPU Time	28806.15	μs
Self Device Time	0.00	μs
CPU Memory Usage	0	B
Device Memory Usage	0	B
Self CPU Memory Usage	0	B
Self Device Memory Usage	0	B
aten::convolution
CPU Time	967159.01	μs
Device Time	355344.08	μs
Self CPU Time	35294.75	μs
Self Device Time	0.00	μs
CPU Memory Usage	0	B
Device Memory Usage	0	B
Self CPU Memory Usage	0	B
Self Device Memory Usage	0	B
aten::_convolution
CPU Time	931864.26	μs
Device Time	355344.08	μs
Self CPU Time	71019.54	μs
Self Device Time	0.00	μs
CPU Memory Usage	0	B
Device Memory Usage	0	B
Self CPU Memory Usage	0	B
Self Device Memory Usage	0	B
aten::cudnn_convolution
CPU Time	710153.51	μs
Device Time	252294.03	μs
Self CPU Time	252572.90	μs
Self Device Time	252294.03	μs
CPU Memory Usage	0	B
Device Memory Usage	0	B
Self CPU Memory Usage	0	B
Self Device Memory Usage	0	B
cudaLaunchKernel
CPU Time	664857.43	μs
Device Time	25679.89	μs
Self CPU Time	664857.43	μs
Self Device Time	25679.89	μs
CPU Memory Usage	0	B
Device Memory Usage	0	B
Self CPU Memory Usage	0	B
Self Device Memory Usage	0	B
void at::native::vectorized_elementwise_kernel<4, at::native::FillFunctor<int>, at::detail::Array<char, 1> >(int, at::native::FillFunctor<int>, at::detail::Array<char, 1>)
CPU Time	0.00	μs
Device Time	992537.61	μs
Self CPU Time	0.00	μs
Self Device Time	992537.61	μs
CPU Memory Usage	0	B
Device Memory Usage	0	B
Self CPU Memory Usage	0	B
Self Device Memory Usage	0	B

Status: Completed

45311 warnings generated when compiling for host.
Suppressed 45328 warnings (45281 in non-user code, 47 NOLINT).
Use -header-filter=.* to display errors from all non-system headers. Use -system-headers to display errors from system headers as well.

/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_21/b3_s1_fused_strided_groupnorm_base/base/base.cu:7:35 bugprone-macro-parentheses

7 | #define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor.")

| ^

| ()

/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_21/b3_s1_fused_strided_groupnorm_base/base/base.cu:8:41: warning: macro argument should be enclosed in parentheses [bugprone-macro-parentheses]

8 | #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous.")

| ^

| ()

/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_21/b3_s1_fused_strided_groupnorm_base/base/base.cu:14:5: warning: 4 adjacent parameters of 'fused_strided_groupnorm_kernel' of similar type ('const float *__restrict') are easily swapped by mistake [bugprone-easily-swappable-parameters]

14 | const float* __restrict__ bias,

| ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

15 | const float* __restrict__ scale,

| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

16 | const float* __restrict__ gn_weight,

| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

17 | const float* __restrict__ gn_bias,

| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_21/b3_s1_fused_strided_groupnorm_base/base/base.cu:14:31: note: the first parameter in the range is 'bias'

14 | const float* __restrict__ bias,

| ^~~~

/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_21/b3_s1_fused_strided_groupnorm_base/base/base.cu:17:31: note: the last parameter in the range is 'gn_bias'

17 | const float* __restrict__ gn_bias,

| ^~~~~~~

/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_21/b3_s1_fused_strided_groupnorm_base/base/base.cu:18:5: warning: 2 adjacent parameters of 'fused_strided_groupnorm_kernel' of similar type ('int') are easily swapped by mistake [bugprone-easily-swappable-parameters]