Skip to content

Replace nested static_for lambdas with compile-time search helper - #3600

Closed
tenpercent wants to merge 6 commits into
developfrom
mpodkory/find-transform-optimization
Closed

Replace nested static_for lambdas with compile-time search helper#3600
tenpercent wants to merge 6 commits into
developfrom
mpodkory/find-transform-optimization

Conversation

@tenpercent

@tenpercent tenpercent commented Jan 16, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

  • Add find_in_tuple_of_sequences compile-time search helper with O(1) template depth
  • Replace nested static_for lambdas in TensorDescriptor::GetTransformAndItsUpperDimension
  • Replace generate_tuple lambda in TensorDescriptor::InitializeElementSize with pack expansion
  • Apply same optimizations to TensorAdaptor

Motivation

The TensorDescriptor and TensorAdaptor classes had excessive template instantiation from:

  1. Nested static_for loops with lambdas (918 applier::operator() instantiations)
  2. generate_tuple with lambdas (78+ instantiations per class)

Why It Works

Each lambda creates a unique closure type, causing separate instantiations at every call site. The find_in_tuple_of_sequences helper uses O(1) template depth via pack expansion instead of O(N) nested static_for recursion, and named functors share a single type across all uses.

Results (example_grouped_conv_fwd_xdl_fp16)

Metric Before After Improvement
Template instantiation time 23.4s 19.1s 18% reduction
applier instantiations 1132 127 89% reduction
generate_tuple lambdas 178 96 46% reduction

Test Plan

  • Added 11 unit tests:
    • 5 tests for sequence_find_value
    • 6 tests for find_in_tuple_of_sequences
  • Waiting for full CI

PR Stack

This PR is part of the build time optimization effort (issue #3575). All PRs now target develop independently:

# PR Description Status
1 #3585 sequence_gen with __make_integer_seq Independent
2 #3628 generate_identity_sequences + named functors New (replaces #3588, #3589)
3 #3590 container_concat optimization Independent
4 #3596 O(1) pack expansion rewrites Independent
5 #3600 TensorDescriptor/TensorAdaptor lambda elimination This PR

Tracking issue: #3575

This was referenced Jan 16, 2026
Merged
Closed
Closed
Closed
Closed
@tenpercent
tenpercent marked this pull request as ready for review January 17, 2026 03:41
@tenpercent
tenpercent force-pushed the mpodkory/recursive-to-pack-expansion branch from f5ada17 to 9942fd6 Compare January 17, 2026 03:51
@tenpercent
tenpercent force-pushed the mpodkory/find-transform-optimization branch 2 times, most recently from e6040e1 to a565d87 Compare January 17, 2026 05:39
@shumway shumway linked an issue Jan 19, 2026 that may be closed by this pull request
[Epic] Reduce CK / CKTile Build Times #3575
Closed
@shumway shumway removed a link to an issue Jan 19, 2026
[Epic] Reduce CK / CKTile Build Times #3575
Closed
@tenpercent tenpercent mentioned this pull request Jan 19, 2026
Closed
@tenpercent
tenpercent force-pushed the mpodkory/recursive-to-pack-expansion branch from c4d95f7 to 631df4f Compare January 21, 2026 23:57
@tenpercent
tenpercent force-pushed the mpodkory/find-transform-optimization branch from 1d351ad to ec8e794 Compare January 21, 2026 23:57
@tenpercent
tenpercent changed the base branch from mpodkory/recursive-to-pack-expansion to develop January 22, 2026 01:05
@tenpercent
Replace nested static_for lambdas with compile-time search helper
f65e9e4 
The GetTransformAndItsUpperDimension function used nested static_for
loops with lambdas to search for a hidden dimension in UpperDimensionIdss.
This caused 918 applier::operator() instantiations (81% of all applier
instantiations).

Replace with find_in_tuple_of_sequences helper that uses constexpr
array lookup and if-constexpr recursion, eliminating the lambda
instantiation overhead.

Results on example_grouped_conv_fwd_xdl_fp16:
- applier instantiations: 1132 -> 127 (89% reduction)
- TensorDescriptor instantiations: 2503 -> 664 (73% reduction)
- Template instantiation time: 23.4s -> 19.4s (17% reduction)
@tenpercent
Replace generate_tuple lambda with pack expansion in InitializeElemen…
1159278 
…tSize

The InitializeElementSize function used generate_tuple with a lambda to
compute visible dimension lengths. Each TensorDescriptor type created
a unique lambda type, causing 78 instantiations (385ms).

Replace with direct pack expansion using helper functions, eliminating
the lambda instantiation overhead entirely.

Results on example_grouped_conv_fwd_xdl_fp16:
- generate_tuple lambdas: 178 -> 100 (44% reduction)
- Template instantiation time: 19.5s -> 19.0s
@tenpercent
Apply same optimization pattern to TensorAdaptor
bc802ff 
TensorAdaptor has identical InitializeElementSize and
GetTransformAndItsUpperDimension patterns as TensorDescriptor.
Apply the same optimization:
- Replace nested static_for lambdas with find_in_tuple_of_sequences
- Replace generate_tuple lambda with pack expansion

Results: generate_tuple lambdas 100 -> 96 (4 events, 17ms eliminated)
@tenpercent
Add unit tests for sequence_find_value and find_in_tuple_of_sequences
83a76d7
@tenpercent
tenpercent force-pushed the mpodkory/find-transform-optimization branch from ec8e794 to 83a76d7 Compare January 22, 2026 01:13
@tenpercent
Add inline documentation for search helper optimizations
8d5da5c 
Detailed comments explain:
- sequence_find_value: Constexpr loop with O(1) template depth vs O(N) recursive
- find_in_tuple_of_sequences: Pack expansion instead of nested static_for loops
- Why constexpr search reduces template instantiations dramatically
- When to apply constexpr search patterns for compile-time operations
- Implementation details for each optimization approach

This documentation helps maintainers understand the compile-time search optimization
strategy without relying on specific benchmark numbers that may vary by use case.
@tenpercent
Apply clang-format with -style=file
19a156a
@tenpercent
tenpercent marked this pull request as draft January 22, 2026 18:49
@cgmillette cgmillette self-assigned this Jan 23, 2026
@assistant-librarian assistant-librarian Bot mentioned this pull request Feb 3, 2026
Closed
2 tasks
@ammallya

ammallya commented Feb 3, 2026

Copy link
Copy Markdown

Imported to ROCm/rocm-libraries

@ammallya ammallya closed this Feb 3, 2026
@tenpercent tenpercent mentioned this pull request Apr 22, 2026
Merged
3 tasks
illsilin pushed a commit to ROCm/rocm-libraries that referenced this pull request Jun 2, 2026
@tenpercent
Replace nested static_for lambdas with compile-time search helper (#6696
9627b91 
)

## Summary

- Add `sequence_find_value` and `find_in_tuple_of_sequences`
compile-time search helpers with O(1) template depth
- Replace nested `static_for` lambdas in
`TensorDescriptor::GetTransformAndItsUpperDimension` and
`InitializeElementSize`
- Apply same optimizations to `TensorAdaptor`

Supersedes #4287. Conflict-resolved rebase of
ROCm/composable_kernel#3600 onto current develop.

## Motivation

The `TensorDescriptor` and `TensorAdaptor` classes had excessive
template instantiation from:
1. Nested `static_for` loops with lambdas creating unique closure types
at every call site
2. `generate_tuple` with lambdas causing per-type instantiation overhead

The new helpers use constexpr array lookup and pack expansion instead of
recursive template patterns, achieving O(1) template depth.

## Results (`example_grouped_conv_fwd_xdl_fp16`, n=10, interleaved,
`-j1`, `-ftime-trace`)

| TU | Baseline (mean) | New (mean) | Delta | Wilcoxon p | Mann-Whitney
p |

|----|-----------------|------------|-------|-----------|---------------|
| `grouped_conv_fwd_xdl_fp16` (host) | 14,886 ms | 13,353 ms |
**-10.3%** | **0.002** | **0.0002** |
| `grouped_conv_fwd_xdl_fp16` (device) | 27,762 ms | 25,629 ms |
**-7.7%** | **0.002** | **0.0002** |
| **Total (all TUs)** | **57,732 ms** | **54,030 ms** | **-6.4%** | | |

Unrelated TUs (`device_memory`, `host_tensor`, `convolution_parameter`)
show no significant difference (p > 0.3), serving as negative controls.

### Methodology

- 10 interleaved runs (baseline₁, new₁, baseline₂, new₂, ...) on the
same node to eliminate ordering/warmup bias
- Wilcoxon signed-rank test (paired, non-parametric) and Mann-Whitney U
test (unpaired)
- Built with patched clang (LLVM 22) on ctr2-alola-compile-11, `-j1` for
accurate per-TU timing
- Raw data available in Slurm job 275230 results

## Test plan

- [x] 11 unit tests added (5 for `sequence_find_value`, 6 for
`find_in_tuple_of_sequences`)
- [x] Compile-time benchmark with statistical significance (p < 0.01)
- [ ] Full CI

Tracking issue: #4229

---
Generated-by: Claude Code (claude-sonnet-4-6)
assistant-librarian Bot pushed a commit that referenced this pull request Jun 2, 2026
@tenpercent @assistant-librarian
[rocm-libraries] ROCm/rocm-libraries#6696 (commit 9627b91)
d574cc4 
Replace nested static_for lambdas with compile-time search
 helper (#6696)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Summary

- Add `sequence_find_value` and `find_in_tuple_of_sequences`
compile-time search helpers with O(1) template depth
- Replace nested `static_for` lambdas in
`TensorDescriptor::GetTransformAndItsUpperDimension` and
`InitializeElementSize`
- Apply same optimizations to `TensorAdaptor`

Supersedes #4287. Conflict-resolved rebase of
#3600 onto current develop.

## Motivation

The `TensorDescriptor` and `TensorAdaptor` classes had excessive
template instantiation from:
1. Nested `static_for` loops with lambdas creating unique closure types
at every call site
2. `generate_tuple` with lambdas causing per-type instantiation overhead

The new helpers use constexpr array lookup and pack expansion instead of
recursive template patterns, achieving O(1) template depth.

## Results (`example_grouped_conv_fwd_xdl_fp16`, n=10, interleaved,
`-j1`, `-ftime-trace`)

| TU | Baseline (mean) | New (mean) | Delta | Wilcoxon p | Mann-Whitney
p |

|----|-----------------|------------|-------|-----------|---------------|
| `grouped_conv_fwd_xdl_fp16` (host) | 14,886 ms | 13,353 ms |
**-10.3%** | **0.002** | **0.0002** |
| `grouped_conv_fwd_xdl_fp16` (device) | 27,762 ms | 25,629 ms |
**-7.7%** | **0.002** | **0.0002** |
| **Total (all TUs)** | **57,732 ms** | **54,030 ms** | **-6.4%** | | |

Unrelated TUs (`device_memory`, `host_tensor`, `convolution_parameter`)
show no significant difference (p > 0.3), serving as negative controls.

### Methodology

- 10 interleaved runs (baseline₁, new₁, baseline₂, new₂, ...) on the
same node to eliminate ordering/warmup bias
- Wilcoxon signed-rank test (paired, non-parametric) and Mann-Whitney U
test (unpaired)
- Built with patched clang (LLVM 22) on ctr2-alola-compile-11, `-j1` for
accurate per-TU timing
- Raw data available in Slurm job 275230 results

## Test plan

- [x] 11 unit tests added (5 for `sequence_find_value`, 6 for
`find_in_tuple_of_sequences`)
- [x] Compile-time benchmark with statistical significance (p < 0.01)
- [ ] Full CI

Tracking issue: #4229
shumway pushed a commit that referenced this pull request Aug 18, 2026
@tenpercent
[rocm-libraries] ROCm/rocm-libraries#6696 (commit 9627b91)
1014634 
Replace nested static_for lambdas with compile-time search helper

## Summary

- Add `sequence_find_value` and `find_in_tuple_of_sequences`
compile-time search helpers with O(1) template depth
- Replace nested `static_for` lambdas in
`TensorDescriptor::GetTransformAndItsUpperDimension` and
`InitializeElementSize`
- Apply same optimizations to `TensorAdaptor`

Supersedes #4287. Conflict-resolved rebase of
#3600 onto current develop.

## Motivation

The `TensorDescriptor` and `TensorAdaptor` classes had excessive
template instantiation from:
1. Nested `static_for` loops with lambdas creating unique closure types
at every call site
2. `generate_tuple` with lambdas causing per-type instantiation overhead

The new helpers use constexpr array lookup and pack expansion instead of
recursive template patterns, achieving O(1) template depth.

## Results (`example_grouped_conv_fwd_xdl_fp16`, n=10, interleaved,
`-j1`, `-ftime-trace`)

| TU | Baseline (mean) | New (mean) | Delta | Wilcoxon p | Mann-Whitney
p |

|----|-----------------|------------|-------|-----------|---------------|
| `grouped_conv_fwd_xdl_fp16` (host) | 14,886 ms | 13,353 ms |
**-10.3%** | **0.002** | **0.0002** |
| `grouped_conv_fwd_xdl_fp16` (device) | 27,762 ms | 25,629 ms |
**-7.7%** | **0.002** | **0.0002** |
| **Total (all TUs)** | **57,732 ms** | **54,030 ms** | **-6.4%** | | |

Unrelated TUs (`device_memory`, `host_tensor`, `convolution_parameter`)
show no significant difference (p > 0.3), serving as negative controls.

### Methodology

- 10 interleaved runs (baseline₁, new₁, baseline₂, new₂, ...) on the
same node to eliminate ordering/warmup bias
- Wilcoxon signed-rank test (paired, non-parametric) and Mann-Whitney U
test (unpaired)
- Built with patched clang (LLVM 22) on ctr2-alola-compile-11, `-j1` for
accurate per-TU timing
- Raw data available in Slurm job 275230 results

## Test plan

- [x] 11 unit tests added (5 for `sequence_find_value`, 6 for
`find_in_tuple_of_sequences`)
- [x] Compile-time benchmark with statistical significance (p < 0.01)
- [ ] Full CI

Tracking issue: #4229

---
Generated-by: Claude Code (claude-sonnet-4-6)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@illsilin illsilin Awaiting requested review from illsilin illsilin is a code owner
@carlushuang carlushuang Awaiting requested review from carlushuang carlushuang is a code owner
@qianfengz qianfengz Awaiting requested review from qianfengz qianfengz is a code owner
@aosewski aosewski Awaiting requested review from aosewski aosewski is a code owner
@poyenc poyenc Awaiting requested review from poyenc poyenc is a code owner
@geyyer geyyer Awaiting requested review from geyyer geyyer is a code owner
@bartekxk bartekxk Awaiting requested review from bartekxk bartekxk is a code owner
@andriy-ca andriy-ca Awaiting requested review from andriy-ca andriy-ca is a code owner
@afagaj afagaj Awaiting requested review from afagaj afagaj is a code owner
@asleepzzz asleepzzz Awaiting requested review from asleepzzz
@ThomasNing ThomasNing Awaiting requested review from ThomasNing
@coderfeli coderfeli Awaiting requested review from coderfeli coderfeli is a code owner
@cgmillette cgmillette Awaiting requested review from cgmillette cgmillette is a code owner
@shumway shumway Awaiting requested review from shumway shumway is a code owner
@vidyasagar-amd vidyasagar-amd Awaiting requested review from vidyasagar-amd vidyasagar-amd is a code owner
@vpietila-amd vpietila-amd Awaiting requested review from vpietila-amd vpietila-amd is a code owner
@Snektron Snektron Awaiting requested review from Snektron Snektron is a code owner
@JiaLuo-CAN JiaLuo-CAN Awaiting requested review from JiaLuo-CAN JiaLuo-CAN will be requested when the pull request is marked ready for review JiaLuo-CAN is a code owner

Assignees

@cgmillette cgmillette

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@tenpercent @ammallya @cgmillette