← Back to Leaderboard

The AI CUDA Engineer 👷

45_Average_Pooling_2Dcombined_avg_pool_base

Level 1 • Task 45
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
) -> torch.Tensor:
    """
    Applies 2D Average Pooling using functional interface.

    Args:
        x (torch.Tensor): Input tensor
        kernel_size (int): Size of pooling window
        stride (int): Stride of pooling operation
        padding (int): Input padding

    Returns:
        torch.Tensor: Output tensor with 2D Average Pooling applied
    """
    return F.avg_pool2d(x, kernel_size=kernel_size, stride=stride, padding=padding)


class Model(nn.Module):
    """
    Simple model that performs 2D Average Pooling.
    """

    def __init__(self, kernel_size: int, stride: int, padding: int):
        """
        Initializes the Average Pooling layer.

        Args:
            kernel_size (int): Size of the pooling window
            stride (int): Stride of the pooling operation
            padding (int): Padding applied to input tensor
        """
        super(Model, self).__init__()
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding

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

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, channels, height, width)
            fn: Function to apply pooling operation, defaults to module_fn

        Returns:
            torch.Tensor: Output tensor with Average Pooling applied
        """
        return fn(
            x,
            self.kernel_size,
            self.stride,
            self.padding,
        )


batch_size = 16
channels = 64
height = 256
width = 256
kernel_size = 3
stride = None  # Defaults to kernel_size
padding = 0


def get_inputs():
    x = torch.randn(batch_size, channels, height, width)
    return [x]


def get_init_inputs():
    return [kernel_size, stride if stride is not None else kernel_size, padding]
import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Simple model that performs 2D Average Pooling.
    """

    def __init__(self, kernel_size: int, stride: int = None, padding: int = 0):
        """
        Initializes the Average Pooling layer.

        Args:
            kernel_size (int): Size of the pooling window.
            stride (int, optional): Stride of the pooling operation. Defaults to None (same as kernel_size).
            padding (int, optional): Padding applied to the input tensor. Defaults to 0.
        """
        super(Model, self).__init__()
        self.avg_pool = nn.AvgPool2d(
            kernel_size=kernel_size, stride=stride, padding=padding
        )

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

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, channels, height, width).

        Returns:
            torch.Tensor: Output tensor with Average Pooling applied.
        """
        return self.avg_pool(x)


batch_size = 16
channels = 64
height = 256
width = 256
kernel_size = 3
stride = None  # Defaults to kernel_size
padding = 0


def get_inputs():
    x = torch.randn(batch_size, channels, height, width)
    return [x]


def get_init_inputs():
    return [kernel_size, stride if stride is not None else kernel_size, padding]

Kernel Information

Related Kernels (Level 1, Task 45 • 45_Average_Pooling_2D)

Rank Kernel Name Runtime (ms) Speedup Native Speedup Compile
🥇 unrolled_avg_pool2d_base 0.11 1.94 3.03
🥇 modular_avg_pool2d_base_base 0.11 1.94 3.03
🥇 grid_stride_manual_unroll_base 0.11 1.94 3.03
🥇 optimized_avg_pool2d_base 0.11 1.94 3.03
🥇 manual_unroll_avg_pool2d_base 0.11 1.94 3.03
🥇 efficient_avg_pool_base 0.11 1.94 3.03
🥇 constant_memory_avg_pool2d_base 0.11 1.94 3.03
🥇 even_workload_avg_pool2d_base 0.11 1.94 3.03
🥇 unrolled_optimized_avg_pool2d_base 0.11 1.94 3.03
10 warp_uniform_flow_base_base 0.11 1.83 2.87
10 optimized_avg_pool2d_base 0.11 1.83 2.87
10 grid_stride_avg_pool2d_base_base 0.11 1.83 2.87
10 warp_uniform_flow_base_edit_1 0.11 1.83 2.87
14 stride_optimized_avg_pool2d_base 0.12 1.82 2.84
14 warp_divergence_avg_pool2d_base 0.12 1.82 2.84
14 stride_loop_avg_pool2d_base 0.12 1.82 2.84
14 grid_unrolled_avg_pool2d_base 0.12 1.82 2.84
18 combined_avg_pool_base 0.12 1.80 2.82
18 spatial_block_optimized_base_base 0.12 1.80 2.82
18 spatial_block_optimized_base_edit_1 0.12 1.80 2.82
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

// Combined 2D average pooling kernel that uses a fast-path when the pooling window is completely in-bound
// and falls back to a safe path when on the borders.

template <typename scalar_t>
__global__ void avg_pool2d_forward_kernel_combined(
    const scalar_t* __restrict__ input,
    scalar_t* __restrict__ output,
    int N,
    int C,
    int H,
    int W,
    int outH,
    int outW,
    int kernel_size,
    int stride,
    int padding
) {
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * C * outH * outW;
    if (index >= total) return;

    // Compute output coordinates
    int w_out = index % outW;
    int h_out = (index / outW) % outH;
    int c = (index / (outW * outH)) % C;
    int n = index / (outW * outH * C);

    // Start indices in input tensor
    int h_start = h_out * stride - padding;
    int w_start = w_out * stride - padding;

    scalar_t sum_val = scalar_t(0);
    
    // Fast path: if the entire pooling window is within bounds, no need for per-element boundary checks.
    bool fast_path = (h_start >= 0) && ((h_start + kernel_size) <= H) &&
                     (w_start >= 0) && ((w_start + kernel_size) <= W);

    if (fast_path) {
        #pragma unroll
        for (int i = 0; i < kernel_size; i++) {
            #pragma unroll
            for (int j = 0; j < kernel_size; j++) {
                int h_in = h_start + i;
                int w_in = w_start + j;
                sum_val += input[((n * C + c) * H + h_in) * W + w_in];
            }
        }
    } else {
        // Safe path for border cases
        for (int i = 0; i < kernel_size; i++) {
            for (int j = 0; j < kernel_size; j++) {
                int h_in = h_start + i;
                int w_in = w_start + j;
                if (h_in >= 0 && h_in < H && w_in >= 0 && w_in < W) {
                    sum_val += input[((n * C + c) * H + h_in) * W + w_in];
                }
            }
        }
    }
    
    // Compute the average and write output
    output[index] = sum_val / static_cast<scalar_t>(kernel_size * kernel_size);
}

// Host function that sets up the kernel launch

torch::Tensor avg_pool2d_forward_combined(
    torch::Tensor x,
    int kernel_size,
    int stride,
    int padding
) {
    TORCH_CHECK(x.dim() == 4, "Input must be a 4D tensor.");
    auto N = x.size(0);
    auto C = x.size(1);
    auto H = x.size(2);
    auto W = x.size(3);

    int outH = (H + 2 * padding - kernel_size) / stride + 1;
    int outW = (W + 2 * padding - kernel_size) / stride + 1;

    auto x_cont = x.contiguous();
    auto options = x.options();
    auto out = torch::empty({N, C, outH, outW}, options);

    const int threads = 256;
    const int blocks = (N * C * outH * outW + threads - 1) / threads;

    AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "avg_pool2d_forward_kernel_combined", ([&] {
        const scalar_t* input_data = x_cont.data_ptr<scalar_t>();
        scalar_t* output_data = out.data_ptr<scalar_t>();

        avg_pool2d_forward_kernel_combined<<<blocks, threads>>>(
            input_data,
            output_data,
            N, C, H, W,
            outH, outW,
            kernel_size, stride, padding
        );
    }));

    cudaError_t err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess, "CUDA Error: ", cudaGetErrorString(err));

    return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &avg_pool2d_forward_combined, "2D Average Pooling forward combined optimized kernel (CUDA)");
}
Performance Metrics
Metric Value Unit Variance Samples
Executed Ipc Active 2.476 inst/cycle 0.000 5
Executed Ipc Elapsed 2.370 inst/cycle 0.000 5
Issue Slots Busy 61.920 % 0.036 5
Issued Ipc Active 2.476 inst/cycle 0.000 5
SM Busy 61.920 % 0.036 5
Memory Throughput 2697438490400.436 byte/second 20138788176125304832.000 5
Mem Busy 45.846 % 0.008 5
Max Bandwidth 80.482 % 0.016 5
L1/TEX Hit Rate 65.590 % 0.000 5
L2 Hit Rate 13.326 % 0.001 5
Mem Pipes Busy 24.466 % 0.004 5
Warp Cycles Per Issued Instruction 22.018 cycle 0.006 5
Warp Cycles Per Executed Instruction 22.034 cycle 0.005 5
Avg. Active Threads Per Warp 32.000 0.000 5
Avg. Not Predicated Off Threads Per Warp 27.610 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 85.056 % 0.000 5
Achieved Active Warps Per SM 54.434 warp 0.000 5
Analysis Rules
Rule Description
INF HighPipeUtilization ALU is the highest-utilized pipeline (41.1%) 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 (85.0%) 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 325913.70 μs
Device Time 27843.04 μs
Self CPU Time 39.64 μ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 325874.05 μs
Device Time 27843.04 μs
Self CPU Time 107.57 μ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 1170864.69 μs
Device Time 53913.49 μs
Self CPU Time 1170864.69 μs
Self Device Time 53913.49 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
void avg_pool2d_forward_kernel_combined<float>(float const*, float*, int, int, int, int, int, int, int, int, int)
CPU Time 0.00 μs
Device Time 803926.48 μs
Self CPU Time 0.00 μs
Self Device Time 803926.48 μ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 737792.79 μs
Device Time 553050.90 μs
Self CPU Time 11796.30 μ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 725997.98 μs
Device Time 553050.90 μs
Self CPU Time 16173.15 μs
Self Device Time 553050.90 μ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 553050.90 μs
Self CPU Time 0.00 μs
Self Device Time 553050.90 μ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
45288 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/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:15:5 bugprone-easily-swappable-parameters
15 | int W,
| ^~~~~~
16 | int outH,
| ~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:15:9: note: the first parameter in the range is 'W'
15 | int W,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:16:9: note: the last parameter in the range is 'outH'
16 | int outH,
| ^~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:17:5: warning: 3 adjacent parameters of 'avg_pool2d_forward_kernel_combined' of similar type ('int') are easily swapped by mistake [bugprone-easily-swappable-parameters]
17 | int outW,
| ^~~~~~~~~
18 | int kernel_size,
| ~~~~~~~~~~~~~~~~
19 | int stride,
| ~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:17:9: note: the first parameter in the range is 'outW'
17 | int outW,
| ^~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:19:9: note: the last parameter in the range is 'stride'
19 | int stride,
| ^~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:22:17: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
22 | int index = blockIdx.x * blockDim.x + threadIdx.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:83:16: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
83 | int outH = (H + 2 * padding - kernel_size) / stride + 1;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:83:21: warning: performing an implicit widening conversion to type 'int64_t' (aka 'long') of a multiplication performed in type 'int' [bugprone-implicit-widening-of-multiplication-result]
83 | int outH = (H + 2 * padding - kernel_size) / stride + 1;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:83:21: note: make conversion explicit to silence this warning
4 | int outH = (H + 2 * padding - kernel_size) / stride + 1;
| ^~~~~~~~~~~
| static_cast<int64_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:83:21: note: perform multiplication in a wider type
83 | int outH = (H + 2 * padding - kernel_size) / stride + 1;
| ^
| static_cast<int64_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:84:16: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
84 | int outW = (W + 2 * padding - kernel_size) / stride + 1;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:84:21: warning: performing an implicit widening conversion to type 'int64_t' (aka 'long') of a multiplication performed in type 'int' [bugprone-implicit-widening-of-multiplication-result]
84 | int outW = (W + 2 * padding - kernel_size) / stride + 1;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:84:21: note: make conversion explicit to silence this warning
84 | int outW = (W + 2 * padding - kernel_size) / stride + 1;
| ^~~~~~~~~~~
| static_cast<int64_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:84:21: note: perform multiplication in a wider type
84 | int outW = (W + 2 * padding - kernel_size) / stride + 1;
| ^
| static_cast<int64_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:91:24: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
91 | const int blocks = (N * C * outH * outW + threads - 1) / threads;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250207_optimize_b5_s4_e1_sweep/level_1/task_45/b4_s1_combined_avg_pool/base/base.cu:93:5: warning: inside a lambda, '__func__' expands to the name of the function call operator; consider capturing the name of the enclosing function explicitly [bugprone-lambda-function-name]
93 | AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "avg_pool2d_forward_kernel_combined", ([&] {
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/ATen/Dispatch.h:237:34: note: expanded from macro 'AT_DISPATCH_FLOATING_TYPES'
237 | AT_DISPATCH_SWITCH(TYPE, NAME, AT_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/ATen/Dispatch.h:233:3: note: expanded from macro 'AT_DISPATCH_CASE_FLOATING_TYPES'
233 | AT_DISPATCH_CASE(at::ScalarType::Double, __VA_ARGS__) \
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/ATen/Dispatch.h:74:3: note: expanded from macro 'AT_DISPATCH_CASE'
74 | AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, scalar_t, __VA_ARGS__)
| ^
note: (skipping 1 expansions in backtrace; use -fmacro-backtrace-limit=0 to see all)
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/ATen/Dispatch.h:58:7: note: expanded from macro 'AT_PRIVATE_CHECK_SELECTIVE_BUILD'
58 | AT_ERROR( \
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/c10/util/Exception.h:711:32: note: expanded from macro 'AT_ERROR'
711 | C10_EXPAND_MSVC_WORKAROUND(TORCH_CHECK(false, ::c10::str(__VA_ARGS__))); \
| ^
/home/robert_sakana_ai/miniconda3/envs/llm2cuda/lib/python3.11/site-packages/torch/include/c10/util/Exception.h:536:9: note: expanded from macro 'TORCH_CHECK'
536 | __func__, \
| ^