Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

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, a pure-Rust GEMM engine. The crate reaches gemmkit through its zero-copy 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). 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:

ItemVisibilityWhat 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 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 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 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 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 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.

The deterministic reductions in 6.3. Parallel Reductions 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 backend, as 6.2.3 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 pairDefaultControls
get_chunk_elems / set_chunk_elems33,554,432element budget for 1 row-chunk of a tiled product
get_cache_resident_max_bytes / set_cache_resident_max_bytes67,108,864cache-resident size threshold, set to your machine’s shared L3
matmul::backend::*see gemmkitevery 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 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.

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:

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:

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. See 7.3. Performance Tuning and Parallelism for a deeper treatment. For the distance kernels alongside this backend, see 6.1. Distance Metrics. For the reductions that share its parallel machinery, see 6.3. Parallel Reductions.