rustyml 0.14.0

A high-performance machine learning & deep learning library in pure Rust, offering ML algorithms and neural network support
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
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
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
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
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
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
# 6.2. Matrix Multiplication

Every dense layer forward pass and every recurrent timestep reduce to 1 operation: a matrix
product. So do linear-model predictions and the pairwise projections inside KNN and t-SNE.
RustyML does not route these through ndarray's `.dot()`.

RustyML delegates them to [`gemmkit`](https://crates.io/crates/gemmkit), a pure-Rust GEMM engine.
The crate reaches gemmkit through its zero-copy
[`gemmkit-ndarray`](https://crates.io/crates/gemmkit-ndarray) adapter. The `math` feature names
only the adapter, which brings the engine with it. RustyML keeps only a thin layer of its own code
in `src/math/matmul.rs`.

This page explains 4 things. It explains what the backend does. It explains why RustyML uses it
instead of `.dot()`. It explains how the backend picks serial versus parallel execution, and how
wide it goes when it runs parallel. It explains the only part you can change: the runtime tuning
surface under `rustyml::tuning::matmul`.

The crate's own matmul entry points are `pub(crate)`. You cannot call `dot_par` from your own
code. The estimators reach `gemmkit_ndarray` directly. They do not go through any type or
function that RustyML re-exports.

You can understand the behavior. It decides how fast your models run. You can also retune the
thresholds for your machine. If you need a matrix product in your own code, use ndarray's
`.dot()` (see [1.3. Working with ndarray](../Chapter-01/1.3._Working_with_ndarray.md)). Do not use
this backend directly.

## 6.2.1. What the backend is, and why it is internal

There are 2 layers here, and keeping them apart helps. The engine is gemmkit. It computes
`C <- alpha*A*B + beta*C` over strided views. It selects its instruction set at run time. It does
its own packing and blocking, and it owns every scheduling decision.

`gemmkit-ndarray` is a thin adapter. It reads the data pointer and the strides straight out of an
`ArrayBase<S, Ix2>`, and forwards them to the engine. It copies nothing, for a C-order view, an
F-order view, a general-stride view, or a negatively-strided view.

This adapter is already the right call-site API. So RustyML's layers and estimators call it
directly. They use `gemmkit_ndarray::dot` for an allocating product, on the backend's automatic
scheduling. They use `gemmkit_ndarray::gemm` where the caller owns the output buffer. They use
`gemmkit_ndarray::gemm_fused` where a bias and an activation ride along in the same pass.

`src/math/matmul.rs` holds, in its own words, "the crate's few additions to the gemmkit backend".
There are exactly 4 items:

| Item | Visibility | What it is |
|---|---|---|
| `dot_par(a, b, par)` | `pub(crate)` | allocating `A @ B` with an explicit `gemmkit_ndarray::Parallelism` (plain `dot` always uses the automatic default) |
| `matvec(a, x, par)` | `pub(crate)` | matvec with `Array1` operands, wraps `x` as a `[k, 1]` column, which gemmkit reroutes to its GEMV path |
| `gemm_chunk_rows(row_len)` | `pub`, `#[doc(hidden)]` | `gemm_chunk_elems() / row_len`, clamped to `[16, 4096]` rows |
| `cache_resident::<T>(rows, cols)` | `pub`, `#[doc(hidden)]` | whether `rows * cols * size_of::<T>()` is under `cache_resident_max_bytes()` |

The first 2 items are generic over `T: gemmkit_ndarray::GemmScalar`. In RustyML's build, this means
exactly `f32` and `f64`. gemmkit also supports `f16` and `bf16` under an optional `half` feature,
and `i8` under an `int8` feature. RustyML turns on neither, so there is no half precision and no
integer matmul here.

The only non-default feature RustyML does turn on is `epilogue`, requested on `gemmkit-ndarray` for
the fused path. If you need `f16`, `bf16`, or `i8` support, write your own code directly against
gemmkit.

The last 2 items are not products at all. They are the caller-side tiling policy for estimators
that would otherwise have to materialize a pairwise projection too large to hold at once. They are
technically reachable as `rustyml::math::matmul::gemm_chunk_rows` and `::cache_resident`. But
`#[doc(hidden)]` means the crate gives no stability guarantee for them. Treat them as internal, and
use the [6.2.5](#625-tuning-the-gates-the-public-surface) knobs that govern them instead.

Here is where each part gets called:

- `Dense::forward` runs a single `gemm_fused` call. It fuses the linear product, the per-column
  bias, and the ReLU activation into 1 pass.
- `Dense::backward` runs 2 plain `dot` calls. The first computes the weight gradient. The second
  computes the input gradient.
- `SimpleRNN`, `LSTM`, and `GRU` project their inputs once with `dot`. Each timestep then fuses its
  recurrent projection with `gemm_fused`. GRU drops to plain `gemm` where it writes into a slice of
  a larger buffer.
- The im2col convolution engine fuses the per-filter bias into its forward GEMM. Its 2 backward
  GEMMs route through `dot_par`. The per-item products go serial once the batch fan alone fills the
  thread pool.
- `LinearRegression`, `LogisticRegression`, `LinearSVC`, and `SVC` use `matvec` for predictions and
  gradients. So do the power iteration and the one-sided Jacobi iteration in
  `machine_learning::linalg` and in LDA. LDA also builds its scatter matrix with `dot_par`.
- PCA, kernel PCA, KMeans, and the kernel-matrix code in `machine_learning::types` use `dot`.
- KNN, t-SNE, and MeanShift use `cache_resident` and `gemm_chunk_rows`. These functions choose
  between a per-row GEMV swarm and a tiled GEMM for a pairwise projection.

You already use this backend when you call any of these models. You do not call it by name.

The products go through `pub(crate)` functions over a private dependency. You cannot call them
directly, and you should not try to work around this. RustyML re-exports only gemmkit's tuning
module. You cannot name a `Parallelism` value through the public API.

Build your layers and estimators on the public API, and you get this backend for free. Write your
own linear algebra, and you use ndarray instead. The knobs in
[6.2.5](#625-tuning-the-gates-the-public-surface) are the only public surface. They change
behavior globally, with no recompile needed.

## 6.2.2. Why not ndarray's `.dot()`

ndarray's `.dot()` in a default build uses the `matrixmultiply` crate. This is a pure-Rust GEMM,
and it works well. Use it in your own code. It is not the right choice for a training loop that
calls it millions of times across many shapes.

matrixmultiply is not a naive scalar kernel. It selects a microkernel at run time, based on the CPU
features it detects. These are FMA plus AVX2, AVX, or SSE2 on x86-64, and NEON on aarch64. The
claim that gemmkit vectorizes while `.dot()` does not is false. The real difference comes from 3 things: the
shape-specialized routes, the fused epilogue, and threading. ndarray does not enable
`matrixmultiply`'s optional threading feature, so `.dot()` in this build runs on 1 thread only.

**Shape-specialized routes.** gemmkit does not run a single blocked algorithm. It picks between
several routes by shape. These include a dedicated matrix-vector path, an in-place path for a
shallow `k` that skips packing, and another path for a small `m` and `n`. A single general kernel
computes these shapes correctly, but slowly. A training loop is full of exactly these shapes.

**The fused epilogue.** `gemm_fused` applies the per-column bias and the activation inside the
kernel, while the output tile still sits in registers. `.dot()` has no such feature. The same
`Dense::forward`, written against `.dot()`, needs 3 passes over the output: the product, the bias,
and the activation. This path needs only 1 pass. [6.2.4](#624-determinism-and-reproducibility)
records the guarantee that makes this safe: the fused result is bit for bit equal to the unfused
sequence.

**Threading.** `matrixmultiply`'s `threading` feature sits behind ndarray's own opt-in
`matrixmultiply-threading` feature, which RustyML does not turn on. gemmkit threads on its own, and
it decides for itself whether threading is worth it.
[6.2.3](#623-how-gemmkit-schedules-a-product) covers that decision in full.

Operand strides pass straight through to the kernel. This is a convenience, not an advantage over
`.dot()`, since `.dot()` also handles strides well. The adapter accepts any `ArrayBase<S, Ix2>`
with `S: Data`. This includes an owned `Array2`, an `ArrayView2`, a transpose (`a.t()`), a
non-contiguous slice, and even a negatively-strided view. Nothing gets copied or physically
transposed first. A transposed view is just a swapped pair of strides, and the engine reads
arbitrary strides directly.

This matters because the backward pass is full of transposed operands. `dot(&input.t(),
&grad_upstream)` is the weight-gradient pattern in `Dense`. A `.dot()`-based path would need to
copy these into contiguous buffers, or lose the fused stride handling. The tests in
`src/math/matmul.rs` confirm this. A `.t()` operand and an `s![..;2, ..]` row-strided slice both
feed the correct strides to the kernel. Both match an independent reference product.

The old hand-rolled backend could not do 2 things that gemmkit now does. The fused epilogue is the
first of these. The second is that gemmkit's blocking and job order do not depend on the worker
count. This independence turns the reproducibility statement in
[6.2.4](#624-determinism-and-reproducibility) from a hedge into a promise.

## 6.2.3. How gemmkit schedules a product

RustyML makes exactly 1 scheduling decision per call site, and this decision is binary. It passes
`Parallelism::Rayon(0)`, which means "decide for yourself". Or it passes `Parallelism::Serial`,
which means "this thread is already inside a rayon region, so do not fork again". The second form
matters. The convolution engine's backward pass and MeanShift's seed loop both use it. It is about
not forking twice, and never about correctness.

Everything past that choice belongs to gemmkit. This covers serial versus parallel, the worker
count, the pool the work runs in, and whether the shape even takes a bandwidth-bound route.

A gemmkit knob resolves in priority order. A per-call argument, such as the `Parallelism` request,
beats a programmatic `set_*` call. A `set_*` call beats a `GEMMKIT_*` environment variable. An
environment variable beats the compiled default. Each environment variable is read once, on the
knob's first access, and then cached for the rest of the process.

A `set_*` call stores its value unconditionally. Once anything in the process calls a setter, the
matching environment variable has no effect for the rest of the run. This is why RustyML never
calls a setter on your behalf. A `GEMMKIT_*` value that fails to parse as a non-negative integer
warns once on stderr, and then falls back to the compiled default. A typo in a performance profile
never crashes the process.

**The work gate.** `parallel_threshold` is the serial-versus-parallel crossover. Its default is
`48 * 48 * 256`, or 589,824. This gate compares the `m * n * k` product, not FLOPs, so there is no
factor of 2 anywhere. Read the units carefully, because an earlier version of this page compared
FLOPs instead. A problem below the gate runs on 1 thread, no matter how many workers you requested.

This band holds the tiny GEMMs: RNN and LSTM timesteps, and small dense layers called in tight
loops. Keeping them serial is the correct choice, not laziness. Dispatching work onto a thread pool
costs more than the multiply itself.

**The worker ramp.** Above the gate, the automatic path does not grab every core at once.
`par_mnk_per_worker` defaults to 2,000,000 on native targets. It sets how much extra `m * n * k`
work each additional worker needs before the product widens by 1. The target worker count is
`mnk / par_mnk_per_worker`, floored at 1 and capped by the core count and the job count.

The ramp is based on work, not on dimension, because the measured optimum tracks total work rather
than linear size. gemmkit's own calibration, from a Ryzen 9950X, shows this. A `128^3` product
(about 2e6) runs fastest serial. A `192^3` product (about 7e6) wants 2 or 3 workers. A `384^3`
product (about 5.7e7) already wants all 32 hardware threads. No single stride along 1 dimension
fits both ends of this curve.

**The pool tiers.** An earlier version of this page said the backend kept no thread pool of its
own. That is no longer true. `pool_classes` builds persistent, exact-fit private rayon pools in
tiers. The tiers halve down from half the machine width: 1 tier is width/2, 2 tiers add width/4,
and 3 tiers add width/8. The automatic worker count snaps to the smallest tier that still holds it.

The reason is rayon's fork-join tax. This tax scales with a pool's slack, which is its width minus
the workers actually doing work, not with the worker count alone. 8 workers in an 8-wide pool beat
the same 8 workers inside a 32-wide global pool, by a wide margin.

The tier pools are built once and reused warm. They are not rebuilt per call. A value of `0`
disables them entirely. The default is arch-split: 2 tiers on x86-64, 1 tier on aarch64, and 0
tiers on every other target, pending on-device validation.

If the calling thread is already a rayon worker, for example inside a nested GEMM or your own
installed pool, gemmkit skips the tiers. It runs inside the current pool instead. This is why these
products compose cleanly inside an outer parallel region. They do not stack a second pool on top of
yours.

**Matvecs are their own cost class.** gemmkit detects the `m == 1` or `n == 1` shape and takes a
dedicated, bandwidth-bound path instead of the general driver. `matmul::matvec` exists to present
an `Array1` as the `[k, 1]` column that triggers this path. This path does not consult
`parallel_threshold` at all.

It stays serial below a byte floor, `gemv_parallel_bytes`, which defaults to `0` (meaning "derive
it from the cache size"). The derived floor is 1 core's private L2. Below it, the touched data is
L2-resident, that core already sees the full L2 bandwidth, and splitting the work only adds
fork-join overhead with no DRAM bandwidth to win back.

Above the floor, the worker count climbs a ladder as the touched bytes grow. Its rungs are the
same exact-fit pool tiers the general driver uses, and it climbs 1 tier per `gemv_tier_step`
factor of bytes above the floor. A matvec that only just clears the floor therefore gets the
narrowest tier, not the full memory-parallel width. `gemv_thread_cap` overrides the ladder: a
non-zero value is the width verbatim, pinned flat at every size. Both default to `0` for auto
mode.

`gemv_axpy_par_min_rows` adds a shape-specific guard on top. A column-major matvec keeps its rows
on 1 worker below that output-row count, because the output-row axis is the inner memory axis
there, and cutting it gives every worker a strided walk over the whole matrix. A row-major matrix
is unaffected, because its workers own whole `k`-contiguous rows. RustyML's operands are
row-major, so `matvec` never consults this guard.

A last knob, `gemv_threshold`, caps how large the vector side may be before the shape falls back
to the general driver. Its default is `usize::MAX - 1`, effectively unbounded. In practice, a
gemv-shaped problem always takes the gemv path, unless you lower this knob yourself.

A dozen more knobs sit behind these: `kc`, `rhs_pack_threshold`, the `lhs_pack_*` family,
`small_k_threshold`, `small_mn_dim`, `prefetch_min_bytes`, and others. This page does not list
them, because such a table would go out of date quickly. They are documented on gemmkit's own
docs.rs page. Each knob is reachable through `rustyml::tuning::matmul::backend`. The
`gemmkit-tune` autotuner sweeps them for you, on your target machine.

This note applies to all of them. gemmkit's reference machines are a Ryzen 9950X (x86-64) and an M4
Max (aarch64). Any knob whose crossover depends on architecture carries a separate default for
each, split by `cfg(target_arch)`. Unless stated otherwise, the numbers quoted on this page are the
x86-64 values.

## 6.2.4. Determinism and reproducibility

An earlier version of this page said results were reproducible on the same machine, but not
necessarily bit for bit. That claim no longer holds, and you should discard it.

The old hedge existed because the crate's own row-split wrapper gave each block a different `m`.
The kernel's internal `k`-blocking depended on `m`, so the summation order moved with the thread
count. That row split is gone, and the hedge went with it. `src/math/matmul.rs` now documents a
direct promise:

> gemmkit's blocking and job order do not depend on the worker count. For a fixed machine and
> configuration, the same product reproduces the same result bit for bit, no matter how many
> threads ran it. The result also repeats from run to run. Fused epilogues (bias and activation)
> are bitwise identical to the plain product followed by the same scalar map.

This is not aspirational. The module's own test suite checks every part of it.

- `dot_par_thread_count_independent_f64` runs a `96^3` shape, a `256 x 64 x 64` shape, and a
  thin-`k` `64 x 8192 x 64` shape. It runs each shape serially, then at `Rayon(2)`, `Rayon(4)`,
  `Rayon(8)`, `Rayon(16)`, and `Rayon(32)`. It asserts `to_bits()` equality across every arm. The
  thin-`k` shape is there because it is the shape most likely to tempt a split-`k` reduction, which
  would break this property.
- `dot_par_thread_count_independent_f32` runs the same check for `f32`.
- `matvec_serial_and_auto_agree_bitwise` covers the bandwidth-bound gemv path. Each output element
  there is reduced over the whole of `k` on 1 worker.
- `dot_run_to_run_deterministic` and `matvec_run_to_run_deterministic` cover repeat calls on the
  same machine.
- `gemm_fused_bias_relu_bitwise_matches_unfused` checks that `gemm_fused`, with a `Bias::PerCol`
  and an `Activation::Relu`, equals a plain `dot` followed by the same scalar add-and-clamp, bit
  for bit. This is what makes fusing the bias and the ReLU into a `Dense` forward pass a free
  optimization, not a numerical trade-off.

The words **fixed machine and configuration** still carry weight. A different CPU picks a
different SIMD width, and so a different accumulation layout. Changing a knob can also change the
blocking. Cross-machine bit-equality is still not promised, and no threaded BLAS promises it
either.

Within 1 binary on 1 machine, though, the worker count is no longer a variable you must reason
about. For a training run you want to replay later, that is the part that matters. For the seeding
side of reproducibility, such as weight initialization, shuffles, and dropout masks, see
[7.1. Reproducibility and Random Seeds](../Chapter-07/7.1._Reproducibility_and_Random_Seeds.md).

The deterministic reductions in [6.3. Parallel Reductions](./6.3._Parallel_Reductions.md) give a
stricter guarantee. They produce the same result by construction, independent of the machine, not
merely independent of the worker count.

## 6.2.5. Tuning the gates: the public surface

This is the part you can call directly. It has 2 layers. The serial-versus-parallel decision
belongs to the [`gemmkit`](https://crates.io/crates/gemmkit) backend, as
[6.2.3](#623-how-gemmkit-schedules-a-product) describes. The per-dtype FLOPs gates that the crate
used to hand-roll and expose are gone.

`rustyml::tuning::matmul` is available with the `math` feature, and so under `full`. It still owns
the caller-side tiling policy. It also re-exports the backend's own knobs, so you never need a
direct `gemmkit` dependency.

The re-export goes through `gemmkit-ndarray`, the adapter RustyML actually calls, and not through a
`gemmkit` dependency of its own. This matters if you add `gemmkit` to your own `Cargo.toml`
regardless. The knobs are process-global atomics, so cargo resolving your `gemmkit` to a different
version than the adapter's would give you a second copy, and a `set_*` call on it would have no
effect on RustyML's products. Going through `rustyml::tuning::matmul::backend` cannot land on the
wrong copy.

| Function pair | Default | Controls |
|---|---|---|
| `get_chunk_elems` / `set_chunk_elems` | 33,554,432 | element budget for 1 row-chunk of a tiled product |
| `get_cache_resident_max_bytes` / `set_cache_resident_max_bytes` | 67,108,864 | cache-resident size threshold, set to your machine's shared L3 |
| `matmul::backend::*` | see gemmkit | every backend knob, each with a matching `GEMMKIT_*` environment variable |

`cache_resident_max_bytes` is the knob you are most likely to change. Set it to your actual shared
L3 size. The default, 64 MiB, is a guess, and the band around it is not calibrated.

For the serial-versus-parallel crossover, use `matmul::backend`. `set_parallel_threshold` gates on
the `m * n * k` product. `set_gemv_threshold` gates the matvec path. Every backend knob also reads
from a `GEMMKIT_*` environment variable. The `gemmkit-tune` autotuner can emit a full machine
profile, so you rarely need to pick numbers by hand.

Set these knobs once at startup, before the hot loop starts. They are global, and they apply to
the whole process. Calling a backend `set_*` function from RustyML silences the matching
`GEMMKIT_*` environment variable, for the rest of the process. This would override a profile you
had set through the environment. That is why RustyML never sets these knobs for you.

See [7.3. Performance Tuning and Parallelism](../Chapter-07/7.3._Performance_Tuning_and_Parallelism.md)
for the full rationale, the calibration workflow, and how these knobs compose with the reduction
and elementwise gates.

## 6.2.6. When parallelism pays, and how to measure it

The gates encode where threading helps. The shape sweep in `benches/benchmarks/matmul_kernels.rs`
shows this. Run it with `cargo bench --bench matmul_kernels`. `Dense::forward` is 1 fused GEMM call
and nothing else. The bias and the activation run inside the kernel's epilogue, not as extra
passes. The sweep times 6 shapes, labeled `batch x in_features x out_features`, which is
`m x k x n`:

- The 4 square-ish rungs are `small_256x256x256`, `medium_512x1024x1024`, `big_1024x2048x2048`, and
  `huge_2048x2048x2048`. They walk the worker ramp from end to end. Even the smallest is
  `16,777,216` in `m*n*k`, about 28 times the work gate, so none of them is a serial case. This
  ladder shows the ramp handing out more workers as the work grows. It also shows the pool tiers
  dropping out at the top, once a problem is large enough to want the full machine width.
- `wide_256x256x8192` is the wide-`n` case. There is an abundance of independent output columns
  here, so the work splits with no trouble at all. This is the easy shape for any threaded GEMM.
- `thin_256x8192x256` is the interesting shape. Its name means thin in `k`'s neighbors, not thin
  overall: `m` and `n` are both 256, while `k` is 8192. This is a deep-`k` product. A common
  intuition says a skinny shape must be bandwidth-bound, but this shape is firmly compute-bound:
  about 1.07 GFLOP against about 17 MB of operands. It is the shape where depth-blocking decisions
  matter most, which is why the sweep includes it.

2 regimes stay outside this bench on purpose.

A genuine matvec never appears in it, because a `Dense` forward is never one. A matvec leaves the
general driver entirely, for gemmkit's gemv path. It gates instead on a byte floor derived from 1
core's private L2, then climbs a worker ladder as the bytes it touches grow, because DRAM saturates
at far fewer workers than a machine has logical cores. This path is bandwidth-bound, so extra cores
start to pay off far earlier there than in a compute-bound GEMM, which can better amortize thread
dispatch.

Sub-gate products are also absent: RNN and LSTM timesteps, and small dense layers. The smallest
shape in this sweep already sits well above the gate. If you profile an RNN and see rayon overhead
dominating, do not lower `parallel_threshold`. Those products stay serial by design, and the
overhead comes from somewhere else.

You can compare serial and parallel execution on your own machine, for a fixed product, by
toggling the gate around it. The example below exercises the backend through a public `Dense`
layer, and times the same product both ways. Treat the printed numbers as a sketch, not a
benchmark, since a single call is noisy. For real figures, use the criterion bench above, which
warms up and repeats.

```rust
use ndarray::Array;
use rustyml::neural_network::layers::{Activation, Dense};
use rustyml::neural_network::traits::Layer;
use rustyml::tuning::matmul;
use std::time::Instant;

fn main() {
    // A Dense forward is 1 backend GEMM: input (batch, in_features) @ weights (in_features, units).
    let (batch, fin, fout) = (256usize, 256usize, 256usize);
    let mut layer = Dense::new(fin, fout, Activation::ReLU)
        .unwrap()
        .with_random_state(42);
    let x = Array::from_elem((batch, fin), 0.5f32).into_dyn();

    // The backend gates on the m*n*k product, not on FLOPs. There is no factor of 2.
    let work = batch * fin * fout;
    println!(
        "backend parallel gate = {}; this product = {} (parallel: {})",
        matmul::backend::parallel_threshold(),
        work,
        work >= matmul::backend::parallel_threshold()
    );

    let warm = layer.forward(&x).unwrap();
    assert_eq!(warm.shape(), &[batch, fout]);

    // Force this exact product serial by lifting the gate just above its work count.
    let saved = matmul::backend::parallel_threshold();
    matmul::backend::set_parallel_threshold(work + 1);
    let t0 = Instant::now();
    for _ in 0..20 {
        let _ = layer.forward(&x).unwrap();
    }
    let serial = t0.elapsed() / 20;

    // Restore the gate so the same product now takes the parallel strategy.
    matmul::backend::set_parallel_threshold(saved);
    let t1 = Instant::now();
    for _ in 0..20 {
        let _ = layer.forward(&x).unwrap();
    }
    let parallel = t1.elapsed() / 20;

    println!("serial  ~ {serial:?} / forward");
    println!("parallel ~ {parallel:?} / forward");
}
```

The gate and the work count are fixed by the defaults and the shape, so those print exactly. The
timings are machine-dependent, so the output below shows their shape and kind rather than numbers:

```text
backend parallel gate = 589824; this product = 16777216 (parallel: true)
serial  ~ <duration> / forward
parallel ~ <duration> / forward
```

Watch the second call to `set_parallel_threshold`, which restores the saved value. A programmatic
setter permanently shadows the matching `GEMMKIT_PARALLEL_THRESHOLD` environment variable. So a
snippet like this pins the knob in code for the rest of the process, even after it "restores" it.
This is harmless here, because the restored value is the one the process started with. It is still
a reason not to scatter setters through a library.

Do not be surprised if the parallel run is not faster, on a small product like this or on a machine
with few cores. That is the whole point of the gate, and it is why the defaults keep sub-gate
products serial. Scale `batch`, `fin`, and `fout` up to the benchmark's larger shapes, and the
parallel arm pulls ahead.

If you want to write your own matrix product instead of routing through a layer, use ndarray. This
is deliberately outside the backend:

```rust
use ndarray::array;

fn main() {
    // RustyML's matmul entry points are crate-internal. Your own matmul code uses ndarray's `.dot()`.
    let a = array![[1.0_f64, 2.0, 3.0], [4.0, 5.0, 6.0]]; // 2x3
    let b = array![[1.0_f64, 0.0], [0.0, 1.0], [1.0, 1.0]]; // 3x2
    let c = a.dot(&b); // 2x2
    assert_eq!(c, array![[4.0, 5.0], [10.0, 11.0]]);
    println!("A.dot(B) shape = {:?}", c.shape());
}
```

Build your models on the public layers and estimators, and you get gemmkit for free. This includes
its runtime ISA dispatch, its work-based scheduling and pool tiers, its fused epilogues, and its
worker-count-independent numerics. Nothing needs configuration.

When a specific machine wants different crossovers, use the gates in
[6.2.5](#625-tuning-the-gates-the-public-surface). See
[7.3. Performance Tuning and Parallelism](../Chapter-07/7.3._Performance_Tuning_and_Parallelism.md)
for a deeper treatment. For the distance kernels alongside this backend, see
[6.1. Distance Metrics](./6.1._Distance_Metrics.md). For the reductions that share its parallel
machinery, see [6.3. Parallel Reductions](./6.3._Parallel_Reductions.md).