← Back to Leaderboard

The AI CUDA Engineer 👷

41_Max_Pooling_1D • balanced_workload_distribution_base

Level 1 • Task 41
import torch
import torch.nn as nn
import torch.nn.functional as F


def module_fn(
    x: torch.Tensor,
    kernel_size: int,
    stride: int,
    padding: int,
    dilation: int,
    return_indices: bool,
) -> torch.Tensor:
    """
    Functional implementation of Max Pooling 1D.

    Args:
        x (torch.Tensor): Input tensor of shape (batch_size, num_features, sequence_length).
        kernel_size (int): Size of the window to take a max over.
        stride (int): Stride of the window.
        padding (int): Implicit zero padding to be added on both sides.
        dilation (int): Spacing between kernel elements.
        return_indices (bool): Whether to return the indices of the maximum values.

    Returns:
        torch.Tensor: Output tensor with Max Pooling 1D applied.
    """
    return F.max_pool1d(
        x,
        kernel_size=kernel_size,
        stride=stride,
        padding=padding,
        dilation=dilation,
        return_indices=return_indices,
    )


class Model(nn.Module):
    """
    Simple model that performs Max Pooling 1D.
    """

    def __init__(
        self,
        kernel_size: int,
        stride: int,
        padding: int,
        dilation: int,
        return_indices: bool,
    ):
        """
        Initializes the Max Pooling 1D layer.

        Args:
            kernel_size (int): Size of the window to take a max over.
            stride (int): Stride of the window.
            padding (int): Implicit zero padding to be added on both sides.
            dilation (int): Spacing between kernel elements.
            return_indices (bool): Whether to return the indices of the maximum values.
        """
        super(Model, self).__init__()
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.return_indices = return_indices

    def forward(self, x: torch.Tensor, fn=module_fn) -> torch.Tensor:
        """
        Applies Max Pooling 1D to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, num_features, sequence_length).
            fn: Function to apply (defaults to module_fn)

        Returns:
            torch.Tensor: Output tensor with Max Pooling 1D applied.
        """
        return fn(
            x,
            self.kernel_size,
            self.stride,
            self.padding,
            self.dilation,
            self.return_indices,
        )


batch_size = 16
features = 64
sequence_length = 128
kernel_size = 4
stride = 2
padding = 2
dilation = 3
return_indices = False


def get_inputs():
    x = torch.randn(batch_size, features, sequence_length)
    return [x]


def get_init_inputs():
    return [kernel_size, stride, padding, dilation, return_indices]

import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs Max Pooling 1D.
    """
    def __init__(self, kernel_size: int, stride: int = None, padding: int = 0, dilation: int = 1, return_indices: bool = False):
        """
        Initializes the Max Pooling 1D layer.

        Args:
            kernel_size (int): Size of the window to take a max over.
            stride (int, optional): Stride of the window. Defaults to None (same as kernel_size).
            padding (int, optional): Implicit zero padding to be added on both sides. Defaults to 0.
            dilation (int, optional): Spacing between kernel elements. Defaults to 1.
            return_indices (bool, optional): Whether to return the indices of the maximum values. Defaults to False.
        """
        super(Model, self).__init__()
        self.maxpool = nn.MaxPool1d(kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, return_indices=return_indices)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Max Pooling 1D to the input tensor.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, num_features, sequence_length).

        Returns:
            torch.Tensor: Output tensor with Max Pooling 1D applied, shape (batch_size, num_features, output_sequence_length).
        """
        return self.maxpool(x)

batch_size = 16
features = 64
sequence_length = 128
kernel_size = 4
stride = 2
padding = 2
dilation = 3
return_indices = False

def get_inputs():
    x = torch.randn(batch_size, features, sequence_length)
    return [x]

def get_init_inputs():
    return [kernel_size, stride, padding, dilation, return_indices]
Download Evaluation Download PyTorch Download CUDA Download Profiles

Kernel Information

Operation Name 41_Max_Pooling_1D
Level ID 1
Task ID 41
Kernel Name balanced_workload_distribution_base
CUDA Speedup (Native) 1.183x
CUDA Speedup (Compile) 5.007x
CUDA Runtime 0.006 ms
PyTorch Runtime (Native) 0.007 ms
PyTorch Runtime (Compile) 0.030 ms
Correct True
Max Diff (vs. Reference) 0.000000
Model azure-gpt-4o-2024-08-06
Temperature 0.00
View Experiment Progress Details

Related Kernels (Level 1, Task 41 • 41_Max_Pooling_1D)

Rank Kernel Name Runtime (ms) Speedup Native Speedup Compile
🥇 41_Max_Pooling_1D 0.01 1.18 5.01
🥇 max_pool1d_shared_opt_base 0.01 1.18 5.01
🥇 optimized_workload_distributed_pool1d_base 0.01 1.18 5.01
🥇 max_pool1d_optimized_grid_base 0.01 1.18 5.01
🥇 optimized_max_pool1d_kernel_base 0.01 1.18 5.01
🥇 max_pool1d_kernel_combined_base 0.01 1.18 5.01
🥇 max_pool1d_nosync_base 0.01 1.18 5.01
🥇 coalesced_writes_edit_1 0.01 1.18 5.01
🥇 aligned_memory_access_base 0.01 1.18 5.01
🥇 aligned_memory_access_edit_1 0.01 1.18 5.01
🥇 loop_unrolling_base 0.01 1.18 5.01
🥇 balanced_workload_distribution_base 0.01 1.18 5.01
🥇 balanced_max_pool1d_base 0.01 1.18 5.01
🥇 balanced_workload_distribution_edit_1 0.01 1.18 5.01
🥇 max_pool1d_fused_kernel_base 0.01 1.18 5.01
🥇 coalesced_max_pool1d_kernel_base_base 0.01 1.18 5.01
🥇 experimental_block_size_pool1d_base_base 0.01 1.18 5.01
🥇 max_pool1d_tunable_base 0.01 1.18 5.01
🥇 coalesced_aligned_pooling_base 0.01 1.18 5.01
🥇 modular_device_functions_edit_1_base 0.01 1.18 5.01 
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

// CUDA kernel with balanced workload distribution
__global__ void max_pool1d_kernel_balanced(
    const float* __restrict__ input,
    float* __restrict__ output,
    int64_t* __restrict__ indices,
    const int batch_size,
    const int num_channels,
    const int input_length,
    const int kernel_size,
    const int stride,
    const int padding,
    const int dilation,
    const int output_length,
    const bool return_indices) {

    const int b = blockIdx.z;
    const int c = blockIdx.y;
    const int i = blockIdx.x * blockDim.x + threadIdx.x;

    if (b >= batch_size || c >= num_channels || i >= output_length) return;

    const int input_start = i * stride - padding;
    float max_val = -INFINITY;
    int max_idx = -1;

    for (int k = 0; k < kernel_size; ++k) {
        const int pos = input_start + k * dilation;
        if (pos >= 0 && pos < input_length) {
            const float val = input[b * num_channels * input_length + c * input_length + pos];
            if (val > max_val) {
                max_val = val;
                max_idx = pos;
            }
        }
    }

    const int out_idx = b * num_channels * output_length + c * output_length + i;
    output[out_idx] = max_val;
    if (return_indices) indices[out_idx] = max_idx;
}

// Host function to launch the CUDA kernel
torch::Tensor forward(
    torch::Tensor x,
    int64_t kernel_size,
    int64_t stride,
    int64_t padding,
    int64_t dilation,
    bool return_indices) {

    TORCH_CHECK(x.dim() == 3, "Input must be 3D");
    TORCH_CHECK(x.is_cuda(), "Input must be on CUDA");
    TORCH_CHECK(x.is_contiguous(), "Input must be contiguous");

    const int batch_size = x.size(0);
    const int num_channels = x.size(1);
    const int input_length = x.size(2);

    const int output_length = ((input_length + 2 * padding - dilation * (kernel_size - 1) - 1) / stride) + 1;
    TORCH_CHECK(output_length > 0, "Output length must be positive");

    auto options = torch::TensorOptions().dtype(x.dtype()).device(x.device());
    auto output = torch::empty({batch_size, num_channels, output_length}, options);
    torch::Tensor indices;
    if (return_indices) {
        indices = torch::empty({batch_size, num_channels, output_length}, 
                                 torch::TensorOptions().dtype(torch::kInt64).device(x.device()));
    }

    const int block_size = 256;
    const dim3 threads(block_size);
    const dim3 blocks((output_length + block_size - 1) / block_size, num_channels, batch_size);

    max_pool1d_kernel_balanced<<<blocks, threads>>>(
        x.data_ptr<float>(),
        output.data_ptr<float>(),
        return_indices ? indices.data_ptr<int64_t>() : nullptr,
        batch_size,
        num_channels,
        input_length,
        kernel_size,
        stride,
        padding,
        dilation,
        output_length,
        return_indices
    );

    return return_indices ? torch::cat({output, indices}, -1) : output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &forward, "MaxPool1D forward with balanced workload distribution (CUDA)");
}
Performance Metrics
Metric Value Unit Variance Samples
Executed Ipc Active 0.924 inst/cycle 0.001 5
Executed Ipc Elapsed 0.340 inst/cycle 0.000 5
Issue Slots Busy 23.934 % 0.660 5
Issued Ipc Active 0.956 inst/cycle 0.001 5
SM Busy 23.934 % 0.660 5
Memory Throughput 140004022121.668 byte/second 563161743537796096.000 5
Mem Busy 10.822 % 0.002 5
Max Bandwidth 7.552 % 0.002 5
L1/TEX Hit Rate 69.070 % 0.001 5
L2 Hit Rate 72.342 % 0.008 5
Mem Pipes Busy 16.128 % 0.015 5
Warp Cycles Per Issued Instruction 30.456 cycle 0.023 5
Warp Cycles Per Executed Instruction 31.550 cycle 0.019 5
Avg. Active Threads Per Warp 31.400 0.000 5
Avg. Not Predicated Off Threads Per Warp 29.920 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 47.776 % 2.801 5
Achieved Active Warps Per SM 30.578 warp 1.147 5
Analysis Rules
Rule Description
WRN HighPipeUtilization All compute pipelines are under-utilized. Either this kernel is very small or it doesn't issue enough warps per scheduler. Check the Launch Statistics and Scheduler Statistics sections for further details.
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.6%) 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 621201.95 μs
Device Time 21.34 μs
Self CPU Time 41.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
aten::_to_copy
CPU Time 621160.00 μs
Device Time 21.34 μs
Self CPU Time 108.89 μ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 620858.06 μs
Device Time 0.00 μs
Self CPU Time 99.82 μ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 613752.93 μs
Device Time 0.00 μs
Self CPU Time 613752.93 μ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 390045.73 μs
Device Time 17061.98 μs
Self CPU Time 390045.73 μs
Self Device Time 17061.98 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
max_pool1d_kernel_balanced(float const*, float*, long*, int, int, int, int, int, int, int, int, bool)
CPU Time 0.00 μs
Device Time 23753.09 μs
Self CPU Time 0.00 μs
Self Device Time 23753.09 μ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 18062.72 μs
Device Time 33891.51 μs
Self CPU Time 18062.72 μs
Self Device Time 33891.51 μ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 60033.52 μs
Device Time 509522.39 μs
Self CPU Time 12010.23 μ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 48025.12 μs
Device Time 509522.39 μs
Self CPU Time 13606.06 μs
Self Device Time 509522.39 μ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 509600.08 μs
Self CPU Time 0.00 μs
Self Device Time 509600.08 μ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 45326 warnings (45279 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/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:12:5 bugprone-easily-swappable-parameters
12 | const int input_length,
| ^~~~~~~~~~~~~~~~~~~~~~~
13 | const int kernel_size,
| ~~~~~~~~~~~~~~~~~~~~~~
14 | const int stride,
| ~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:12:15: note: the first parameter in the range is 'input_length'
12 | const int input_length,
| ^~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:14:15: note: the last parameter in the range is 'stride'
14 | const int stride,
| ^~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:15:5: warning: 3 adjacent parameters of 'max_pool1d_kernel_balanced' of similar type ('const int') are easily swapped by mistake [bugprone-easily-swappable-parameters]
15 | const int padding,
| ^~~~~~~~~~~~~~~~~~
16 | const int dilation,
| ~~~~~~~~~~~~~~~~~~~
17 | const int output_length,
| ~~~~~~~~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:15:15: note: the first parameter in the range is 'padding'
15 | const int padding,
| ^~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:17:15: note: the last parameter in the range is 'output_length'
17 | const int output_length,
| ^~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:20:19: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
20 | const int b = blockIdx.z;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:21:19: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
21 | const int c = blockIdx.y;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:22:19: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
22 | const int i = blockIdx.x * blockDim.x + threadIdx.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:48:19: warning: the parameter 'x' is copied for each invocation but only used as a const reference; consider making it a const reference [performance-unnecessary-value-param]
48 | torch::Tensor x,
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:59:28: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
59 | const int batch_size = x.size(0);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:60:30: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
60 | const int num_channels = x.size(1);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:61:30: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
61 | const int input_length = x.size(2);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:63:31: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
63 | const int output_length = ((input_length + 2 * padding - dilation * (kernel_size - 1) - 1) / stride) + 1;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:85:9: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
85 | kernel_size,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:86:9: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
86 | stride,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:87:9: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
87 | padding,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_41/b2_s3_balanced_workload_distribution/base/base.cu:88:9: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
88 | dilation,
| ^