Skip to main content

Module matmul_lower

Module matmul_lower 

Source
Expand description

Genuine LOGOS → Tensix matrix-engine lowering — an AFFINE GEMM-idiom recognizer.

This is the compiler generating the matrix-engine kernel, not returning a canned template. It reads an ordinary LOGOS program — a triple-nested loop computing C[i][j] += A[i][k]*B[k][j] over row-major Seq of Int — recognizes it as a matmul by checking that the three index expressions are AFFINE in the induction variables and match the row-major GEMM layout (verified against the loop bounds), extracts the dimensions M, N, K, and lowers to a tiled Tensix kernel (matmul_tiles + tile-register management) whose tile bounds are the recognized program’s dimensions. A program that is not a matmul is REFUSED (typed), never miscompiled.

The genuineness contract, enforced by the lock: change the program’s dimensions → the emitted kernel’s tile bounds change; a non-matmul program → a typed refusal. A template can do neither.

Re-exports§

pub use crate::spatial::passes::tile::TileConfig;
pub use crate::tensor_recognize::*;

Structs§

ChipKnobs
Per-chip capability knobs — the hardware limits that make the fastest tiling chip-specific. A bigger DST register file lets a core hold a larger reuse sub-block; a bigger worker grid lets the work spread across more cores; more L1 lets wider inner blocks stay resident. The autotuner searches within these, so the SAME matmul compiles to a DIFFERENT (faster) kernel on a different part.
ComputeKnobs
The compute NUMERICS knobs of an emitted matmul program. They are part of the program’s identity: the host’s ComputeConfig carries them, the program prints them on its LOGOS_MM_PASS line, and the ledger refuses a row whose printed knobs differ from what the harness requested — so a LoFi run can never again be compared to a HiFi4 run as if they were one config.

Enums§

MathFidelity
The Tensix matrix engine’s math fidelity: how many FPU passes a bf16 multiply gets (LoFi = 1 … HiFi4 = 4). ttnn picks it by a heuristic (HiFi2 with no program config, LoFi the moment one is passed); LOGOS makes it an explicit knob, printed by the emitted program and recorded on every ledger row.
MatmulStrategy
The program STRATEGY the compiler picks for a matmul — the same decision the vendor’s op-library makes. Both are correct; they trade differently: Simple streams one tile at a time (best overlap of DRAM and compute — wins when the problem is compute/latency-bound), BlockReuse holds operand blocks in L1 and reuses them across a DST sub-block (fewer DRAM reads — wins when the problem is DRAM-bandwidth-bound, i.e. large per-core output blocks). Choosing wrong is a real perf bug: at 640³ BlockReuse measures ~12,097 device cycles vs Simple’s 9,735, so the compiler must pick Simple there.

Constants§

MAX_DST_TILES
The DST register file holds at most this many bf16 output tiles in one acquire (half-sync mode) on Wormhole; other parts differ — see ChipKnobs. Kept for the default-chip helpers + locks.
MAX_GRID
The Wormhole Tensix worker grid is 8×8; other parts differ — see ChipKnobs.

Functions§

choose_strategy
Choose the program strategy from the recognized shape + chip — the compiler making the op-library’s call. Heuristic (calibrated to the measured 640³ crossover): BlockReuse only pays once each core owns a large enough output block that saved DRAM re-reads outweigh its reduced overlap; below that, Simple’s streaming overlap wins. per_core_out ≥ 64 tiles (an 8×8-tile per-core block) is the switch.
choose_tiling
Auto-tune for the default part (Wormhole). Prefer tune_tiling with an explicit ChipKnobs to get the fastest kernel for a specific chip.
dram_read_tiles
DRAM operand-tile reads modeled for each strategy (in 32×32 tiles) — the traffic that separates the three matmul plans. Simple re-streams both operands for every output tile (no reuse). BlockReuse holds A/B blocks in L1 and reuses them across a DST sub-block, but EVERY core still reads its blocks from DRAM, so an operand tile is fetched once per grid column (A) / row (B) — redundantly. ReuseMcast reads each operand block once and NOC-multicasts it across its row/column, so every A and B tile is read from DRAM exactly once — the fewest possible. This is the mechanism behind the vendor’s win; the device-cycle confirmation is a ttsim run (traffic is necessary, not sufficient — overlap also matters). A thin facade over crate::spatial::cost::dram_tiles, which owns this arithmetic now. The NAME stays here on purpose: spatial_cost_model_lock::no_strategy_enum_input_exists greps spatial/ for it, so the computation moves and the strategy-shaped vocabulary does not follow it.
emit_matmul_vendor_dispatch
tune_tiling
Smart-lattice autotune for a named Tenstorrent part: the part’s limits, then spatial::passes::tile::best_tile_config.
tune_tiling_with
tune_tiling under explicit compute knobs: fp32 dest-accumulate halves the DST sub-block ceiling.
vendor_dispatch_blocking
Emit a vendor-dispatch program: the compiler recognized the matmul, so it routes the primitive to the target’s optimal op — ttnn.matmul on Tenstorrent — the same way a C compiler recognizes a matmul loop and calls BLAS. It passes the compiler’s OWN auto-tuned config (grid, per-core, sub-block, inner block), so the dispatched program runs at the compiler’s tuning on the vendor’s optimized implementation. For 640³ this measured 6154 device cycles — beating ttnn’s own default (6745). Generated from (spec, cfg). The secondary blocking the vendor dispatch runs with, for a recognized shape and the autotuner’s tiling: (in0_block_w, out_subblock_h, out_subblock_w). ONE definition, used by the dispatch emitter and by the measurement harness’s tiling printer, so “our config” can never mean two different things.