Commit b1845b1
[rocm-libraries] ROCm/rocm-libraries#12173 (commit 4606f83)
feat(ck-tile): make dispatcher LDS capacity budget
architecture-aware (#12173)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
JIRA ID : AICK-2249
## Summary
The dispatcher's codegen rejects any GEMM tile whose LDS staging
footprint exceeds a cap. The cap was keyed on the **pipeline only** and
carried no architecture term, so newer targets with substantially more
LDS than gfx942 were all held to gfx942's budget.
The effect is that the largest and deepest tiles — the ones most likely
to win on large GEMM shapes — were never generated, never benchmarked
and never selectable on those targets. On a gfx950 fp16 sweep, **160 of
288 offered configurations (56%) were rejected by this cap alone**.
This PR makes the budget architecture-aware, adds gfx1250 to the spec,
fixes both the Python and C++ validators together, and makes the check
account for ping-pong LDS staging.
**Scope:** the check is not GEMM-specific. `_validate_lds_capacity` sits
in the common `validate_kernel` chain, ungated by operator, and
`OperatorType` covers 11 operators (5 GEMM variants, 6 conv variants).
In Python it is reached by `unified_gemm_codegen.py:1595` (universal
GEMM) and `unified_grouped_conv_codegen.py:1941` (grouped convolution);
in C++ by `Registry::filter_by_arch()` (`registry.cpp:160`), which is
operator-agnostic and covers every registered kernel. The ctypes bridge
paths (quant variants, batched contraction, multi-ABD, FMHA) validate
via `validate_kernel_config` in `python/ctypes_utils.py`, which has no
LDS check at all, so they are unaffected either way.
Implements AICK-2249. Baseline evidence from AICK-2244.
## Motivation
`arch_specs.json` had `pipeline_lds_limits` as a **top-level** key — a
sibling of `architectures`, not nested inside it:
```json
"pipeline_lds_limits": {
"mem": <bytes>, "compv3": <bytes>, "compv4": <bytes>, ... "default": <bytes>
}
```
Every value was a fixed byte count, derived from gfx942 and applied to
every target.
There was no architecture axis, and nowhere to put one.
`_validate_lds_capacity` did `LDS_CAPACITY_LIMITS.get(config.pipeline,
...)` and nothing else.
Meanwhile `arch.hpp` has declared a per-architecture LDS capacity all
along, via `get_lds_size()`. On gfx942 the hardcoded cap and that
capacity coincide; on the newer targets they do not, and the shortfall
is large. This PR consumes the existing declaration rather than
introducing any new hardware figure.
This is the usual shape of this defect: a value that was correct when
there was exactly one architecture, frozen into a schema with no slot
for a second one. It survived because on gfx942 the correct answer and
the hardcoded answer are the same number, so it only ever fails
*silently, by generating less*.
The evidence that it is binding rather than theoretical: on the gfx950
sweep the cutoff falls **exactly** on the hardcoded cap for compv3/mem,
and exactly on the tighter cap for compv4, with zero exceptions in
either direction. And the winning kernels on that target were
`128x256x64` and `128x128x128` under compv3 — precisely tiles that
compv4 was forbidden from using. When the measured optimum sits on the
constraint boundary, the constraint is probably binding.
## Design note
**I did not invent a new schema.** A per-architecture `lds_capacity_kb`
field already existed in `arch_specs.json` (gfx950 already correctly
said `160`), and `ADDING_NEW_GPU.md` already documented it as a required
onboarding field. Nothing read it. Every new-GPU onboarding has been
filling in a mandatory field that no code consumed — the contract was
documented and unhonoured. So the fix is to honour it.
Pipelines now declare a **basis** rather than a byte count:
```json
"compv3": { "basis": "fraction", "value": 1.0 },
"compv4": { "basis": "fraction", "value": 0.5 }
```
resolved against each architecture's capacity at generation time. A new
GPU needs **one** number, not ten. The alternative — nesting the byte
table under each architecture — would have reintroduced the same failure
mode one level down: 8 arches x 10 pipelines of hand-maintained
constants free to drift from `arch.hpp` independently.
Resolution happens in the generator, once, so the emitted Python and C++
get literal byte counts and cannot disagree about how to read the
schema.
### On the tighter-capped pipelines — I read the sources rather than
guessing
Four pipelines sat at the tighter cap and the ticket flagged it as an
open question: genuine double-buffering that should scale with capacity,
or an independent absolute limit? On gfx942 the two readings are
numerically identical, which is why it was never forced. **The answer is
not uniform across the four:**
| pipeline | doubles LDS? | evidence |
|---|---|---|
| `compv4` | **yes** | `GemmPipelineAgBgCrCompV4::GetSmemSize()` returns
`2 * Policy::GetSmemSize<Problem>()` |
| `preshufflev2` | **yes** |
`WeightPreshufflePipelineAGmemBGmemCRegV2::GetSmemSize()` returns
`DoubleSmemBuffer ? 2 * smem_size : smem_size` |
| `compv6` | **no** | returns the policy size unmultiplied |
| `preshufflev1` | n/a | no such pipeline exists;
`preshuffle_pipelines.supported` lists only `preshufflev2` |
This is corroborated by `DOUBLE_SMEM_PIPELINES = {"compv4",
"preshufflev2", "comp_async"}`, which already exists in
`unified_batched_contraction_codegen.py` and names exactly the same two.
So for compv4 and preshufflev2, half-of-capacity is the **exact** model,
not a safety margin: the validator checks `A+B` against `capacity/2`,
which is precisely equivalent to checking the real `2*(A+B)` allocation
against full capacity.
For compv6 and preshufflev1 the halved cap is **unexplained by
buffering**. I preserved their ratio rather than widening them, because
no evidence supports a larger budget and a blind raise that regresses is
worse than the status quo. Widening those two is a separate change that
needs a measurement behind it. The reasoning is recorded in
`arch_specs.json` next to each value.
### Ping-pong staging had to be threaded in
The check modelled a *single* staging buffer. But
`unified_grouped_conv_codegen.py:148` documents that `mem`, `compv3`,
`compv5` and `compv6` make double buffering a **configuration choice**
(`--double-smem-buffer`), not a property of the pipeline — and the conv
path passed `pipeline=` into the validator without that flag, so the
budget could not see it.
That was harmless while the cap matched gfx942's capacity, because twice
that still fits the larger parts. **Widening the budget removes the
accidental headroom**, so I had to close it in the same PR:
Writing `C` for a target's LDS capacity:
| | compv3 budget | actually allocated | vs capacity | |
|---|---|---|---|---|
| develop | gfx942's cap | 2x that | below `C` | fits, by accident |
| arch-aware alone | `C` | `2C` | **over `C`** | **overflows** |
| with this fix | `C/2` | `C` | exactly `C` | fits exactly |
So the flag is now threaded into both validators. Conv passes the value
it already tracks; C++ reads `algorithm.double_buffer`, which registered
kernels populate from `SelectedKernel::DoubleSmemBuffer`
(`kernel_registration.hpp:60`).
Pipelines that *always* double already carry the halving in their
per-pipeline budget, so the two signals are combined with a `min()` —
`compv4` and `preshufflev2` are never halved twice, which is what keeps
gfx942 byte-identical. Independent confirmation that those are exactly
the two: `unified_gemm_codegen.py:1531` sets `double_buffer = pipeline
in ("compv4", "preshufflev2")`.
One deliberate non-use: `KernelConfig::build_key()`
(`kernel_config.hpp:280`) and `utils.hpp:676` hardcode `double_buffer =
true` for *any* pipeline. Those build **query** keys, which never reach
`validate_lds` — only registered-instance keys do — so the flag is
trustworthy at the one site that reads it. Worth fixing separately.
### A landmine worth calling out
`arch_specs_generated.py` **already contained gfx1250** — family, warp
configs, warp tile combos — but `arch_specs.json`, the file it is
generated from, did not. The generated C++ header did not either.
Three-way drift, in a file stamped `AUTO-GENERATED - DO NOT EDIT
DIRECTLY`.
This means following the documented workflow (edit JSON, regenerate,
commit) would have **silently deleted gfx1250's warp tables**. I
transplanted them into the JSON verbatim rather than authoring anything,
and the regeneration check below confirms every table came back
byte-identical.
## Test plan
- [x] New regression suite passes (15/15). Asserts budgets differ across
gfx942/gfx950/gfx1250, that gfx942's budget is byte-identical to the
historical table, that no budget exceeds the declared hardware capacity,
that unknown targets get the *smallest* budget rather than the largest,
and end-to-end that a large staging tile is rejected on gfx942 and
accepted on gfx950
- [x] Double-buffer cases covered: configurable pipelines halve, `2 x
budget <= capacity` on every arch/pipeline pair, always-double pipelines
are not halved twice, single-buffered remains the default, and
end-to-end a large staging tile is accepted on gfx950 single-buffered
but rejected double-buffered
- [x] Existing suites pass: `test_arch_filter_constraints`,
`test_gemm_utils`, `test_codegen_common`, `test_dispatcher_common`,
`test_tile_math`, `test_grouped_conv_codegen`, `test_grouped_conv_utils`
- [x] Generator is idempotent; regeneration leaves all non-LDS tables
**byte-identical** (`ARCH_FAMILY_MAP`, `WARP_SUPPORTED_COMBINATIONS`,
`WARP_TILE_SUPPORTED_COMBINATIONS`,
`PRESHUFFLE_WARP_TILE_SUPPORTED_COMBINATIONS`,
`TRAIT_UNSUPPORTED_COMBINATIONS`, `ELEMENT_SIZE_MAP`,
`DTYPE_COMBINATIONS`, `PRESHUFFLE_PIPELINES`)
- [x] C++ compiles standalone and returns values identical to Python for
every architecture and pipeline, single- and double-buffered, including
the unknown-architecture fallback
- [x] `clang-format-18 -style=file` clean on both headers
- [ ] **GPU validation pending** — see below
### Equivalence proof
This edits a shared table, so I enumerated the validator's survivor set
per architecture, before and after. The axes now include **every
pipeline the validators can see** — `comp_async` and `wavelet` as well
as the nine that were in the old table — and both settings of the
ping-pong staging flag.
| arch | develop | this PR | gained | lost |
|---|---:|---:|---:|---:|
| gfx908 | 4,613 | 7,386 | +3,003 | **230** |
| gfx90a | 2,701 | 4,138 | +1,595 | **158** |
| gfx942 | 5,324 | 8,360 | +3,322 | **286** |
| gfx950 | 7,116 | 19,016 | +11,900 | 0 |
| gfx1100 | 1,155 | 1,870 | +770 | **55** |
| gfx1200 | 442 | 668 | +253 | **27** |
| gfx1201 | 442 | 668 | +253 | **27** |
| gfx1250 | 2,012 | 10,926 | +8,914 | 0 |
| **total** | | | **+30,010** | **783** |
**There are losses, and an earlier version of this section claimed there
were none.** That claim was measured over the old parameter space, which
could not contain `comp_async` because the pipeline had no entry in the
old table to enumerate. It was true of what it measured and wrong as a
general statement.
**All 783 losses are `comp_async`, on the six architectures whose budget
is unchanged**, and they are intended. Attribution by pipeline across
those six:
```
comp_async 783
(no other pipeline appears)
```
`comp_async` had no entry before, so it inherited the full-capacity
default. It allocates two LDS buffers unconditionally — `GetSmemSize()`
returns `num_lds_buffers * smem_size` with `num_lds_buffers = 2` — so it
was budgeted for twice the staging it can actually use. A tile such as
conv `128x64x128` at fp16 needs 48 KB of staging and was accepted, while
the kernel would then have asked the hardware for twice that. Rejecting
it is the fix, not a regression.
The two architectures with a raised budget lose nothing, and no default
configuration set in tree generates `comp_async`, so no shipped kernel
disappears.
## Explicitly out of scope
Two other filters in the same validation chain share this defect class
(gfx942-derived constraints applied to every architecture). Neither
rejected anything in the baseline campaign, so I left them alone:
- `TRAIT_UNSUPPORTED_COMBINATIONS` — 10 tuples, no arch key
- `_cshuffle_store_ok` — docstring says "GPU-verified on gfx942",
applied to all architectures
A wider audit of the validation chain found this same pattern in several
more places, including a hand-written warp-tile lookup on the native
side that disagrees with the generated Python table on a majority of
supported architectures. That work is tracked separately in AICK-2264
and AICK-2266, under epic AICK-2265, and is deliberately not part of
this PR: it is a different table on a different axis, and unlike this
change it removes configurations, so it needs its own before/after
enumeration and device validation.
Also left alone: `get_smem_capacity()` in `arch.hpp` special-cases a
single architecture and returns a fixed value otherwise, which disagrees
with `get_lds_size()` for at least one target. It has three live
consumers outside this ticket's scope (`cshuffle_epilogue.hpp`,
`moe_sorting_kernel.hpp`, `grouped_convolution_forward_kernel.hpp`) and
may legitimately encode a per-workgroup addressable limit rather than
per-CU capacity. Flagging rather than changing it.1 parent 6fdf2b0 commit b1845b1
13 files changed
Lines changed: 1358 additions & 83 deletions
File tree
- dispatcher
- codegen
- include/ck_tile/dispatcher
- tests
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
159 | 159 | | |
160 | 160 | | |
161 | 161 | | |
162 | | - | |
| 162 | + | |
| 163 | + | |
163 | 164 | | |
164 | 165 | | |
165 | 166 | | |
| |||
241 | 242 | | |
242 | 243 | | |
243 | 244 | | |
244 | | - | |
| 245 | + | |
| 246 | + | |
| 247 | + | |
| 248 | + | |
| 249 | + | |
| 250 | + | |
| 251 | + | |
| 252 | + | |
| 253 | + | |
| 254 | + | |
| 255 | + | |
| 256 | + | |
| 257 | + | |
| 258 | + | |
| 259 | + | |
| 260 | + | |
| 261 | + | |
| 262 | + | |
| 263 | + | |
| 264 | + | |
| 265 | + | |
| 266 | + | |
| 267 | + | |
| 268 | + | |
| 269 | + | |
| 270 | + | |
| 271 | + | |
| 272 | + | |
| 273 | + | |
| 274 | + | |
| 275 | + | |
| 276 | + | |
| 277 | + | |
| 278 | + | |
| 279 | + | |
| 280 | + | |
| 281 | + | |
245 | 282 | | |
246 | 283 | | |
247 | 284 | | |
| |||
353 | 390 | | |
354 | 391 | | |
355 | 392 | | |
| 393 | + | |
| 394 | + | |
| 395 | + | |
| 396 | + | |
| 397 | + | |
356 | 398 | | |
357 | 399 | | |
358 | 400 | | |
| |||
533 | 575 | | |
534 | 576 | | |
535 | 577 | | |
| 578 | + | |
536 | 579 | | |
537 | 580 | | |
538 | 581 | | |
| |||
544 | 587 | | |
545 | 588 | | |
546 | 589 | | |
| 590 | + | |
| 591 | + | |
547 | 592 | | |
548 | 593 | | |
549 | 594 | | |
| |||
568 | 613 | | |
569 | 614 | | |
570 | 615 | | |
| 616 | + | |
571 | 617 | | |
572 | 618 | | |
573 | 619 | | |
| |||
709 | 755 | | |
710 | 756 | | |
711 | 757 | | |
| 758 | + | |
| 759 | + | |
| 760 | + | |
| 761 | + | |
| 762 | + | |
| 763 | + | |
| 764 | + | |
712 | 765 | | |
713 | | - | |
| 766 | + | |
714 | 767 | | |
715 | 768 | | |
716 | | - | |
717 | | - | |
| 769 | + | |
| 770 | + | |
| 771 | + | |
| 772 | + | |
718 | 773 | | |
719 | 774 | | |
720 | 775 | | |
| 776 | + | |
721 | 777 | | |
722 | | - | |
| 778 | + | |
| 779 | + | |
| 780 | + | |
723 | 781 | | |
724 | 782 | | |
725 | 783 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
172 | 172 | | |
173 | 173 | | |
174 | 174 | | |
| 175 | + | |
| 176 | + | |
| 177 | + | |
| 178 | + | |
| 179 | + | |
| 180 | + | |
| 181 | + | |
| 182 | + | |
| 183 | + | |
| 184 | + | |
| 185 | + | |
| 186 | + | |
| 187 | + | |
| 188 | + | |
| 189 | + | |
| 190 | + | |
| 191 | + | |
| 192 | + | |
| 193 | + | |
| 194 | + | |
| 195 | + | |
| 196 | + | |
| 197 | + | |
| 198 | + | |
| 199 | + | |
| 200 | + | |
| 201 | + | |
| 202 | + | |
| 203 | + | |
| 204 | + | |
| 205 | + | |
| 206 | + | |
| 207 | + | |
| 208 | + | |
| 209 | + | |
| 210 | + | |
| 211 | + | |
| 212 | + | |
| 213 | + | |
| 214 | + | |
| 215 | + | |
| 216 | + | |
| 217 | + | |
| 218 | + | |
| 219 | + | |
| 220 | + | |
| 221 | + | |
| 222 | + | |
175 | 223 | | |
176 | 224 | | |
177 | | - | |
| 225 | + | |
178 | 226 | | |
179 | 227 | | |
180 | 228 | | |
| |||
221 | 269 | | |
222 | 270 | | |
223 | 271 | | |
224 | | - | |
225 | | - | |
226 | | - | |
227 | | - | |
228 | | - | |
229 | | - | |
230 | | - | |
231 | | - | |
232 | | - | |
233 | | - | |
234 | | - | |
235 | | - | |
| 272 | + | |
| 273 | + | |
| 274 | + | |
| 275 | + | |
| 276 | + | |
| 277 | + | |
| 278 | + | |
| 279 | + | |
| 280 | + | |
| 281 | + | |
| 282 | + | |
| 283 | + | |
| 284 | + | |
| 285 | + | |
| 286 | + | |
| 287 | + | |
| 288 | + | |
| 289 | + | |
| 290 | + | |
| 291 | + | |
| 292 | + | |
| 293 | + | |
| 294 | + | |
| 295 | + | |
| 296 | + | |
| 297 | + | |
| 298 | + | |
| 299 | + | |
| 300 | + | |
| 301 | + | |
| 302 | + | |
| 303 | + | |
| 304 | + | |
| 305 | + | |
| 306 | + | |
| 307 | + | |
| 308 | + | |
| 309 | + | |
| 310 | + | |
| 311 | + | |
| 312 | + | |
| 313 | + | |
| 314 | + | |
| 315 | + | |
| 316 | + | |
| 317 | + | |
| 318 | + | |
| 319 | + | |
| 320 | + | |
| 321 | + | |
| 322 | + | |
| 323 | + | |
| 324 | + | |
| 325 | + | |
| 326 | + | |
| 327 | + | |
| 328 | + | |
| 329 | + | |
| 330 | + | |
| 331 | + | |
| 332 | + | |
| 333 | + | |
| 334 | + | |
| 335 | + | |
| 336 | + | |
| 337 | + | |
| 338 | + | |
| 339 | + | |
| 340 | + | |
| 341 | + | |
| 342 | + | |
| 343 | + | |
| 344 | + | |
| 345 | + | |
| 346 | + | |
| 347 | + | |
| 348 | + | |
| 349 | + | |
| 350 | + | |
| 351 | + | |
| 352 | + | |
| 353 | + | |
| 354 | + | |
| 355 | + | |
| 356 | + | |
| 357 | + | |
| 358 | + | |
236 | 359 | | |
237 | 360 | | |
238 | 361 | | |
| |||
0 commit comments