← Back to Leaderboard

The AI CUDA Engineer 👷

54_Conv2d_Multiply_LeakyReLU_GELUbalanced_thread_distribution_base

Level 2 • Task 54
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,
    multiplier: torch.Tensor,
) -> torch.Tensor:
    """
    Applies convolution, scalar multiplication, LeakyReLU and GELU.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width)
        conv_weight (torch.Tensor): Convolution weights of shape (out_channels, in_channels, kernel_size, kernel_size)
        conv_bias (torch.Tensor): Convolution bias of shape (out_channels)
        multiplier (torch.Tensor): Learnable scalar of shape (out_channels, 1, 1)

    Returns:
        torch.Tensor: Output tensor after applying convolution, multiplication, LeakyReLU and GELU
    """
    x = F.conv2d(x, conv_weight, bias=conv_bias)
    x = x * multiplier
    x = F.leaky_relu(x)
    x = F.gelu(x)
    return x


class Model(nn.Module):
    """
    Model that performs a convolution, multiplies by a learnable scalar, applies LeakyReLU, and then GELU.
    """

    def __init__(self, in_channels, out_channels, kernel_size, multiplier_shape):
        super(Model, self).__init__()
        conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.conv_weight = nn.Parameter(conv.weight)
        self.conv_bias = nn.Parameter(conv.bias)
        self.multiplier = nn.Parameter(torch.randn(multiplier_shape) * 0.02)

    def forward(self, x, fn=module_fn):
        return fn(x, self.conv_weight, self.conv_bias, self.multiplier)


batch_size = 128
in_channels = 3
out_channels = 16
height, width = 32, 32
kernel_size = 3
multiplier_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, multiplier_shape]
import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a convolution, multiplies by a learnable scalar, applies LeakyReLU, and then GELU.
    """
    def __init__(self, in_channels, out_channels, kernel_size, multiplier_shape):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.multiplier = nn.Parameter(torch.randn(multiplier_shape) * 0.02) 
        self.leaky_relu = nn.LeakyReLU()

    def forward(self, x):
        x = self.conv(x)
        x = x * self.multiplier
        x = self.leaky_relu(x)
        x = torch.nn.functional.gelu(x)
        return x

batch_size = 128
in_channels = 3
out_channels = 16
height, width = 32, 32
kernel_size = 3
multiplier_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, multiplier_shape]

Kernel Information

Related Kernels (Level 2, Task 54 • 54_Conv2d_Multiply_LeakyReLU_GELU)

Rank Kernel Name Runtime (ms) Speedup Native Speedup Compile
🥇 54_Conv2d_Multiply_LeakyReLU_GELU 0.04 1.28 1.44
🥇 balanced_workload_distribution_base 0.04 1.28 1.44
🥇 warp_divergence_optimized_base 0.04 1.28 1.44
🥇 optimized_block_size_128_base 0.04 1.28 1.44
🥇 optimized_convolution_with_tunable_blocksize_base 0.04 1.28 1.44
🥇 direct_3d_indexing_opt_base 0.04 1.28 1.44
🥇 direct_3d_indexing_base 0.04 1.28 1.44
🥇 unroll_loops_54conv_edit_1 0.04 1.28 1.44
🥇 dynamic_block_size_54conv_base 0.04 1.28 1.44
🥇 threadblock_3d_mapping_base 0.04 1.28 1.44
🥇 balanced_thread_distribution_base 0.04 1.28 1.44
🥇 branchless_no_divergence_54conv_base 0.04 1.28 1.44
🥇 modular_device_functions_base 0.04 1.28 1.44
🥇 tile_based_2d_indexing_base 0.04 1.28 1.44
15 combined_conv_act_base 0.04 1.25 1.40
15 optimized_stride_loop_base 0.04 1.25 1.40
15 unroll_loops_54conv_base 0.04 1.25 1.40
15 54_Conv2d_Multiply_LeakyReLU_GELU_warp_divergence_reduction_base 0.04 1.25 1.40
15 dynamic_block_size_54conv_edit_1 0.04 1.25 1.40
15 modular_device_functions_edit_1 0.04 1.25 1.40
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#include <cmath>
#include <stdio.h>

// Device function: GELU approximation
__device__ __forceinline__ float gelu(float x) {
    const float k0 = 0.7978845608028654f; // sqrt(2/pi)
    return 0.5f * x * (1.0f + tanhf(k0 * (x + 0.044715f * x * x * x)));
}

// CUDA kernel that performs convolution, scalar multiplication, LeakyReLU and GELU.
// - input: [batch_size, in_channels, input_h, input_w]
// - weight: [out_channels, in_channels, kernel_size, kernel_size]
// - bias: [out_channels]
// - multiplier: [out_channels] (broadcast over spatial dims)
// - output: [batch_size, out_channels, output_h, output_w]
__global__ void conv_forward_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    const float* __restrict__ multiplier,
    float* __restrict__ output,
    int batch_size,
    int in_channels,
    int input_h,
    int input_w,
    int out_channels,
    int kernel_size,
    int output_h,
    int output_w
) {
    // Calculate the linear index
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * out_channels * output_h * output_w;
    if (index >= total) return;

    // Calculate the indices for output: (n, oc, oh, ow)
    int ow = index % output_w;
    int tmp = index / output_w;
    int oh = tmp % output_h;
    tmp = tmp / output_h;
    int oc = tmp % out_channels;
    int n = tmp / out_channels;

    // Start with the bias for output channel oc.
    float sum = bias[oc];
    
    // Convolution: iterate over input channels and kernel window.
    for (int ic = 0; ic < in_channels; ic++) {
        for (int i = 0; i < kernel_size; i++) {
            for (int j = 0; j < kernel_size; j++) {
                int in_h = oh + i; // stride = 1, no padding
                int in_w = ow + j;
                int input_index = ((n * in_channels + ic) * input_h + in_h) * input_w + in_w;
                int weight_index = ((oc * in_channels + ic) * kernel_size + i) * kernel_size + j;
                sum += input[input_index] * weight[weight_index];
            }
        }
    }
    
    // Multiply with the channel-specific multiplier
    sum *= multiplier[oc];
    
    // Apply LeakyReLU activation (negative slope = 0.01)
    sum = (sum > 0.0f) ? sum : 0.01f * sum;
    
    // Apply GELU activation
    float out_val = gelu(sum);
    
    output[index] = out_val;
}

// C++ interface (to be called from Python)
torch::Tensor forward_cuda(
    torch::Tensor input,
    torch::Tensor conv_weight,
    torch::Tensor conv_bias,
    torch::Tensor multiplier
) {
    // Get input dimensions
    const int batch_size = input.size(0);
    const int in_channels = input.size(1);
    const int input_h = input.size(2);
    const int input_w = input.size(3);
    
    // Get convolution parameters
    const int out_channels = conv_weight.size(0);
    const int kernel_size = conv_weight.size(2);
    const int output_h = input_h - kernel_size + 1;
    const int output_w = input_w - kernel_size + 1;
    
    // Allocate output tensor
    auto output = torch::empty({batch_size, out_channels, output_h, output_w}, input.options());
    
    // Calculate the total number of elements to process
    const int total_elements = batch_size * out_channels * output_h * output_w;
    const int threads = 256;
    const int blocks = (total_elements + threads - 1) / threads;
    
    // Launch CUDA kernel
    conv_forward_kernel<<<blocks, threads>>>(
        input.data_ptr<float>(),
        conv_weight.data_ptr<float>(),
        conv_bias.data_ptr<float>(),
        multiplier.data_ptr<float>(),
        output.data_ptr<float>(),
        batch_size,
        in_channels,
        input_h,
        input_w,
        out_channels,
        kernel_size,
        output_h,
        output_w
    );
    
    // Check for kernel launch errors
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        printf("CUDA kernel failed: %s\n", cudaGetErrorString(err));
    }
    
    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward_cuda, "Convolution, scalar multiplication, LeakyReLU and GELU (CUDA) with balanced thread distribution");
}
Performance Metrics
Metric Value Unit Variance Samples
Executed Ipc Active 3.190 inst/cycle 0.000 5
Executed Ipc Elapsed 2.962 inst/cycle 0.000 5
Issue Slots Busy 79.824 % 0.032 5
Issued Ipc Active 3.196 inst/cycle 0.000 5
SM Busy 79.824 % 0.032 5
Memory Throughput 36095013854.068 byte/second 15006783629789324.000 5
Mem Busy 54.180 % 0.033 5
Max Bandwidth 36.992 % 0.016 5
L1/TEX Hit Rate 87.946 % 0.000 5
L2 Hit Rate 92.540 % 0.145 5
Mem Pipes Busy 47.950 % 0.026 5
Warp Cycles Per Issued Instruction 16.618 cycle 0.000 5
Warp Cycles Per Executed Instruction 16.630 cycle 0.000 5
Avg. Active Threads Per Warp 32.000 0.000 5
Avg. Not Predicated Off Threads Per Warp 29.340 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 32.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 83.198 % 0.003 5
Achieved Active Warps Per SM 53.246 warp 0.001 5
Analysis Rules
Rule Description
INF HighPipeUtilization ALU is the highest-utilized pipeline (47.5%) 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 (83.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::to
CPU Time 877998.75 μs
Device Time 82.08 μs
Self CPU Time 68.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::_to_copy
CPU Time 877930.16 μs
Device Time 82.08 μs
Self CPU Time 145.11 μ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::empty_strided
CPU Time 877382.75 μs
Device Time 0.00 μs
Self CPU Time 163.44 μ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
cudaDeviceGetStreamPriorityRange
CPU Time 868741.95 μs
Device Time 0.00 μs
Self CPU Time 868741.95 μ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
cudaLaunchKernel
CPU Time 154971.57 μs
Device Time 862.27 μs
Self CPU Time 154971.57 μs
Self Device Time 862.27 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
conv_forward_kernel(float const*, float const*, float const*, float const*, float*, int, int, int, int, int, int, int, int)
CPU Time 0.00 μs
Device Time 14698.10 μs
Self CPU Time 0.00 μs
Self Device Time 14698.10 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
cudaEventRecord
CPU Time 1115.36 μs
Device Time 1564.31 μs
Self CPU Time 1115.36 μs
Self Device Time 1564.31 μ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 127150.65 μs
Device Time 32015.68 μs
Self CPU Time 748.96 μ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::fill_
CPU Time 126402.73 μs
Device Time 32015.68 μs
Self CPU Time 905.64 μs
Self Device Time 32015.68 μ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 32015.68 μs
Self CPU Time 0.00 μs
Self Device Time 32015.68 μ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
45291 warnings generated when compiling for host.
Suppressed 45324 warnings (45277 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/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:22:5 bugprone-easily-swappable-parameters
22 | const float* __restrict__ weight,
| ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
23 | const float* __restrict__ bias,
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
24 | const float* __restrict__ multiplier,
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:22:31: note: the first parameter in the range is 'weight'
22 | const float* __restrict__ weight,
| ^~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:24:31: note: the last parameter in the range is 'multiplier'
24 | const float* __restrict__ multiplier,
| ^~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:26:5: warning: 2 adjacent parameters of 'conv_forward_kernel' of similar type ('int') are easily swapped by mistake [bugprone-easily-swappable-parameters]
26 | int batch_size,
| ^~~~~~~~~~~~~~~
27 | int in_channels,
| ~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:26:9: note: the first parameter in the range is 'batch_size'
26 | int batch_size,
| ^~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:27:9: note: the last parameter in the range is 'in_channels'
27 | int in_channels,
| ^~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:29:5: warning: 3 adjacent parameters of 'conv_forward_kernel' of similar type ('int') are easily swapped by mistake [bugprone-easily-swappable-parameters]
29 | int input_w,
| ^~~~~~~~~~~~
30 | int out_channels,
| ~~~~~~~~~~~~~~~~~
31 | int kernel_size,
| ~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:29:9: note: the first parameter in the range is 'input_w'
29 | int input_w,
| ^~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:31:9: note: the last parameter in the range is 'kernel_size'
31 | int kernel_size,
| ^~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:36:17: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
36 | int index = blockIdx.x * blockDim.x + threadIdx.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:78:19: warning: the parameter 'input' is copied for each invocation but only used as a const reference; consider making it a const reference [performance-unnecessary-value-param]
78 | torch::Tensor input,
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:79:19: warning: the parameter 'conv_weight' is copied for each invocation but only used as a const reference; consider making it a const reference [performance-unnecessary-value-param]
79 | torch::Tensor conv_weight,
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:80:19: warning: the parameter 'conv_bias' is copied for each invocation but only used as a const reference; consider making it a const reference [performance-unnecessary-value-param]
80 | torch::Tensor conv_bias,
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:81:19: warning: the parameter 'multiplier' is copied for each invocation but only used as a const reference; consider making it a const reference [performance-unnecessary-value-param]
81 | torch::Tensor multiplier
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:84:28: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
84 | const int batch_size = input.size(0);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:85:29: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
85 | const int in_channels = input.size(1);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:86:25: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
86 | const int input_h = input.size(2);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:87:25: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
87 | const int input_w = input.size(3);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:90:30: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
90 | const int out_channels = conv_weight.size(0);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250212_optimize_b5_s4_e1_v2/level_2/task_54/b5_s0_balanced_thread_distribution/base/base.cu:91:29: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
91 | const int kernel_size = conv_weight.size(2);
| ^