mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-03-17 21:57:38 +00:00
56e1d5da08
[CK_TILE] Optimize ck_tile::sequence to reduce template instantiation depth [2A] (#5028) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary ### Rationale `ck_tile::sequence` is the most fundamental metaprogramming type in ck_tile — it underpins tensor dimensions, strides, loop bounds, and index calculations. Six of its metafunctions use recursive template instantiation, producing O(N) to O(N²) intermediate types that the compiler must process. When these are used inside deeply nested GEMM pipelines with large dimension counts, the cumulative instantiation overhead becomes a significant contributor to frontend compile time. Measurements on `test_gemm_pipeline_compv6` show 84,288 `InstantiateFunction` calls in the frontend alone. Reducing template instantiation depth in these core utilities has a multiplicative effect because they are called from hundreds of sites. ### What changed | Metafunction | Before | After | |---|---|---| | `sequence::modify` | O(N) recursive split/merge | O(1) pack expansion | | `sequence_gen` | O(log N) recursive binary split | O(1) via `__make_integer_seq` | | `uniform_sequence_gen` | Delegates to `sequence_gen` | O(1) via `__make_integer_seq` | | `sequence_reverse_inclusive_scan` | O(N) recursive | O(1) constexpr for-loop + pack expansion | | `sequence_inclusive_scan` | Computed via reverse + flip | O(1) constexpr for-loop (unified impl) | | `sequence_exclusive_scan` | O(N) recursive merge chain | O(1) constexpr for-loop + pack expansion | | `sequence_map_inverse` | O(N²) recursive modify calls | O(1) constexpr for-loop + pack expansion | Supporting changes: - Portable `__type_pack_element` fallback with `__has_builtin` guard (hipRTC-safe, no `<tuple>` dependency) - Renamed reserved `__integer_sequence` to `integer_sequence_wrapper` - Adopted `static_array` from develop (PR #4355) for constexpr computation - Unified forward and reverse inclusive scan into a single `sequence_inclusive_scan_impl` with `bool Reverse` template parameter - Added `sequence_inclusive_scan` struct (new public API for forward scan direction) - Replaced recursive `sequence_exclusive_scan` (3 template specializations) with `sequence_exclusive_scan_impl` using the same constexpr for-loop pattern as inclusive scan - Rewired `exclusive_scan_sequence` and `prefix_sum_sequence` to use new impl - Added `CK_TILE_HOST_DEVICE` to `exclusive_scan_sequence` and `prefix_sum_sequence` to match sibling scan function annotations ### Technical debt and housekeeping - Unified all `namespace impl` to `namespace detail` across sequence.hpp for consistency - Removed dead comment block (orphaned `integer_sequence` alternative) - Added defensive `static_assert(sizeof...(Is) > 0)` in `sequence_map_inverse::build_inverse` - Converted all multi-line Doxygen blocks from `///` to `/** */` per style guide - Corrected `constexpr static` to `static constexpr` keyword ordering in `static_array` - Added blank line between `#pragma once` and first `#include` in `static_array.hpp` - Trimmed redundant 4-line comment on `sequence_gen_helper` to a one-liner - Moved `sequence_gen` Doxygen comment below `namespace detail` block so it directly precedes the struct it documents - Added Doxygen `@brief`/`@tparam`/`@pre` documentation for `sequence_gen` and `sequence_map_inverse` public APIs - Added `@brief` documentation to `static_array` explaining relationship to `ck_tile::array` - Added scope comment at `namespace detail` openings **Note:** `private:`/`public:` access modifier indentation is enforced at 4 spaces by `.clang-format`. The style guide calls for left-alignment, but the formatter overrides this. Requires a `.clang-format` config change to resolve — not addressable in code. ### `static_array` hardening (from develop's PR #4355) - Added zero-length array guard (`T elems[N > 0 ? N : 1]`) - Added `CK_TILE_HOST_DEVICE` annotations to `operator[]` and `size()` - Added `#include "ck_tile/core/config.hpp"` (IWYU for `CK_TILE_HOST_DEVICE`) ### Value Combined with the `static_ford` changes, measured impact on `test_gemm_pipeline_compv6`: - **Frontend: -28.9%** (InstantiateFunction: 84,288 → 69,439) - **Backend: -13.1%** (CodeGen Functions: 3,170 → 2,203) - **Wall-clock: -16.3%** (611.6s → 512.2s) ### Files changed (4) - `sequence.hpp`: Metafunction optimizations, namespace unification, documentation, style fixes - `static_array.hpp`: Zero-length guard, `CK_TILE_HOST_DEVICE`, documentation, style fixes - `test_sequence.cpp`: 50 unit tests with runtime `EXPECT_EQ` assertions (new file) - `CMakeLists.txt`: Register new test target ## Test plan - [x] 50 runtime unit tests covering all optimized and pre-existing sequence APIs - [x] Edge cases: empty sequences, single-element, larger sizes (N=8), negative values, non-trivial init values - [x] Both functor signatures tested (`operator()(index_t)` and `operator()(number<I>)`) - [x] Both scan reducers (`plus`, `multiplies`) with forward, reverse, inclusive, and exclusive directions - [x] Exclusive scan: sum, product, single, empty, non-zero init - [x] Prefix sum: N+1 output verification, single, empty - [x] Permutation round-trip verification for `sequence_map_inverse` - [x] Full sequence public API coverage: modify, gen, uniform_gen, scans (inclusive, exclusive, prefix sum), map_inverse, make_index_sequence, size/sum/product, push/pop, reverse, extract, merge, arithmetic operators, equality, transform - [x] Portable `__type_pack_element` fallback tested implicitly (same `at_index_t` interface) 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.