Kernel Details - fused_optimized_warp

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_optimized_warp_base
CUDA Speedup (Native)	1.744x
CUDA Speedup (Compile)	1.493x
CUDA Runtime	0.043 ms
PyTorch Runtime (Native)	0.075 ms
PyTorch Runtime (Compile)	0.064 ms
Correct	True
Max Diff (vs. Reference)	0.006000
Model	o3-mini-2025-01-31
Temperature	1.00

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)

// Warp-level reduction for a float value
__inline__ __device__ float warpReduceSum(float val) {
    // All threads in a warp participate
    for (int offset = warpSize/2; offset > 0; offset /= 2)
        val += __shfl_down_sync(0xffffffff, val, offset);
    return val;
}

// Fused kernel with warp-level reduction for improved efficiency
// This kernel performs bias addition, scaling, sigmoid activation and group normalization
// in a single pass, using warp shuffles to reduce shared memory overhead during reduction.
__global__ void fused_optimized_warp_kernel(
    const float* __restrict__ x,    // Input tensor after conv2d: [N, C, H, W]
    float* __restrict__ y,          // Output tensor: [N, C, H, W]
    const float* __restrict__ bias, // Bias for elementwise op (broadcastable)
    const float* __restrict__ scale,// Scale for elementwise op (broadcastable)
    const float* __restrict__ gn_weight, // GroupNorm weight, shape [C]
    const float* __restrict__ gn_bias,   // GroupNorm bias, shape [C]
    int N, int C, int H, int W,
    int num_groups,
    bool bias_broadcast,
    bool scale_broadcast,
    float eps) {

    // Determine which sample and which group this block is processing
    int group_idx = blockIdx.x % num_groups;
    int sample_idx = blockIdx.x / num_groups;
    int channels_per_group = C / num_groups;
    int group_size = channels_per_group * H * W;  // Total number of elements in this group

    // Offsets to index into the input/output tensor
    int sample_offset = sample_idx * C * H * W;
    int group_channel_offset = group_idx * channels_per_group;

    // Each thread computes partial sum and sum of squares over its assigned elements
    float local_sum = 0.0f;
    float local_sum_sq = 0.0f;
    
    // First pass: apply elementwise operations and compute activation (sigmoid)
    for (int i = threadIdx.x; i < group_size; i += blockDim.x) {
        int c_local = i / (H * W);
        int hw = i % (H * W);
        int c = group_channel_offset + c_local;
        int idx = sample_offset + c * (H * W) + hw;

        float in_val = x[idx];
        float b_val = bias_broadcast ? bias[0] : bias[c];
        float s_val = scale_broadcast ? scale[0] : scale[c];
        float pre_act = (in_val + b_val) * s_val;
        // Sigmoid activation
        float v = 1.0f / (1.0f + expf(-pre_act));
        y[idx] = v;  // store intermediate result for later use
        local_sum += v;
        local_sum_sq += v * v;
    }

    // Use warp-level reduction to reduce local sums within each warp
    float sum = warpReduceSum(local_sum);
    float sum_sq = warpReduceSum(local_sum_sq);

    int lane = threadIdx.x % warpSize;       // lane index within the warp
    int warpId = threadIdx.x / warpSize;       // current warp index

    // Allocate shared memory to store the reduction results from each warp
    extern __shared__ float shared[]; // size will be 2 * (blockDim.x/warpSize) floats
    float* warp_sum = shared;                         // first half: partial sums
    float* warp_sum_sq = shared + (blockDim.x / warpSize); // second half: partial sum of squares

    // Write warp-level results to shared memory
    if (lane == 0) {
        warp_sum[warpId] = sum;
        warp_sum_sq[warpId] = sum_sq;
    }
    __syncthreads();

    // Let first warp perform reduction over warp-level partial results
    float block_sum = 0.0f;
    float block_sum_sq = 0.0f;
    if (threadIdx.x < (blockDim.x / warpSize)) {
        block_sum = warp_sum[threadIdx.x];
        block_sum_sq = warp_sum_sq[threadIdx.x];
    }
    // Only threads in first warp participate in final reduction
    if (threadIdx.x < warpSize) {
        float val_sum = (threadIdx.x < (blockDim.x / warpSize)) ? block_sum : 0.0f;
        float val_sum_sq = (threadIdx.x < (blockDim.x / warpSize)) ? block_sum_sq : 0.0f;
        val_sum = warpReduceSum(val_sum);
        val_sum_sq = warpReduceSum(val_sum_sq);
        if (threadIdx.x == 0) {
            // Compute group mean and variance
            float group_mean = val_sum / group_size;
            float group_var = val_sum_sq / group_size - group_mean * group_mean;
            // Store these in shared memory for broadcast
            shared[0] = group_mean;
            shared[1] = group_var;
        }
    }
    __syncthreads();

    // Retrieve computed group mean and variance
    float group_mean = shared[0];
    float group_var = shared[1];
    float inv_std = 1.0f / sqrtf(group_var + eps);

    // Second pass: normalize each element and apply group normalization affine transform
    for (int i = threadIdx.x; i < group_size; i += blockDim.x) {
        int c_local = i / (H * W);
        int hw = i % (H * W);
        int c = group_channel_offset + c_local;
        int idx = sample_offset + c * (H * W) + hw;

        float v = y[idx];
        float normalized = (v - group_mean) * inv_std;
        float gamma = gn_weight[c];
        float beta = gn_bias[c];
        y[idx] = gamma * normalized + beta;
    }
}


// Launcher for the fused optimized warp kernel
void fused_optimized_warp_cuda(
    at::Tensor x,          // Input from conv2d
    at::Tensor bias,       // Bias for elementwise op
    at::Tensor scale,      // Scale for elementwise op
    at::Tensor y,          // Output tensor
    at::Tensor gn_weight,  // GroupNorm weight
    at::Tensor gn_bias,    // GroupNorm bias
    int64_t num_groups,
    bool bias_broadcast,
    bool scale_broadcast,
    float eps) {

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

    // Each block processes one group of one sample
    int total_blocks = N * num_groups;
    int threads = 256;  // Optimal thread count per block
    int warps_per_block = threads / 32;
    size_t shared_mem_size = 2 * warps_per_block * sizeof(float); // For warp-level reduction

    fused_optimized_warp_kernel<<<total_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));
    }
}

// Forward function: applies convolution then the fused optimized warp kernel
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);

    // Perform convolution
    x = at::conv2d(x, conv_weight, conv_bias);

    // Prepare output tensor
    at::Tensor y = at::empty_like(x);

    bool bias_broadcast = (bias.numel() == 1);
    bool scale_broadcast = (scale.numel() == 1);
    float eps = 1e-5;

    fused_optimized_warp_cuda(x, bias, scale, y, gn_weight, gn_bias,
                              num_groups, bias_broadcast, scale_broadcast, eps);

    return y;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &module_fn_forward, "Fused optimized warp kernel (CUDA)");
}

Performance Metrics

Metric	Value	Unit	Variance	Samples
Executed Ipc Active	2.454	inst/cycle	0.000	5
Executed Ipc Elapsed	1.972	inst/cycle	0.000	5
Issue Slots Busy	61.518	%	0.271	5
Issued Ipc Active	2.462	inst/cycle	0.000	5
SM Busy	61.518	%	0.271	5
Memory Throughput	475392179057.796	byte/second	26896129338345553920.000	5
Mem Busy	22.658	%	0.059	5
Max Bandwidth	27.662	%	0.100	5
L1/TEX Hit Rate	65.524	%	0.001	5
L2 Hit Rate	70.538	%	0.000	5
Mem Pipes Busy	19.000	%	0.050	5
Warp Cycles Per Issued Instruction	19.686	cycle	0.013	5
Warp Cycles Per Executed Instruction	19.742	cycle	0.013	5
Avg. Active Threads Per Warp	31.140		0.000	5
Avg. Not Predicated Off Threads Per Warp	27.950		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	8.000	block	0.000	5
Block Limit Shared Mem	28.000	block	0.000	5
Block Limit Warps	8.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	76.294	%	0.122	5
Achieved Active Warps Per SM	48.828	warp	0.050	5

Analysis Rules

Rule	Description
INF HighPipeUtilization	ALU is the highest-utilized pipeline (44.7%) 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 (76.0%) 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.7%) 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 (76.0%) 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	166063.80	μs
Device Time	1067315.88	μs
Self CPU Time	30689.10	μs
Self Device Time	1067315.88	μ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	189262.90	μs
Device Time	1067315.88	μs
Self CPU Time	23232.21	μ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	1171241.12	μs
Device Time	380719.85	μs
Self CPU Time	22582.24	μ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	1148658.88	μs
Device Time	380719.85	μs
Self CPU Time	27994.68	μ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	1120664.20	μs
Device Time	380719.85	μs
Self CPU Time	55387.62	μ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	933751.53	μs
Device Time	271647.75	μs
Self CPU Time	222948.83	μs
Self Device Time	271647.75	μ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	954330.56	μs
Device Time	27619.95	μs
Self CPU Time	954330.56	μs
Self Device Time	27619.95	μ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	1067315.88	μs
Self CPU Time	0.00	μs
Self Device Time	1067315.88	μ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

45314 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/b8_s2_fused_optimized_warp/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/b8_s2_fused_optimized_warp/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/b8_s2_fused_optimized_warp/base/base.cu:25:5: warning: 4 adjacent parameters of 'fused_optimized_warp_kernel' of similar type ('const float *__restrict') are easily swapped by mistake [bugprone-easily-swappable-parameters]

25 | const float* __restrict__ bias, // Bias for elementwise op (broadcastable)

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

26 | const float* __restrict__ scale,// Scale for elementwise op (broadcastable)

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

27 | const float* __restrict__ gn_weight, // GroupNorm weight, shape [C]

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

28 | const float* __restrict__ gn_bias, // GroupNorm bias, shape [C]

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

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

25 | const float* __restrict__ bias, // Bias for elementwise op (broadcastable)

| ^~~~

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

28 | const float* __restrict__ gn_bias, // GroupNorm bias, shape [C]

| ^~~~~~~

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