← Back to Leaderboard

The AI CUDA Engineer 👷

61_ConvTranspose3d_ReLU_GroupNorm • fused_rg_const_base_base

Level 2 • Task 61
Download Evaluation Download PyTorch Download CUDA Download Profiles

Kernel Information

Operation Name 61_ConvTranspose3d_ReLU_GroupNorm
Level ID 2
Task ID 61
Kernel Name fused_rg_const_base_base
CUDA Speedup (Native) 1.160x
CUDA Speedup (Compile) 0.989x
CUDA Runtime 0.213 ms
PyTorch Runtime (Native) 0.247 ms
PyTorch Runtime (Compile) 0.211 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 61 • 61_ConvTranspose3d_ReLU_GroupNorm)

Rank Kernel Name Runtime (ms) Speedup Native Speedup Compile
🥇 fused_rg_atomic_opt_base_base 0.18 1.34 1.14
🥈 fused_rg_warp_opt_base_base 0.18 1.34 1.14
🥉 optimized_fused_relu_groupnorm_base 0.19 1.31 1.11
4 fused_rg_atomic_min_base_base 0.19 1.29 1.10
5 fused_rg_dynamic_bs_opt_base 0.20 1.27 1.08
6 fused_rg_aligned_base_base 0.20 1.22 1.04
7 fused_rg_even_work_base 0.21 1.20 1.02
8 fused_rg_coalesced_base 0.21 1.19 1.02
9 fused_rg_const_base_base 0.21 1.19 1.01
10 fused_rg_const_base_base 0.21 1.16 0.99
10 fused_rg_no_divergence_base 0.21 1.16 0.99
12 fused_rg_modular_base 0.21 1.15 0.98
13 61_ConvTranspose3d_ReLU_GroupNorm_modular_base 0.23 1.09 0.93
14 61_ConvTranspose3d_ReLU_GroupNorm_modular_edit_1 0.23 1.09 0.93
15 61_ConvTranspose3d_ReLU_GroupNorm_grid_stride_base 0.23 1.08 0.92
15 61_ConvTranspose3d_ReLU_GroupNorm_aligned_memory_edit_1 0.23 1.08 0.92
15 61_ConvTranspose3d_ReLU_GroupNorm_modular_base 0.23 1.08 0.92
15 61_ConvTranspose3d_ReLU_GroupNorm_aligned_memory_base 0.23 1.08 0.92
19 61_convtranspose3d_relu_groupnorm_unroll_edit_1 0.23 1.07 0.92
19 61_ConvTranspose3d_ReLU_GroupNorm_grid_stride_edit_1 0.23 1.07 0.92 
#include <pybind11/pybind11.h>
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

#define BLOCK_SIZE 256
#define WARP_SIZE 32
#define MAX_CHANNELS 4096  // Maximum number of channels stored in constant memory

// Declare constant memory arrays for GroupNorm parameters (gamma and beta)
__constant__ float c_gamma[MAX_CHANNELS];
__constant__ float c_beta[MAX_CHANNELS];

// Kernel: Fused ConvTranspose3D, ReLU, and GroupNorm with constant memory for gamma and beta
// Data layout: [N, C, D, H, W]
__global__ void fused_relu_groupnorm_const_kernel(
    float* __restrict__ data,   // input/output tensor
    int N, int C, int D, int H, int W,
    int G, float eps)           // number of groups and epsilon
{
    // Each block processes one (sample, group) pair
    int n = blockIdx.x;        // sample index
    int g = blockIdx.y;        // group index
    int channels_per_group = C / G;
    int c_start = g * channels_per_group;

    int spatial_size = D * H * W;
    int group_elems = channels_per_group * spatial_size;
    int tid = threadIdx.x;

    // Use vectorized loads for data processing (float4 for 128-bit loads)
    int aligned_elements = (group_elems / 4) * 4;

    float4 local_sum4 = make_float4(0.f, 0.f, 0.f, 0.f);
    float local_sum = 0.f;
    float local_sumsq = 0.f;

    // Process aligned elements in chunks of 4
    for (int i = tid * 4; i < aligned_elements; i += blockDim.x * 4) {
        int base_idx = n * (C * spatial_size) + (c_start * spatial_size) + i;
        float4 data4;
        // Load 4 elements
        data4.x = data[base_idx];
        data4.y = data[base_idx + 1];
        data4.z = data[base_idx + 2];
        data4.w = data[base_idx + 3];

        // Apply branchless ReLU
        data4.x = fmaxf(data4.x, 0.f);
        data4.y = fmaxf(data4.y, 0.f);
        data4.z = fmaxf(data4.z, 0.f);
        data4.w = fmaxf(data4.w, 0.f);

        // Write back activated values
        data[base_idx]     = data4.x;
        data[base_idx + 1] = data4.y;
        data[base_idx + 2] = data4.z;
        data[base_idx + 3] = data4.w;

        // Accumulate local sums for mean and variance computation
        local_sum4.x += data4.x;
        local_sum4.y += data4.y;
        local_sum4.z += data4.z;
        local_sum4.w += data4.w;

        local_sumsq += data4.x * data4.x + data4.y * data4.y +
                      data4.z * data4.z + data4.w * data4.w;
    }

    // Process remaining elements
    for (int i = aligned_elements + tid; i < group_elems; i += blockDim.x) {
        int base_idx = n * (C * spatial_size) + (c_start * spatial_size) + i;
        float val = data[base_idx];
        val = fmaxf(val, 0.f);
        data[base_idx] = val;
        local_sum += val;
        local_sumsq += val * val;
    }

    // Combine vectorized accumulations
    local_sum += local_sum4.x + local_sum4.y + local_sum4.z + local_sum4.w;

    // Warp-level reduction using shuffle intrinsics
    unsigned int mask = 0xffffffff;
    int lane = tid & (WARP_SIZE - 1);
    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
        local_sum  += __shfl_down_sync(mask, local_sum, offset);
        local_sumsq += __shfl_down_sync(mask, local_sumsq, offset);
    }

    // Use shared memory to aggregate warp-level results
    __shared__ float s_sum[WARP_SIZE];
    __shared__ float s_sumsq[WARP_SIZE];
    int warp_id = tid / WARP_SIZE;
    if (lane == 0) {
        s_sum[warp_id] = local_sum;
        s_sumsq[warp_id] = local_sumsq;
    }
    __syncthreads();

    float group_sum = 0.f;
    float group_sumsq = 0.f;
    if (tid == 0) {
        int num_warps = (BLOCK_SIZE + WARP_SIZE - 1) / WARP_SIZE;
        for (int i = 0; i < num_warps; i++) {
            group_sum += s_sum[i];
            group_sumsq += s_sumsq[i];
        }
        float group_mean = group_sum / group_elems;
        float var = group_sumsq / group_elems - group_mean * group_mean;
        float inv_std = rsqrtf(var + eps);
        s_sum[0] = group_mean;
        s_sumsq[0] = inv_std;
    }
    __syncthreads();

    float group_mean = s_sum[0];
    float inv_std = s_sumsq[0];

    // Second pass: Normalize the data using computed statistics and constant memory parameters
    for (int i = tid; i < group_elems; i += blockDim.x) {
        int channel_offset = i / spatial_size;
        int c = c_start + channel_offset;
        int base_idx = n * (C * spatial_size) + (c_start * spatial_size) + i;
        float val = data[base_idx];
        val = (val - group_mean) * inv_std;
        // Retrieve gamma and beta from constant memory
        float gamma_val = c_gamma[c];
        float beta_val  = c_beta[c];
        val = val * gamma_val + beta_val;
        data[base_idx] = val;
    }
}

// Forward function: applies ConvTranspose3d then launches the fused CUDA kernel
// Copies GroupNorm parameters to constant memory before kernel launch

torch::Tensor forward(
    torch::Tensor x,
    torch::Tensor conv_transpose,
    torch::Tensor group_norm_weight,
    torch::Tensor group_norm_bias,
    int64_t groups,
    double eps) {

    // Perform 3D transposed convolution using ATen
    auto y = at::conv_transpose3d(
        x,
        conv_transpose,
        /*bias=*/c10::nullopt,
        /*stride=*/{1, 1, 1},
        /*padding=*/{0, 0, 0},
        /*output_padding=*/{0, 0, 0},
        /*groups=*/1,
        /*dilation=*/{1, 1, 1}
    );

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

    // Copy GroupNorm parameters to constant memory (ensure C does not exceed MAX_CHANNELS)
    cudaMemcpyToSymbol(c_gamma, group_norm_weight.data_ptr<float>(), C * sizeof(float));
    cudaMemcpyToSymbol(c_beta, group_norm_bias.data_ptr<float>(), C * sizeof(float));

    dim3 grid(N, G);
    dim3 block(BLOCK_SIZE);

    fused_relu_groupnorm_const_kernel<<<grid, block>>>(
         y.data_ptr<float>(),
         N, C, D, H, W,
         G, static_cast<float>(eps)
    );

    cudaDeviceSynchronize();
    return y;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "Fused ConvTranspose3D + ReLU + GroupNorm with constant memory (CUDA)");
}