Proto/activation registry - #2644
Draft
kali wants to merge 13 commits into
Draft
Conversation
Register each activation kernel as an inventory descriptor carrying its target, gating feature and quality tier (native / via-f32 / generic), selected by tier then ISA rank. Sigmoid is wired across generic, aarch64, armv7 and x86_64 and wasm behind a `registry-all-targets` feature so a dev host can compile every target's descriptors and draw the whole matrix; it runs alongside the existing plug() dispatch and a test asserts the two pick the same kernel. Adds a `tract activations` subcommand rendering the quality matrix by feature-tier deployment target, bolding the column this host runs.
Generalize `Tier::ViaF32` into `Tier::Via(&'static str)` so the composition a
kernel routes through ("f32", later "sigmoid") is data rather than an enum
variant, and drop the `!has_fp16()` gate on the f16 via-f32 fallbacks: a via-f32
kernel is always available and tier ranking picks a native fp16 kernel over it
only when one exists. Silu needs this — it has no native fp16 kernel, so its f16
uses via-f32 even on fp16 hardware. Adds silu descriptors across generic,
aarch64, armv7 and x86_64 (none on wasm), and extends the equivalence test to
silu f32 and f16.
Collaborator
Author
@czoli1976 @cverrier WDYT ? |
Contributor
|
Contributor
|
Deliberate non-gaps: the wasm Sigmoid generic cells aren't an oversight — #2591 > wrote simd128 sigmoid and tanh, measured them, and deleted them because the generic polynomial auto-vectorizes faster |
Each backend registers its own descriptors from a co-located `activation.rs` (generic/, arm64/, arm32/, x86_64_fma/, wasm/) instead of one central block, so inventory does what it is for — decentralized registration. The per-file modules are declared from the always-compiled `activation` module via `#[path]`, because the arch modules are `#[cfg(target_arch)]`-gated and would otherwise hide their descriptors from a host drawing the whole matrix. Also drops the redundant `matrix()` renderer (the CLI owns rendering).
Extends the activation registry to the remaining parameterless element-wise functions across generic, aarch64, armv7, x86_64 and wasm, and widens the equivalence test to all of them. LeakyRelu and MulByScalar are left out: they are parameterized (ElementWise<T, Params>) and need an ActFactory that carries the runtime param first.
Route the six parameterless element-wise activations (sigmoid, silu, tanh, erf, hardswish, gelu) through `activation::kernel_f32` / `kernel_f16` in core, and delete their field assignments from every `plug()` and their fields from `Ops`. The registry is now the sole source of truth for activation dispatch; `plug()` keeps only mmm, reducers, leaky_relu and mul_by_scalar (the latter two are parameterized and not yet ported). Registry/plug equivalence was verified green at the previous commit before removing the plug side.
Adds `ActFactory::{F32Param, F16Param}` and `kernel_f32_param` / `kernel_f16_param`
accessors for activations that take a runtime scalar (leaky-relu's alpha, an
ElementWise<T, T>), registers leaky-relu across generic, aarch64 and x86_64, routes
core through the registry and drops its fields from Ops and its assignments from
plug(). The native avx512fp16 leaky kernel stays unregistered — it is a measured
regression that plug() never dispatched. Registry/plug equivalence was verified
green before removing the plug side.
Registers mul-by-scalar (the last parameterized element-wise op) across generic, aarch64 and x86_64, routes softmax through the registry and removes its Ops fields and plug() assignments. With this, plug() no longer touches any element-wise activation: dispatch for the whole family (sigmoid/silu/tanh/erf/hardswish/gelu + leaky_relu/mul_by_scalar) goes through the activation registry. Drops the now-unused ElementWiseKer imports from lib.rs and each arch module. Equivalence verified green before removing the plug side.
Contributor
|
@kali BTW, did you ever consider Google Highway as a stopgap at least? |
Extends the registry to reductions: adds Reduce and MapReduce factory variants (F32Reduce/F16Reduce/F32MapReduce/F16MapReduce) plus reduce_f32/reduce_f16/ map_reduce_f32/map_reduce_f16 accessors, registers max/min/sum/softmax across generic, aarch64 and x86_64, routes softmax and reduce ops in core through the registry, and removes their Ops fields and plug() assignments. plug() now touches only mmm (+ panel-extract + rms_norm); every element-wise and reduction kernel dispatches through the registry. Equivalence verified green before cutting plug. The ActivationFn enum and the `activation` module now also cover reductions, so both names are misnomers — worth renaming to a neutral `Op` / kernel registry once mmm joins.
The registry now covers pointwise activations and reductions (and soon more), so `activation`/`ActivationFn`/`ActivationImpl`/`ActFactory` were misnomers. Rename to `routines` module, `Routine` op-id enum, `RoutineImpl`, `RoutineFactory`, and the CLI subcommand `tract routines`. Mechanical rename only; no behaviour change.
Collaborator
Author
Yeah, good point, I will allow annotations for "this is generic, but we can't do best". or this is "via-f32" but we can't do best. |
Register the binary element-wise ops (by-scalar and unicast layouts of
min/max/add/mul/sub/subf) as `Routine::BinByScalar(BinOp)` / `BinUnicast(BinOp)`
with a new `RoutineFactory::Bin` variant and a `routines::bin(func, dt)` accessor,
and delete the ad-hoc `LinalgRegistry` HashMap, `register_all_*`, `bin_by_scalar`
and `bin_unicast` in favour of it (core/ops/binary.rs rewired). f16 binops stay
gated on `has_fp16()`, preserving the old fallback-to-scalar behaviour. The matrix
now exposes that binops are arm64-native only (generic on every x86 tier).
Also rename the reduce ops `Max/Min/Sum` -> `ReduceMax/ReduceMin/ReduceSum` to stay
distinct from `BinOp::{Max,Min}`, and render `tract routines` grouped by kernel
prototype (unary / scalar-param / reduction / map-reduction / binary), alphabetical
within each group.
kali
marked this pull request as draft
August 17, 2026 16:42
…rgets) Widen the arm64/arm32/wasm module gates to `any(target_arch = X, feature = registry-all-targets)` so a foreign host compiles the module shell + its `routines` descriptors (folding the `#[path]` routine bridges into `arch::routines`). All native code (asm/intrinsic kernel submodules, CPU detection, the mmm-pool builder renamed `plug_mmm`) is gated to its own arch — the kernel submodule decls in place, the rest in a single `#[cfg(target_arch = X)] mod native`. No behaviour change on a native build; on a non-native all-targets build only the descriptors compile. x86_64_fma is left native-gated for now (its interleaved kernel-macro body isn't verifiable off-x86 on this host), so its column shows when drawing on x86 but not from a non-x86 host — to be flipped once CI can certify it.
… super:: refs) Pull the kernel submodule declarations into the per-arch `mod native` alongside the detection + `plug_mmm` (via `#[path = "arm64"]` / `"arm32"` / `"wasm"` so the child `mod`s still resolve to the arch directory). This also fixes a real bug in the previous flip: kernels left at the module root referenced `super::KIND` / `super::has_neon`, which no longer resolved once those items moved into `native` — now `super` is `native`, so they resolve again. The arch root is now just the shell (`pub mod routines;` + `pub use native::*` + `mod native`). Cross-verified compiling for aarch64, armv7 and wasm32 (musl cross toolchains), alongside x86 default and all-targets builds.
Rename the `x86_64_fma` module to `x64` (files + `::` path refs; kernel symbol names like `x86_64_fma_max_f32_32n` are unchanged) and give it the same flip as the other arches: gate widened to `any(target_arch = "x86_64", registry-all-targets)`, all native code (kernel submodules, `vendor()`/feature detection, `plug_mmm`) collapsed into one `#[path = "x64"] mod native`, root reduced to the shell. The x86 routine column now shows when drawing the matrix from any host, not just x86. Cross-verified: x86 default + all-targets builds, aarch64 foreign compile, core, routines tests.
Collaborator
Author
Nope. When I incepted tract, I had a look at possible candidates for low-level op delegation, but nothing was covering the targets that mattered to us, and never really revisited. I'm a bit reluctant to include a big dependency tbh. |
Collaborator
Author
|
Not sure if I want to merge this or not. |
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
2 participants
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.