Commit Graph

2 Commits

Author SHA1 Message Date
John Shumway
d65ad35b23 [rocm-libraries] ROCm/rocm-libraries#7180 (commit 54aed1e)
[CK] Add rocm_ck spec factories: GemmSpec, makeSpec() (#7180)

## What this PR does

This is the third PR in the rocm_ck schema stack:

1. **#7150** — Foundation types (DataType, Layout, Args, Ops)
2. **#7163** — Schema engine (Signature, resolve(), ArchProperties)
3. **#7180 (this)** — Spec factories (GemmSpec, makeSpec())

`makeSpec()` is the bridge between user intent and kernel instantiation.
It takes a **Signature** (WHAT to compute — operator graph, dtypes,
layouts) and a **GemmAlgorithm** (HOW to compute it — tile sizes,
pipeline, partitioning) and produces a validated `GemmSpec` — a
structural type ready to use as a non-type template parameter.

The key property: **every constraint is enforced at compile time.** An
invalid GEMM configuration is a compile error, not a runtime crash or
silent corruption. The 33 compile-fail tests are the executable
specification of what's allowed.

## What's interesting

**Physical tensor table.** Not every tensor in a compute graph needs
device memory. The intermediate result of `C = A * B` in a fused
GEMM+Add+ReLU lives only in registers. `makeSpec()` walks the operator
chain and determines which tensors are physical (need Args slots) and
which are intermediate. The output is a fixed-layout table: `[lhs, rhs,
output, D0?, D1?, scale?]`.

**Epilogue composition.** Instead of a combinatorial explosion of named
patterns (GemmAdd, GemmAddRelu, GemmMulSilu, ...), the epilogue is a
composable chain of ops. `{GemmOp, AddOp, ReluOp}` produces
`epilogue_ops = {Add, Relu}` with the bias tensor automatically slotted
as D0. Two consecutive AddOps fold into a single Add with two D tensors
via CK Tile's parameter pack.

**Signature/Algorithm split.** The same Signature can pair with multiple
GemmAlgorithms to produce different tuning variants without changing the
mathematical result. This is the foundation for the dispatcher — one
operation description, many tile configurations.

## New types

| Type | Role |
|------|------|
| `GemmSpec` | Validated NTTP kernel descriptor — physical tensors, tile
geometry, epilogue chain |
| `GemmAlgorithm` | User-facing tuning input — tile sizes, pipeline,
partitioning, padding |
| `EpilogueOp` | NTTP-compatible projection of the Op variant for
epilogue chains |
| `Dim3` | M x N x K triple for tile geometry |

## Test coverage

- **69 unit tests** — happy paths, layouts, dtypes, quantization,
epilogue chains, algorithm variants
- **33 compile-fail tests** — one per constraint (tile divisibility,
INT8 rules, pipeline restrictions, etc.)
- **6 schema compatibility baselines** — frozen specs that break if the
schema changes

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-05-26 12:07:31 +02:00
John Shumway
b1975951d4 [rocm-libraries] ROCm/rocm-libraries#7179 (commit 05edc86)
[CK] Add rocm_ck schema engine: Signature, resolve(), ArchProperties (#7179)

## Summary

A `Signature` is a directed compute graph: tensors are nodes, operators
are edges. Shared names between operator outputs and inputs form the
graph structure. `resolve()` walks this graph at compile time
(`consteval`), inferring dtype, rank, and layout for every tensor —
invalid configs become compiler errors, not runtime crashes.

**Key design decisions:**

- **Operators teach the system about tensors.** `GemmOp` implies rank 2
and Row/Col/Row layout. `AddOp` and `ReluOp` propagate from connected
slots. The dtype cascade fills in the rest: per-tensor → signature-wide
→ error.

- **Adding a new op is zero lines in the resolution engine** if it's
structurally binary (`lhs/rhs/out`) or unary (`in/out`) — C++20 concepts
handle dispatch automatically. Only ops with special semantics need
explicit branches.

- **TargetSet is a compile-time bitset over GPU targets.** The wave tile
validation table is the single source of truth for valid instruction
shapes, traced from CK Tile's WarpGemmDispatcher. FP8 tiles are
available on gfx942+ via IterateK composition, not gfx950-only.

**Reading order:** signature.hpp (data model) → arch_properties.hpp
(TargetSet, wave tiles) → resolve.hpp (resolution engine).

3 new headers, 3 unit tests (including diamond DAG coverage), 3
compile-fail tests. Introduces tests/compile_fail/ infrastructure.

**Stack**: PR 2 of 3 porting the rocm_ck constexpr schema from
experimental to production.
1. Foundation types — DataType, Layout, Args, Ops (#7114)
2. **This PR** — Schema engine (graph resolution)
3. Spec factories — GemmSpec, makeSpec() (#7180 )

Note: We also removed `FmhaBwdOp` for clarity, since that was introduced
early and doesn't have tests set up.

**Depends on**: #7114

## Test plan

- [x] ctest --test-dir build --output-on-failure — unit tests +
compile-fail tests pass
- [x] Compile-fail tests correctly reject: mixed CDNA+RDNA TargetSet,
conflicting layouts, empty quantization scale names

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-05-22 09:43:45 -07:00