Refactor/mmm pool and candidates - #2696
Open
kali wants to merge 6 commits into
Open
Conversation
Every plug pushed its own kernels into `ops.mmm_impls`, so what the pool holds was spread over ten functions and, in three places, gated on conditions the kernels did not carry: the avx tier stood down for the fma one, armvfpv2 for NEON, the 16x16 VNNI tile for dual-FMA cores, and `fma_mmm_f32_32x1` wanted f16c. Give those four their real `where` predicate and the pool follows from the inventory — every routine this build assembled that `is_supported_here`, portable kernels first and each group by name, since `inventory` yields link order and tie-breaking downstream reads pool position. The two generic 3x2 kernels stay out of the inventory, which is what keeps them out of dispatch.
Both matmul lowerings walked the kernel pool themselves — einsum in `list_impls`, conv in its block-quant branch — with filters that had drifted apart: conv checked neither `stores()` nor, before the pool carried it, `is_supported_here`, and neither knew about the other. `Ops::candidates` now answers "how can this build multiply these two operands", kernel, packing and panel extractor, and both callers ask it; `Ops::filter_impls`, a third unused variant of the same query, goes. Which candidate wins is still each caller's own business.
|
🔴 Bench vs main — 1 speed regression(s) · Reference: 2026-08-21 morning nightly run (0d old) · full report → run Speed — evaltime · prefill · decode
Improvements
|
| Δ | metric | device | main → PR |
|---|---|---|---|
| arm_ml_kws_cnn_m load · pass |
cortex-a9 |
83 ms → 106 ms | |
| arm_ml_kws_cnn_m load+optimize · pass |
cortex-a9 |
129 ms → 157 ms | |
| arm_ml_kws_cnn_m load · pass |
cortex-a7 |
77 ms → 86 ms |
`best()` picked the arch policy with a `#[cfg]` per architecture, so which policies exist was a compile-time fact no host could enumerate — and a policy could not be examined from anywhere but the machine it was written for. Each arch tree now submits a `PlatformSelector` carrying its plug and whether this build is native to it, and `best()` applies the native ones, cfg-free. `platform::inspect` pairs a target's policy with its kernels, from the new `mmm_routines::pool_for`, so one host can ask what another platform would dispatch: faithful to the letter on the native target, and elsewhere as faithful as the `TRACT_CPU_*` knobs make the hardware probes. The cost model's fallback stops unwrapping on a pool that lacks its fitted kernels, which is exactly what a foreign pool is.
`Target` named only the arches carrying a kernel tree, so a riscv64 build could not say what it runs on and a riscv64 kernel could not be declared at all. `Target::RiscV64` fills that in — naming a platform is now separate from having kernels for it, `native()` answers `Some` on riscv64 with no tree and no selector behind it — and every arch-keyed codegen macro grows a riscv64 arm, so the tree can land a kernel at a time.
|
🔴 Bench vs main — 4 speed regression(s) · Reference: 2026-08-22 morning nightly run (0d old) · full report → run Speed — evaltime · prefill · decode
|
| Δ | metric | device | main → PR |
|---|---|---|---|
| arm_ml_kws_cnn_m load · pass |
cortex-a9 |
84 ms → 104 ms | |
| arm_ml_kws_cnn_m load+optimize · pass |
cortex-a9 |
130 ms → 155 ms | |
| hey_snips_v4_model17 load+optimize · 2sec |
cortex-a7 |
4.95 s → 5.73 s | |
| hey_snips_v4_model17 load · 2sec |
cortex-a7 |
4.25 s → 4.89 s |
`candidates()` took four positional arguments and no notion of shape, so the enumeration could never consider a dim and each caller assembled the operand types on the spot. It now takes a `Query` holding the operand types plus the dims the caller can pin, and returns `Candidate`, the triple core's `Impl` alias was spelling out by hand. Selection is unchanged — the query's dims are not read yet.
The rules any platform falls back on — keep the best quality tier, then choose on n — lived in einsum, so linalg could not answer a selection question without a caller to score for it. `retain_best_quality` and `pick_by_shape` now sit beside `candidates`, and `pick_by_shape` abstains on a symbolic n, which is precisely where einsum's packing-group reasoning takes over. `list_impls` becomes `query`: the caller needs the query itself, and enumerating is one call.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.