n150 n300 T3000 p100 p150 p300c Galaxy Sim 45 min Draft

Module 8: Matrix Math and Matmul Labs

Introduction: The Kernel Pattern Behind Modern AI

C = A × B looks simple. In practice, matrix multiply is the dominant kernel in transformer inference and training.

This module ports the full Lab 1 → Lab 2 → Lab 3 arc from TT-Metalium™ docs into the interactive lesson style used by the rest of this collection.

You will move from:

  1. One core + tiled matmul (correctness and dataflow)
  2. Many cores + split work (parallel decomposition)
  3. Multicast + reuse (network-aware optimization)

By the end, you should be able to explain not just what command to run, but why each optimization exists.

What You'll Learn

Upstream Labs (Source of Truth)

Use these as the canonical references while you work through this module:

Part 1: Concept Primer (Before You Run Anything)

Matmul refresher

For A[M×K], B[K×N], output C[M×N]:

C[i, j] = Σ (A[i, k] * B[k, j]) for k in [0, K)

A naive kernel streams every operand from DRAM repeatedly. The labs show why that is too expensive and how to move reuse closer to compute.

Why tiles matter

In these labs, matmul is expressed as tile operations (commonly 32×32):

Think of each output tile as an independent unit of work that can be assigned to one core (Lab 1) or many cores (Lab 2/3).

Dataflow mental model

DRAM A/B tiles -> Reader kernel -> Circular buffers in L1
                          -> Compute kernel -> Accumulator tiles
                          -> Writer kernel -> DRAM C tiles

In Lab 3, the first arrow is optimized further with NoC multicast so multiple cores can consume one producer's data movement.

Part 2: Environment Setup (Hardware or Simulator)

If ~/tt-metal is missing (common on TT-QuietBox® 2 preconfigured images), start here first:

🛠️ Build TT-Metalium from Source

Build programming examples used by this module:

cd ~/tt-metal
./build_metal.sh --build-programming-examples

🔨 Build Programming Examples
VS Code
cd ~/tt-metal && ./build_metal.sh --build-programming-examples

Optional simulator path (hardware-free)

🧪 Set Up ttsim Simulator
VS Code
mkdir -p ~/sim
wget -q https://github.com/tenstorrent/ttsim/releases/download/v1.8.4/libttsim_wh.so -O ~/sim/libttsim_wh.so || { echo "ERROR: failed to download libttsim_wh.so"; exit 1; }
wget -q https://github.com/tenstorrent/ttsim/releases/download/v1.8.4/libttsim_bh.so -O ~/sim/libttsim_bh.so || { echo "ERROR: failed to download libttsim_bh.so"; exit 1; }
wget -q https://github.com/tenstorrent/ttsim/releases/download/v1.8.4/libttsim_wh_x2.so -O ~/sim/libttsim_wh_x2.so || { echo "ERROR: failed to download libttsim_wh_x2.so"; exit 1; }
wget -q https://github.com/tenstorrent/ttsim/releases/download/v1.8.4/libttsim_bh_x2.so -O ~/sim/libttsim_bh_x2.so || { echo "ERROR: failed to download libttsim_bh_x2.so"; exit 1; }
wget -q https://github.com/tenstorrent/ttsim/releases/download/v1.8.4/libttsim_wh_x8.so -O ~/sim/libttsim_wh_x8.so || { echo "ERROR: failed to download libttsim_wh_x8.so"; exit 1; }
if [ -n "$TT_METAL_HOME" ]; then
  cp $TT_METAL_HOME/tt_metal/soc_descriptors/wormhole_b0_80_arch.yaml ~/sim/soc_descriptor.yaml || { echo "ERROR: failed to copy SOC descriptor"; exit 1; }
  cp $TT_METAL_HOME/tests/tt_metal/tt_fabric/custom_mock_cluster_descriptors/n300_cluster_desc.yaml ~/sim/n300_cluster_desc.yaml || { echo "WARNING: n300 cluster desc copy skipped (optional for N300 sim)"; }
else
  echo "TT_METAL_HOME not set — SOC descriptor copy skipped"
fi
echo "ttsim v1.8.4 ready (wh + bh + wh_x2 + bh_x2 + wh_x8)"

# Wormhole simulator
export TT_METAL_SIMULATOR=$HOME/sim/libttsim_wh.so

# Blackhole simulator (P-series)
# export TT_METAL_SIMULATOR=$HOME/sim/libttsim_bh.so

# Required for many simulator workflows
export TT_METAL_SLOW_DISPATCH_MODE=1

cd ~/tt-metal

More simulator workflows: Twenty-and-Ten Things You Can Do with ttsim

Part 3: Lab 1 Deep Dive — Single-Core Tiled Matmul

Goal of Lab 1

Build intuition for tiled data movement and kernel decomposition on a single core.

Run it

cd ~/tt-metal
./build/programming_examples/metal_example_matmul_single_core

What success looks like

Metalium vs Golden -- PCC = ...
Test Passed

Code walkthrough: exact host program path

Source file:

uint32_t Mt = M / TILE_HEIGHT;
uint32_t Kt = K / TILE_WIDTH;
uint32_t Nt = N / TILE_WIDTH;

auto reader_id = tt_metal::CreateKernel(
    program,
    "matmul/matmul_single_core/kernels/dataflow/reader_single_core_mm.cpp",
    core,
    tt_metal::DataMovementConfig{...});

auto writer_id = tt_metal::CreateKernel(
    program,
    "matmul/matmul_single_core/kernels/dataflow/writer_single_core_mm.cpp",
    core,
    tt_metal::DataMovementConfig{...});

tt_metal::CreateKernel(
    program,
    "matmul/matmul_single_core/kernels/compute/mm.cpp",
    core,
    tt_metal::ComputeConfig{.math_fidelity = math_fidelity, .compile_args = {Mt, Kt, Nt}});

tt_metal::SetRuntimeArgs(program, reader_id, core, {src0_addr, src1_addr, Mt, Kt, Nt});
tt_metal::SetRuntimeArgs(program, writer_id, core, {dst_addr, Mt, Kt, Nt});

Read this snippet line-by-line (iota-by-iota):

  1. Mt/Kt/Nt converts element dimensions into tile dimensions.
  2. Reader/writer kernels are distinct dataflow kernels.
  3. Compute kernel receives tile-shape constants at compile time ({Mt, Kt, Nt}).
  4. Reader runtime args bind input addresses + tile extents for this launch.
  5. Writer runtime args bind output address + tile extents for this launch.

What to inspect while reading Lab 1 code

Why this matters

Lab 1 is the correctness baseline. If you cannot reason about tile lifecycle here, Lab 2 and Lab 3 optimizations become opaque.

Checkpoint questions

Part 4: Lab 2 Deep Dive — Multi-Core Work Distribution

Goal of Lab 2

Scale the same tiled algorithm across many cores by partitioning output tiles.

Run baseline multi-core

cd ~/tt-metal
./build/programming_examples/metal_example_matmul_multi_core

Run reuse-oriented variant

cd ~/tt-metal
./build/programming_examples/metal_example_matmul_multicore_reuse

Core idea

split_work_to_cores(...) (or equivalent decomposition logic) partitions output tiles so each core receives a contiguous share of the output space.

Code walkthrough: work split + per-core args

Source file:

auto [num_cores, all_cores, core_group_1, core_group_2, work_per_core1, work_per_core2] =
    split_work_to_cores(core_grid, num_output_tiles_total);

for (const auto& [ranges, work_per_core] : work_groups) {
    for (const auto& range : ranges.ranges()) {
        for (const auto& core : range) {
            tt_metal::SetRuntimeArgs(program, reader_id, core, {
                src0_dram_buffer->address(), src1_dram_buffer->address(),
                Mt, Kt, Nt, work_offset, work_per_core});
            tt_metal::SetRuntimeArgs(program, writer_id, core, {
                dst_dram_buffer->address(), work_per_core, work_offset});
            tt_metal::SetRuntimeArgs(program, compute_kernel_id, core, {
                work_per_core, Kt});
            work_offset += work_per_core;
        }
    }
}

Read this snippet line-by-line (iota-by-iota):

  1. split_work_to_cores forms two groups when work is uneven.
  2. The host loops every core in both groups and writes core-specific runtime args.
  3. work_offset is the start tile index for that core's output region.
  4. work_per_core is the tile count assigned to that core.
  5. Reader + writer + compute all consume the same partition contract.

What to inspect in Lab 2

Performance intuition

Lab 2 often improves throughput because:

  1. More cores compute output tiles concurrently
  2. Per-core work can stay local in L1 buffers
  3. The kernel pipeline overlaps data movement and compute better

But scaling is never free: load imbalance and DRAM contention can flatten speedup.

Checkpoint questions

Part 5: Lab 3 Deep Dive — Multicast Data Reuse

Goal of Lab 3

Reduce redundant memory traffic by sharing operand tiles across cores through NoC multicast rather than issuing repeated DRAM reads.

Conceptual flow

Producer core reads source tile once
-> multicasts tile to receiver core set
-> semaphore/handshake guarantees readiness
-> receivers consume tile for compute

Toolkit runnable multicast check

cd ~/tt-metal
if [ -x ./build/programming_examples/contributed/multicast ]; then
  ./build/programming_examples/contributed/multicast
else
  echo "multicast example not found yet. Rebuild with ./build_metal.sh --build-programming-examples and follow Lab 3 source walkthrough."
fi

Code walkthrough: sender/receiver synchronization

Source file:

CoreCoord sender_core_device = prog_state.mesh_device->worker_core_from_logical_core(sender_core_logical);
uint32_t receivers_ready_semaphore = CreateSemaphore(prog_state.program, all_cores_logical, 0);
uint32_t tile_sent_semaphore = CreateSemaphore(prog_state.program, all_cores_logical, INVALID);

SetRuntimeArgs(
    prog_state.program, mcast_sender_id, sender_core_logical,
    {receiver_cores_device.start_coord.x, receiver_cores_device.start_coord.y,
     receiver_cores_device.end_coord.x, receiver_cores_device.end_coord.y,
     receivers_ready_semaphore, tile_sent_semaphore, src_mesh_buffer.address(), n_tiles, num_dests});

SetRuntimeArgs(
    prog_state.program, mcast_receiver_id, receiver_cores_logical,
    {sender_core_device.x, sender_core_device.y, receivers_ready_semaphore, tile_sent_semaphore, n_tiles});

Read this snippet line-by-line (iota-by-iota):

  1. Host converts logical → device coords before NoC-facing runtime args.
  2. Two semaphores coordinate readiness and "tile sent" signaling.
  3. Sender runtime args include the multicast destination rectangle.
  4. Receiver runtime args include sender device coordinates + semaphore IDs.
  5. Same semaphore IDs are shared across sender/receiver kernels for handshake correctness.

What to inspect in Lab 3

Why this matters for transformers

Attention and MLP blocks repeatedly reuse shared weights/activations. Multicast patterns directly target those reuse opportunities and reduce bandwidth pressure.

Checkpoint questions

Part 6: Guided RST-to-Interactive Study Plan

Use this sequence to get the "full thrust" of the original labs while staying hands-on:

  1. Read Lab 1 objective + run single-core binary
  2. Trace one output tile end-to-end (reader -> compute -> writer)
  3. Read Lab 2 partition section + run multi-core binaries
  4. Map output tile ranges to cores and confirm balance assumptions
  5. Read Lab 3 multicast section + run multicast example / source walk-through
  6. Sketch traffic difference: repeated DRAM reads vs one read + multicast fanout

If you're teaching this module, require students to answer all checkpoint questions before advancing.

Part 7: Troubleshooting and Validation Signals

Build or binary missing

cd ~/tt-metal
./build_metal.sh --build-programming-examples

Simulator runs are very slow

That is expected with TT_METAL_SLOW_DISPATCH_MODE=1. Use it for understanding correctness and dataflow, not throughput benchmarking.

Numerical mismatch or failed validation

Repro checklist

Upstream drift check (tt-metal main)

Use this test to detect drift between this lesson's tracked matmul upstream sources and tenstorrent/tt-metal main:

npm run check:matmul-lesson-drift

This reports:

Part 8: Mastery Checklist

Before moving on, confirm you can explain:

What's Next?

→ Continue to Bounty Program: Model Bring-Up