← Back to Leaderboard

The AI CUDA Engineer 👷

54_Conv2d_Multiply_LeakyReLU_GELUtile_based_2d_indexing_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>

// Define tile dimensions for 2D block indexing
#define TILE_W 16
#define TILE_H 16

// Device function: GELU approximation using tanhf
__device__ inline 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 using tile-based 2D indexing
// The grid is organized as follows:
//   gridDim.x covers tiles across the output width,
//   gridDim.y covers tiles across the output height,
//   gridDim.z covers batch_size * out_channels.
// Each thread computes one output element corresponding to (n, oc, oh, ow).
__global__ void conv_forward_kernel_tile(
    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
) {
    // Compute output spatial indices within the tile
    int ow = blockIdx.x * TILE_W + threadIdx.x;
    int oh = blockIdx.y * TILE_H + threadIdx.y;

    // Determine batch and output channel from the third grid dimension
    int bc_index = blockIdx.z;
    int n = bc_index / out_channels;
    int oc = bc_index % out_channels;

    // Process only if within the output bounds
    if (ow < output_w && oh < output_h) {
        float sum = bias[oc];

        // Perform the convolution over input channels and kernel window
        for (int ic = 0; ic < in_channels; ++ic) {
            for (int i = 0; i < kernel_size; ++i) {
                #pragma unroll
                for (int j = 0; j < kernel_size; ++j) {
                    int in_h = oh + i;  // stride = 1, no padding assumed
                    int in_w = ow + j;
                    int input_idx = ((n * in_channels + ic) * input_h + in_h) * input_w + in_w;
                    int weight_idx = ((oc * in_channels + ic) * kernel_size + i) * kernel_size + j;
                    sum += input[input_idx] * weight[weight_idx];
                }
            }
        }

        // Apply channel-specific multiplier
        sum *= multiplier[oc];

        // Apply branchless LeakyReLU activation (negative slope = 0.01)
        sum = fmaxf(sum, 0.01f * sum);

        // Apply GELU activation
        float out_val = gelu(sum);

        // Compute the linear index for the output tensor
        int output_idx = ((n * out_channels + oc) * output_h + oh) * output_w + ow;
        output[output_idx] = out_val;
    }
}

// C++ interface (to be called from Python via pybind11)
// This function extracts tensor dimensions, allocates the output tensor, and launches the CUDA kernel
torch::Tensor forward_cuda(
    torch::Tensor input,
    torch::Tensor conv_weight,
    torch::Tensor conv_bias,
    torch::Tensor multiplier
) {
    // Get input dimensions (batch_size, in_channels, input_h, input_w)
    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);

    // Convolution parameters (out_channels, kernel dimensions)
    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());

    // Set up block dimensions (2D tile for output spatial dimensions)
    dim3 block(TILE_W, TILE_H, 1);
    
    // Grid dimensions:
    //   gridDim.x: number of tiles covering the output width
    //   gridDim.y: number of tiles covering the output height
    //   gridDim.z: covers each (batch, out_channel) pair
    dim3 grid((output_w + TILE_W - 1) / TILE_W,
              (output_h + TILE_H - 1) / TILE_H,
              batch_size * out_channels);

    // Launch the CUDA kernel
    conv_forward_kernel_tile<<<grid, block>>>(
        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 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, "Tile-based 2D index optimized convolution, scalar multiplication, LeakyReLU, and GELU (CUDA)");
}
Performance Metrics
Metric Value Unit Variance Samples
Executed Ipc Active 3.034 inst/cycle 0.000 5
Executed Ipc Elapsed 2.822 inst/cycle 0.000 5
Issue Slots Busy 75.924 % 0.017 5
Issued Ipc Active 3.036 inst/cycle 0.000 5
SM Busy 75.924 % 0.017 5
Memory Throughput 36986577817.254 byte/second 2392144318895155.500 5
Mem Busy 61.782 % 0.009 5
Max Bandwidth 42.142 % 0.005 5
L1/TEX Hit Rate 87.244 % 0.003 5
L2 Hit Rate 92.162 % 0.012 5
Mem Pipes Busy 55.776 % 0.009 5
Warp Cycles Per Issued Instruction 16.348 cycle 0.001 5
Warp Cycles Per Executed Instruction 16.354 cycle 0.001 5
Avg. Active Threads Per Warp 30.170 0.000 5
Avg. Not Predicated Off Threads Per Warp 28.040 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 77.838 % 0.025 5
Achieved Active Warps Per SM 49.818 warp 0.011 5
Analysis Rules
Rule Description
INF HighPipeUtilization ALU is the highest-utilized pipeline (38.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.
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 (78.1%) 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 330484.71 μs
Device Time 83.61 μs
Self CPU Time 48.81 μ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 330435.89 μs
Device Time 83.61 μs
Self CPU Time 93.05 μ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 329995.99 μs
Device Time 0.00 μs
Self CPU Time 89.97 μ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 329532.91 μs
Device Time 0.00 μs
Self CPU Time 329532.91 μ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 735575.00 μs
Device Time 704.76 μs
Self CPU Time 735575.00 μs
Self Device Time 704.76 μ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_tile(float const*, float const*, float const*, float const*, float*, int, int, int, int, int, int, int, int)
CPU Time 0.00 μs
Device Time 289017.05 μs
Self CPU Time 0.00 μs
Self Device Time 289017.05 μ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 246759.38 μs
Device Time 632701.32 μs
Self CPU Time 13006.79 μ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 233753.30 μs
Device Time 632701.32 μs
Self CPU Time 17928.71 μs
Self Device Time 632701.32 μ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 632701.32 μs
Self CPU Time 0.00 μs
Self Device Time 632701.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
45293 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/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:26:5 bugprone-easily-swappable-parameters
26 | const float* __restrict__ weight,
| ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
27 | const float* __restrict__ bias,
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
28 | const float* __restrict__ multiplier,
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:26:31: note: the first parameter in the range is 'weight'
26 | const float* __restrict__ weight,
| ^~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:28:31: note: the last parameter in the range is 'multiplier'
28 | const float* __restrict__ multiplier,
| ^~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:30:5: warning: 2 adjacent parameters of 'conv_forward_kernel_tile' of similar type ('int') are easily swapped by mistake [bugprone-easily-swappable-parameters]
30 | int batch_size,
| ^~~~~~~~~~~~~~~
31 | int in_channels,
| ~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:30:9: note: the first parameter in the range is 'batch_size'
30 | int batch_size,
| ^~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:31:9: note: the last parameter in the range is 'in_channels'
31 | int in_channels,
| ^~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:33:5: warning: 3 adjacent parameters of 'conv_forward_kernel_tile' of similar type ('int') are easily swapped by mistake [bugprone-easily-swappable-parameters]
33 | int input_w,
| ^~~~~~~~~~~~
34 | int out_channels,
| ~~~~~~~~~~~~~~~~~
35 | int kernel_size,
| ~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:33:9: note: the first parameter in the range is 'input_w'
33 | int input_w,
| ^~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:35:9: note: the last parameter in the range is 'kernel_size'
35 | int kernel_size,
| ^~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:40:14: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
40 | int ow = blockIdx.x * TILE_W + threadIdx.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:41:14: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
41 | int oh = blockIdx.y * TILE_H + threadIdx.y;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:44:20: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
44 | int bc_index = blockIdx.z;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:84: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]
84 | torch::Tensor input,
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:85: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]
85 | torch::Tensor conv_weight,
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:86: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]
86 | torch::Tensor conv_bias,
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:87: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]
87 | torch::Tensor multiplier
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:90:28: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
90 | const int batch_size = input.size(0);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/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 in_channels = input.size(1);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:92:25: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
92 | const int input_h = input.size(2);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:93:25: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
93 | const int input_w = input.size(3);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:96:30: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
96 | const int out_channels = conv_weight.size(0);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_54/b7_s3_tile_based_2d_indexing/base/base.cu:97:29: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
97 | const int kernel_size = conv_weight.size(2);
| ^