[CK_TILE] Add pooling to ckTileEngine by aledudek · Pull Request #3416 · ROCm/composable_kernel

aledudek · 2025-12-12T10:36:24Z

Proposed changes

Add pooling to ckTileEngine

Checklist

Please put an x into the boxes that apply. You can also fill these out after creating the PR. If you're not sure, please don't hesitate to ask.

I have added tests relevant to the introduced functionality, and the unit tests are passing locally
I have added the test to REGRESSION_TESTS list defined at the top of CMakeLists.txt in tests/CMakeLists.txt, IF the test takes more than 30 seconds to run.
I have added inline documentation which enables the maintainers with understanding the motivation
I have removed the stale documentation which is no longer relevant after this pull request
(If this change is user-facing) I have added release notes which provide the end users with a brief summary of the improvement from this pull request
I have run clang-format on all changed files
Any dependent changes have been merged

…ions

… part3

Copilot

Pull request overview

This PR adds comprehensive pooling operations support to the ckTileEngine, implementing 2D and 3D pooling with configurable reduction operations (max, min, average). The implementation follows the existing GEMM module architecture with individual kernel compilation for improved build parallelism and testing capabilities.

Key Changes:

Added pooling kernel generation system with Python-based instance builder
Implemented benchmark infrastructure for pooling operations
Created configuration system for kernel parameters and GPU target selection

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
tile_engine/ops/pooling/pool_profiler.hpp	Profiler class for benchmarking pooling kernels with verification and performance metrics
tile_engine/ops/pooling/pool_instance_builder.py	Python script to generate C++ kernel instances from JSON configurations
tile_engine/ops/pooling/pool_common.hpp	Common type traits and data type definitions for pooling operations
tile_engine/ops/pooling/pool_benchmark_single.cpp	Single kernel benchmark executable with argument parsing and verification
tile_engine/ops/pooling/pool_benchmark.py	Python orchestration script for running and analyzing benchmark sweeps
tile_engine/ops/pooling/pool_benchmark.hpp	Core data structures and enums for pooling benchmarks
tile_engine/ops/pooling/configs/default_config.json	Default configuration for block sizes, warp configs, and traits
tile_engine/ops/pooling/README.md	Comprehensive documentation for building, running, and configuring pooling benchmarks
tile_engine/ops/pooling/CMakeLists.txt	CMake build system for generating and compiling individual pooling kernels
tile_engine/ops/gemm/README.md	Added missing gpu_target argument to example command
tile_engine/ops/commons/test_benchmark.sh	Extended test script to support both GEMM and pooling benchmarks
tile_engine/ops/CMakeLists.txt	Added pooling subdirectory to build system

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

tile_engine/ops/pooling/pool_instance_builder.py

tile_engine/ops/pooling/pool_profiler.hpp

Copilot · 2025-12-18T10:45:21Z

tile_engine/ops/pooling/pool_benchmark.py

+            "Sz": stride_z,
+            "Sy": stride_y,
+            "Sx": stride_x,
+            "pool_dim": pool_dim,


The parameter 'pool_dim' is included in the params dictionary but is not a valid command-line parameter for the benchmark executables according to pool_benchmark_single.cpp (which doesn't define a 'pool_dim' argument). This will be silently ignored or could cause errors. The pool dimension is determined by the kernel itself (POOL_DIM constant), not passed as an argument.

Suggested change

"pool_dim": pool_dim,

tile_engine/ops/pooling/CMakeLists.txt

tile_engine/ops/pooling/pool_benchmark.py

tile_engine/ops/pooling/README.md

msaffari-amd · 2026-01-07T10:51:03Z

tile_engine/ops/gemm/README.md

    --config_json configs/user_provided_config.json \
-    --gen_all_individual
+    --gen_all_individual \
+    --gpu_target gfx942


Gemm TileEngine target is just gfx942?
And this change is related anyway to the adding pooling to TileEngine?

Great comment. TileEngine supports on gfx90a and gfx1_ series.

msaffari-amd · 2026-01-07T11:48:13Z

tile_engine/ops/pooling/configs/default_config.json

+        "block_n": {
+            "values": [1]


If block_n is always 1 for pooling, does it make sense to hard code it in instance_builder or can theoretically select other values?

Block_n could also be 2 in pooling.

msaffari-amd · 2026-01-07T13:27:22Z

tile_engine/ops/pooling/pool_instance_builder.py

I have a general question about the instance builder. In the GEMM tile engine, there’s a gemm_validation_utils.py module that is used by gemm_instance_builder to generate only valid instances that can later be compiled. However, it seems that for pooling, there isn’t a similar validation step. and, some checks are performed directly within pool_instance_builder. Would it make sense to adopt a similar structure for these tile engine operations to keep things same?

ammallya · 2026-02-03T22:05:40Z

Error importing due to merge conflicts – please reopen the PR on ROCm/rocm-libraries

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

ThomasNing · 2026-02-19T19:24:13Z

tile_engine/ops/pooling/configs/default_config.json

+        "block_n": {
+            "values": [1]


Block_n could also be 2 in pooling.

ThomasNing · 2026-02-19T19:26:00Z

tile_engine/ops/gemm/README.md

    --config_json configs/user_provided_config.json \
-    --gen_all_individual
+    --gen_all_individual \
+    --gpu_target gfx942


Great comment. TileEngine supports on gfx90a and gfx1_ series.

ThomasNing · 2026-02-19T19:27:22Z

tile_engine/ops/pooling/CMakeLists.txt

+# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+# SPDX-License-Identifier: MIT
+
+set(POOL_DATATYPE "fp16;fp32" CACHE STRING "List of datatypes for Pool (semicolon-separated)")


We could also add fp8 here.

ThomasNing · 2026-02-19T19:34:22Z

tile_engine/ops/pooling/CMakeLists.txt

+    string(REPLACE "_" ";" trait_parts ${trait})
+    list(GET trait_parts 0 output_index)
+    list(GET trait_parts 1 propagate_nan)
+    list(GET trait_parts 2 pool_dim)


We have both the pool_dim in command line and the json input. Could we unify that and only keep 1? I suggest we only keep it in json.

ThomasNing · 2026-02-19T19:39:58Z

tile_engine/ops/pooling/pool_profiler.hpp

+#include "ck_tile/host/reference/reference_pool.hpp"
+#include "pool_benchmark.hpp"
+
+class PoolProfiler


Could I know what is the difference between the pool_profiler and pool_benchmark_single.cpp? I didn't see the compilation of the pool_profiler.hpp.

aledudek added 5 commits November 27, 2025 11:31

[CK_TILE] Add pooling to ckTileEngine part1

99f5c2f

[CK_TILE] Add pooling to ckTileEngine part2

990f132

[CK_TILE] Add pooling to ckTileEngine part3

07c078d

[CK_TILE] Add pooling to ckTileEngine part4 fix suppported configurat…

f6d2243

…ions

Merge branch 'develop' into ckTileEnginePooling

d5cef00

aledudek requested review from a team, ThomasNing, afagaj, andriy-ca, aosewski, asleepzzz, bartekxk, carlushuang, cgmillette, coderfeli, ddembeckAMD, geyyer, illsilin, poyenc, qianfengz, shumway, tenpercent and vidyasagar-amd as code owners December 12, 2025 10:36

[CK_TILE] fix formatting of pooling in ckTileEngine with clang-format

1bccd37

aledudek force-pushed the ckTileEnginePooling branch from b7bbf03 to 1bccd37 Compare December 16, 2025 10:40

aledudek added 3 commits December 16, 2025 10:42

Merge branch 'develop' into ckTileEnginePooling

4c8dbd3

[CK_TILE] fix formatting of pooling in ckTileEngine with clang-format…

ef33d16

… part3

[CK_TILE] Add pooling tests for ckTileEngine

39ad789

spolifroni-amd previously approved these changes Dec 17, 2025

View reviewed changes

aosewski requested a review from Copilot December 18, 2025 10:43

Copilot AI reviewed Dec 18, 2025

View reviewed changes

msaffari-amd reviewed Jan 7, 2026

View reviewed changes

illsilin assigned ThomasNing Jan 27, 2026

Update tile_engine/ops/pooling/pool_instance_builder.py

9e97c2d

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

ThomasNing dismissed spolifroni-amd’s stale review via 9e97c2d February 19, 2026 19:17

ThomasNing requested review from Snektron and vpietila-amd as code owners February 19, 2026 19:17

Update tile_engine/ops/pooling/pool_profiler.hpp

99b4da1

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

ThomasNing requested changes Feb 19, 2026

View reviewed changes

ThomasNing mentioned this pull request Feb 19, 2026

[CK_TILE] Add pooling in tile_engine ROCm/rocm-libraries#4469

Open

1 task

Comments

Conversation

aledudek commented Dec 12, 2025

Proposed changes

Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Copilot AI Dec 18, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ammallya commented Feb 3, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants