← Back to Leaderboard

The AI CUDA Engineer 👷

41_Max_Pooling_1Doptimized_workload_distributed_pool1d_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]

Kernel Information

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>

__global__ void optimized_workload_distributed_pool1d_kernel(
    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,
    bool return_indices)
{
    const int tid = threadIdx.x + blockIdx.x * blockDim.x;
    const int total_outputs = output_length * num_channels * batch_size;

    if (tid >= total_outputs) return;

    const int i = tid % output_length;
    const int c = (tid / output_length) % num_channels;
    const int b = tid / (output_length * num_channels);

    const int input_start = i * stride - padding;
    const int base_idx = b * num_channels * input_length + c * input_length;
    float max_val = -INFINITY;
    int max_idx = -1;

    // Loop unrolling to improve performance
    #pragma unroll
    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[base_idx + pos];
            if (val > max_val) {
                max_val = val;
                max_idx = pos;
            }
        }
    }

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

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}, 
            options.dtype(torch::kInt64));
    }

    const int total_elements = batch_size * num_channels * output_length;
    const int threads_per_block = 512;  // Optimized to better utilize GPU capabilities
    const int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block;

    optimized_workload_distributed_pool1d_kernel<<<num_blocks, threads_per_block>>>(
        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 (CUDA)");
}
Performance Metrics
Metric Value Unit Variance Samples
Executed Ipc Active 1.034 inst/cycle 0.001 5
Executed Ipc Elapsed 0.402 inst/cycle 0.000 5
Issue Slots Busy 26.880 % 0.773 5
Issued Ipc Active 1.074 inst/cycle 0.001 5
SM Busy 26.880 % 0.773 5
Memory Throughput 139982328045.906 byte/second 24696796611886718976.000 5
Mem Busy 10.594 % 0.139 5
Max Bandwidth 6.932 % 0.056 5
L1/TEX Hit Rate 69.000 % 0.000 5
L2 Hit Rate 72.530 % 0.007 5
Mem Pipes Busy 7.070 % 0.056 5
Warp Cycles Per Issued Instruction 14.566 cycle 0.927 5
Warp Cycles Per Executed Instruction 15.164 cycle 1.009 5
Avg. Active Threads Per Warp 32.000 0.000 5
Avg. Not Predicated Off Threads Per Warp 28.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 4.000 block 0.000 5
Block Limit Shared Mem 16.000 block 0.000 5
Block Limit Warps 4.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 23.602 % 0.003 5
Achieved Active Warps Per SM 15.104 warp 0.001 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.
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 (23.5%) 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 254756.23 μs
Device Time 21.34 μs
Self CPU Time 37.39 μ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 254718.84 μs
Device Time 21.34 μs
Self CPU Time 86.56 μ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 254434.40 μs
Device Time 0.00 μs
Self CPU Time 89.63 μ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 254046.17 μs
Device Time 0.00 μs
Self CPU Time 254046.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
cudaLaunchKernel
CPU Time 462228.28 μs
Device Time 21155.12 μs
Self CPU Time 462228.28 μs
Self Device Time 21155.12 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
optimized_workload_distributed_pool1d_kernel(float const*, float*, long*, int, int, int, int, int, int, int, int, bool)
CPU Time 0.00 μs
Device Time 30294.25 μs
Self CPU Time 0.00 μs
Self Device Time 30294.25 μ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 24421.44 μs
Device Time 41993.32 μs
Self CPU Time 24421.44 μs
Self Device Time 41993.32 μ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 73032.97 μs
Device Time 628118.07 μs
Self CPU Time 15796.42 μ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 57238.42 μs
Device Time 628118.07 μs
Self CPU Time 19580.82 μs
Self Device Time 628118.07 μ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 628118.07 μs
Self CPU Time 0.00 μs
Self Device Time 628118.07 μ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 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/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:11:5 bugprone-easily-swappable-parameters
11 | const int input_length,
| ^~~~~~~~~~~~~~~~~~~~~~~
12 | const int kernel_size,
| ~~~~~~~~~~~~~~~~~~~~~~
13 | const int stride,
| ~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:11:15: note: the first parameter in the range is 'input_length'
11 | const int input_length,
| ^~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:13:15: note: the last parameter in the range is 'stride'
13 | const int stride,
| ^~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:14:5: warning: 3 adjacent parameters of 'optimized_workload_distributed_pool1d_kernel' of similar type ('const int') are easily swapped by mistake [bugprone-easily-swappable-parameters]
14 | const int padding,
| ^~~~~~~~~~~~~~~~~~
15 | const int dilation,
| ~~~~~~~~~~~~~~~~~~~
16 | const int output_length,
| ~~~~~~~~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:14:15: note: the first parameter in the range is 'padding'
14 | const int padding,
| ^~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:16:15: note: the last parameter in the range is 'output_length'
16 | const int output_length,
| ^~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:19:21: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
19 | const int tid = threadIdx.x + blockIdx.x * blockDim.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:52: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]
52 | torch::Tensor x,
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:63:28: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
63 | const int batch_size = x.size(0);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:64:30: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
64 | const int num_channels = x.size(1);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:65:30: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
65 | const int input_length = x.size(2);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:67:31: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
67 | const int output_length = ((input_length + 2 * padding - dilation * (kernel_size - 1) - 1) / stride) + 1;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:90:9: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
90 | kernel_size,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:91:9: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
91 | stride,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:92:9: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
92 | padding,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250201_optimize_b10_s4_e0_sweep/level_1/task_41/b5_s2_optimized_workload_distributed_pool1d/base/base.cu:93:9: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
93 | dilation,
| ^