Kernel Details - fused_vec_gridstride

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


def module_fn(
    x: torch.Tensor,
    stride: int,
    padding: int,
    groups: int,
    eps: float,
    conv_transpose: torch.Tensor,
    conv_transpose_bias: torch.Tensor,
    group_norm_weight: torch.Tensor,
    group_norm_bias: torch.Tensor,
) -> torch.Tensor:
    """
    Applies 3D transposed convolution, Swish activation, group normalization and HardSwish activation.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, in_channels, depth, height, width)
        stride (int): Stride of the transposed convolution
        padding (int): Padding of the transposed convolution
        groups (int): Number of groups for group normalization
        eps (float): Epsilon value for group normalization
        conv_transpose (torch.Tensor): Transposed convolution weight tensor
        conv_transpose_bias (torch.Tensor): Bias tensor for transposed convolution
        group_norm_weight (torch.Tensor): Weight tensor for group normalization
        group_norm_bias (torch.Tensor): Bias tensor for group normalization

    Returns:
        torch.Tensor: Output tensor after applying all operations
    """
    x = F.conv_transpose3d(
        x, conv_transpose, bias=conv_transpose_bias, stride=stride, padding=padding
    )
    x = torch.sigmoid(x) * x  # Swish activation
    x = F.group_norm(
        x, num_groups=groups, weight=group_norm_weight, bias=group_norm_bias, eps=eps
    )
    x = F.hardswish(x)
    return x


class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, applies Swish activation,
    group normalization, and then HardSwish activation.
    """

    def __init__(
        self, in_channels, out_channels, kernel_size, stride, padding, groups, eps
    ):
        super(Model, self).__init__()
        conv = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size, stride=stride, padding=padding
        )
        self.conv_transpose_parameter = nn.Parameter(conv.weight)
        self.conv_transpose_bias = nn.Parameter(conv.bias)
        gn = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps)
        self.group_norm_weight = nn.Parameter(gn.weight)
        self.group_norm_bias = nn.Parameter(gn.bias + torch.randn(out_channels) * 0.02)

    def forward(self, x, stride, padding, groups, eps, fn=module_fn):
        return fn(
            x,
            stride,
            padding,
            groups,
            eps,
            self.conv_transpose_parameter,
            self.conv_transpose_bias,
            self.group_norm_weight,
            self.group_norm_bias,
        )


batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1
groups = 4
eps = 1e-5


def get_inputs():
    return [
        torch.randn(batch_size, in_channels, depth, height, width),
        stride,
        padding,
        groups,
        eps,
    ]


def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, groups, eps]

import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, applies Swish activation, 
    group normalization, and then HardSwish activation.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, groups, eps, bias=True):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, bias=bias)
        self.group_norm = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps)
        # Add noise to group norm bias to match functional implementation
        self.group_norm.bias = nn.Parameter(self.group_norm.bias + torch.randn(out_channels) * 0.02)

    def forward(self, x):
        x = self.conv_transpose(x)
        x = torch.sigmoid(x) * x  # Swish activation
        x = self.group_norm(x)
        x = torch.nn.functional.hardswish(x)  # HardSwish activation
        return x

batch_size = 128
in_channels = 3
out_channels = 16
depth, height, width = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1
groups = 4
eps = 1e-5

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding, groups, eps]

Download Evaluation Download PyTorch Download CUDA Download Profiles

Kernel Information

Operation Name	60_ConvTranspose3d_Swish_GroupNorm_HardSwish
Level ID	2
Task ID	60
Kernel Name	fused_vec_gridstride_base
CUDA Speedup (Native)	1.417x
CUDA Speedup (Compile)	0.836x
CUDA Runtime	5.427 ms
PyTorch Runtime (Native)	7.690 ms
PyTorch Runtime (Compile)	4.539 ms
Correct	True
Max Diff (vs. Reference)	0.000000
Model	o3-mini-2025-01-31
Temperature	1.00

View Experiment Progress Details

Related Kernels (Level 2, Task 60 • 60_ConvTranspose3d_Swish_GroupNorm_HardSwish)

Rank	Kernel Name	Runtime (ms)	Speedup Native	Speedup Compile
🥇	warp_optimized_fused_base_base	5.40	1.42	0.84
🥈	fused_min_sync_base	5.41	1.42	0.84
🥉	optimized_fused_activation_norm_base	5.41	1.42	0.84
4	fused_swish_gn_hardswish_base	5.42	1.42	0.84
5	fused_vec_gridstride_base	5.43	1.42	0.84
6	aligned_memory_fused_swish_grpnorm_hardswish_base_base	5.43	1.42	0.84
7	optimized_reduction_fused_actnorm_base	5.44	1.41	0.84
8	aligned_vectorized_fused_kernel_base_base	5.44	1.41	0.83
9	streamed_fused_activation_norm_base	5.45	1.41	0.83
10	fused_swish_grpnorm_hardswish_fast_red_base	6.00	1.28	0.76
10	optimized_fused_swish_grpnorm_hardswish_sync_reduction_base	6.00	1.28	0.76
10	fused_swish_grpnorm_hardswish_stride_base	6.00	1.28	0.76
13	optimized_fused_activation_grpnorm_base	6.00	1.28	0.76
14	fused_swish_groupnorm_hardswish_optimized_base	6.00	1.28	0.76
15	fused_swish_grpnorm_hardswish_optimized_base	6.00	1.28	0.76
16	optimized_fused_swish_grpnorm_hardswish_coalesced_base_base	6.00	1.28	0.76
17	fused_swish_grpnorm_hardswish_warp_base	6.01	1.28	0.76
18	fused_swish_grpnorm_hardswish_stride_base	6.01	1.28	0.76
19	efficient_fused_kernel_base	6.01	1.28	0.76
20	fused_swish_grpnorm_hardswish_unrolled_base	6.02	1.28	0.75

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

#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)

// This fused kernel combines vectorized memory accesses (float4 loads/stores) with grid-stride loops
// to handle arbitrary group sizes. It fuses Swish activation, Group Normalization (with warp-level reduction),
// and HardSwish activation in two phases: first, a reduction phase to compute mean/variance using vectorized loads
// where possible; second, a normalization and activation phase that re-computes Swish and applies the affine transform.

__global__ void fused_vec_gridstride_kernel(
    const float* __restrict__ input,
    float* __restrict__ output,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    int N, int C, int D, int H, int W,
    int groups, float eps
) {
    // Each block processes one group from one sample
    int n = blockIdx.x;
    int g = blockIdx.y;

    int channels_per_group = C / groups;
    int group_elements = channels_per_group * D * H * W;
    int base = n * (C * D * H * W) + g * group_elements;

    const int VECTOR_SIZE = 4;
    int aligned_elements = (group_elements / VECTOR_SIZE) * VECTOR_SIZE;  // number of elements that can be processed vectorized

    int tid = threadIdx.x;
    int blockSize = blockDim.x;

    // Phase 1: Compute partial sums and sums-of-squares for the Swish activated values
    float local_sum = 0.0f;
    float local_sumsq = 0.0f;

    // Use vectorized loads for the aligned portion
    for (int i = tid * VECTOR_SIZE; i < aligned_elements; i += blockSize * VECTOR_SIZE) {
        int idx = base + i;
        // Load four consecutive floats assuming alignment
        float4 vec = *reinterpret_cast<const float4*>(&input[idx]);
        float x0 = vec.x, x1 = vec.y, x2 = vec.z, x3 = vec.w;
        float sw0 = x0 / (1.f + expf(-x0));
        float sw1 = x1 / (1.f + expf(-x1));
        float sw2 = x2 / (1.f + expf(-x2));
        float sw3 = x3 / (1.f + expf(-x3));
        local_sum += sw0 + sw1 + sw2 + sw3;
        local_sumsq += sw0 * sw0 + sw1 * sw1 + sw2 * sw2 + sw3 * sw3;
    }

    // Handle any tail elements that are not vectorized
    for (int i = aligned_elements + tid; i < group_elements; i += blockSize) {
        int idx = base + i;
        float x = __ldg(&input[idx]);
        float sw = x / (1.f + expf(-x));
        local_sum += sw;
        local_sumsq += sw * sw;
    }

    // Warp-level reduction using shuffles
    unsigned int mask = 0xffffffff;
    for (int offset = warpSize/2; offset > 0; offset /= 2) {
        local_sum += __shfl_down_sync(mask, local_sum, offset);
        local_sumsq += __shfl_down_sync(mask, local_sumsq, offset);
    }

    // Shared memory for storing warp-level partial results
    extern __shared__ float shared[]; // shared memory: first half for sums, second half for sumsq
    int numWarps = (blockSize + warpSize - 1) / warpSize;
    float* warpSums = shared;
    float* warpSumsSq = shared + numWarps;

    int warpId = tid / warpSize;
    int laneId = tid % warpSize;
    if(laneId == 0) {
        warpSums[warpId] = local_sum;
        warpSumsSq[warpId] = local_sumsq;
    }
    __syncthreads();

    float total_sum = 0.f;
    float total_sumsq = 0.f;
    if(tid < numWarps) {
        total_sum = warpSums[tid];
        total_sumsq = warpSumsSq[tid];
    }
    __syncthreads();
    if(tid == 0) {
        for (int i = 1; i < numWarps; i++) {
            total_sum += warpSums[i];
            total_sumsq += warpSumsSq[i];
        }
        warpSums[0] = total_sum;  // store final reduction in shared memory
        warpSumsSq[0] = total_sumsq;
    }
    __syncthreads();

    total_sum = warpSums[0];
    total_sumsq = warpSumsSq[0];
    float mean = total_sum / group_elements;
    float var = total_sumsq / group_elements - mean * mean;
    float inv_std = rsqrtf(var + eps);

    // Phase 2: Recompute the activation and apply normalization + HardSwish
    // Use vectorized stores for the aligned part
    for (int i = tid * VECTOR_SIZE; i < aligned_elements; i += blockSize * VECTOR_SIZE) {
        int idx = base + i;
        float4 in_vec = *reinterpret_cast<const float4*>(&input[idx]);
        float4 out_vec;
        #pragma unroll
        for (int j = 0; j < VECTOR_SIZE; j++) {
            float x = ((&in_vec.x)[j]);
            float sw = x / (1.f + expf(-x));
            int local_channel = (i + j) / (D * H * W);
            int global_channel = g * channels_per_group + local_channel;
            float norm = (sw - mean) * inv_std;
            float y = norm * __ldg(&gamma[global_channel]) + __ldg(&beta[global_channel]);
            float hs = y * fminf(fmaxf(y + 3.f, 0.f), 6.f) / 6.f;
            (&out_vec.x)[j] = hs;
        }
        *reinterpret_cast<float4*>(&output[idx]) = out_vec;
    }

    // Process remaining tail elements
    for (int i = aligned_elements + tid; i < group_elements; i += blockSize) {
        int idx = base + i;
        float x = __ldg(&input[idx]);
        float sw = x / (1.f + expf(-x));
        int local_channel = i / (D * H * W);
        int global_channel = g * channels_per_group + local_channel;
        float norm = (sw - mean) * inv_std;
        float y = norm * __ldg(&gamma[global_channel]) + __ldg(&beta[global_channel]);
        output[idx] = y * fminf(fmaxf(y + 3.f, 0.f), 6.f) / 6.f;
    }
}

// The forward function first applies a conv_transpose3d then launches the fused kernel
torch::Tensor forward(
    torch::Tensor x,
    int stride,
    int padding,
    int groups,
    float eps,
    torch::Tensor conv_transpose,
    torch::Tensor conv_transpose_bias,
    torch::Tensor group_norm_weight,
    torch::Tensor group_norm_bias
) {
    CHECK_INPUT(x);
    CHECK_INPUT(conv_transpose);
    CHECK_INPUT(conv_transpose_bias);
    CHECK_INPUT(group_norm_weight);
    CHECK_INPUT(group_norm_bias);

    // Perform the transposed convolution
    x = torch::conv_transpose3d(x, conv_transpose, conv_transpose_bias, stride, padding);
    torch::Tensor output = torch::empty_like(x);

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

    // Each block processes one (sample, group) pair
    dim3 grid(N, groups);
    int threads = 256;
    int numWarps = (threads + 31) / 32;
    size_t shared_mem = 2 * numWarps * sizeof(float);

    fused_vec_gridstride_kernel<<<grid, threads, shared_mem>>>(
        x.data_ptr<float>(),
        output.data_ptr<float>(),
        group_norm_weight.data_ptr<float>(),
        group_norm_bias.data_ptr<float>(),
        N, C, D, H, W,
        groups, eps
    );

    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Fused vectorized and grid-stride kernel for Swish, GroupNorm, and HardSwish");
}

Performance Metrics

Metric	Value	Unit	Variance	Samples
Executed Ipc Active	1.450	inst/cycle	0.000	5
Executed Ipc Elapsed	1.418	inst/cycle	0.000	5
Issue Slots Busy	36.318	%	0.001	5
Issued Ipc Active	1.450	inst/cycle	0.000	5
SM Busy	36.318	%	0.001	5
Memory Throughput	2475401238144.612	byte/second	30102299843454754816.000	5
Mem Busy	40.378	%	0.008	5
Max Bandwidth	73.844	%	0.028	5
L1/TEX Hit Rate	16.732	%	0.001	5
L2 Hit Rate	35.658	%	0.000	5
Mem Pipes Busy	8.434	%	0.000	5
Warp Cycles Per Issued Instruction	20.762	cycle	0.001	5
Warp Cycles Per Executed Instruction	20.762	cycle	0.001	5
Avg. Active Threads Per Warp	31.990		0.000	5
Avg. Not Predicated Off Threads Per Warp	29.220		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	47.172	%	0.000	5
Achieved Active Warps Per SM	30.190	warp	0.000	5

Analysis Rules

Rule	Description
INF HighPipeUtilization	ALU is the highest-utilized pipeline (25.9%) 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.
INF CPIStall	Check the Warp Stall Sampling (All Cycles) table for the top stall locations in your source based on sampling data. The Kernel Profiling Guide (https://docs.nvidia.com/nsight-compute/ProfilingGuide/index.html#metrics-reference) provides more details on each stall reason.
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 (47.2%) 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::conv_transpose3d
CPU Time	3770321.47	μs
Device Time	6752810.09	μs
Self CPU Time	3173.17	μ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	3767148.29	μs
Device Time	6752810.09	μs
Self CPU Time	4449.00	μ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	3762699.29	μs
Device Time	6752810.09	μs
Self CPU Time	9234.59	μ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_transpose
CPU Time	406661.07	μs
Device Time	5380651.77	μs
Self CPU Time	103036.97	μs
Self Device Time	5380651.77	μ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	6684988.92	μs
Device Time	62101.50	μs
Self CPU Time	6684988.92	μs
Self Device Time	62101.50	μs
CPU Memory Usage	0	B
Device Memory Usage	0	B
Self CPU Memory Usage	0	B
Self Device Memory Usage	0	B
sm90_xmma_dgrad_implicit_gemm_indexed_f32f32_tf32f32_f32_nhwckrsc_nhwc_tilesize256x64x32_warpgroupsize1x1x1_g1_strided_execute_kernel__5x_cudnn
CPU Time	0.00	μs
Device Time	3554116.68	μs
Self CPU Time	0.00	μs
Self Device Time	3554116.68	μs
CPU Memory Usage	0	B
Device Memory Usage	0	B
Self CPU Memory Usage	0	B
Self Device Memory Usage	0	B
aten::add_
CPU Time	3344170.66	μs
Device Time	1372158.32	μs
Self CPU Time	8375.91	μs
Self Device Time	1372158.32	μ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

45317 warnings generated when compiling for host.
Suppressed 45344 warnings (45297 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_60/b8_s3_fused_vec_gridstride/base/base.cu:9:35 bugprone-macro-parentheses

9 | #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_60/b8_s3_fused_vec_gridstride/base/base.cu:10:41: warning: macro argument should be enclosed in parentheses [bugprone-macro-parentheses]

10 | #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_60/b8_s3_fused_vec_gridstride/base/base.cu:23:5: warning: 2 adjacent parameters of 'fused_vec_gridstride_kernel' of similar type ('int') are easily swapped by mistake [bugprone-easily-swappable-parameters]