Part 3. Data · lesson 21 of 42

21. The matrix unit¶

The core's arithmetic has a unit built to multiply two 32 x 32 tiles as matrices in one call, and using it is the same program as the addition with a different operation. This lesson relies on lesson 20, Computing on a tile.

Why a matrix unit¶

The addition of the last lesson does one operation per number: 1024 of them for a tile. A matrix product of two 32 x 32 tiles does 32 x 32 x 32 = 32,768 multiplications and as many additions, and nearly all of the work of a neural network is products of this kind. The card has a unit that takes the two tiles and produces the product in a single call, instead of the loops you would write for 32,768 steps. That is why the data is in tiles: the unit's size is the tile.

The product of tiles is the matrix product: number (row, column) of the result is the sum, over k from 0 to 31, of number (row, k) of the left tile times number (k, column) of the right tile. It is not the product of numbers in the same place.

The general form¶

compute_kernel_hw_startup<SrcOrder::Reverse>(LEFT, RIGHT, OUTPUT);
matmul_init(LEFT, RIGHT);
...
tile_regs_acquire();
matmul_tiles(LEFT, RIGHT, 0, 0, 0);

and in Python, in km.Program:

compute = {"math_fidelity": "HiFi4"}

It is the program of lesson 20, Computing on a tile with three changes. matmul_init replaces add_init. matmul_tiles replaces add_tiles, with the same five arguments.
SrcOrder::Reverse is a setting that the matrix operation needs, and the others do not. Write it as shown.
LEFT is the buffer of the left tile of the product and RIGHT that of the right: the order matters, A x B is not B x A.
math_fidelity says how many passes the unit makes over the numbers. "HiFi4" is the most accurate and the slowest; for bfloat16 products it keeps them exact where the operands allow.

The call that replaces the addition.

One line at a time¶

compute.run([&](const km_command_t& command) {
    if (km::handler(command) != kHandlerMultiply) {
        return;
    }
    if (
        !compute.wait_front(kFirstInput, 1) || !compute.wait_front(kSecondInput, 1)
    ) {
        return;
    }
    tile_regs_acquire();
    matmul_tiles(kFirstInput, kSecondInput, 0, 0, 0);
    tile_regs_commit();
    tile_regs_wait();
    if (compute.reserve_back(kOutput, 1)) {
        pack_tile(0, kOutput);
        cb_push_back(kOutput, 1);
    }
    tile_regs_release();
    cb_pop_front(kFirstInput, 1);
    cb_pop_front(kSecondInput, 1);
});

The body of the compute kernel is that of the last lesson. The line with matmul_tiles is the only one that differs, and kHandlerMultiply is the handler for the product.

numbers = np.random.default_rng(1).integers(0, 3, (4 * 32, 32))
source = session.tensor(numbers, "bfloat16", "tile")
destination = session.tensor(np.zeros((32, 32)), "bfloat16", "tile")

The numbers are now 0, 1 and 2. Whole numbers keep the product exact: the largest possible number of the result is 32 x 2 x 2 = 128, and "bfloat16" holds whole numbers up to 256 exactly.

total = destination.read()
expected = numbers[32:64] @ numbers[64:96]
print("tile 1, first row:      ", numbers[32, :8].tolist(), "...")
print("tile 2, first column:   ", numbers[64:72, 0].tolist(), "...")
print("product, first number:  ", int(total[0, 0]))
assert (total == expected).all()
print("every number of the product is the matrix product of tile 1 and tile 2")

numpy's @ is the matrix product. The script compares all 1024 numbers of the result with it. The three prints show a row, a column and the first number of the product: the first number is the first row dotted with the first column.

The whole program¶

the_matrix_unit.pykernels/handlers.hkernels/compute.cpp

"""The matrix unit: the compute kernel multiplies two tiles as matrices."""

from pathlib import Path

import numpy as np

import komorebi as km

HERE = Path(__file__).resolve().parent

with km.Session() as session:
    numbers = np.random.default_rng(1).integers(0, 3, (4 * 32, 32))
    source = session.tensor(numbers, "bfloat16", "tile")
    destination = session.tensor(np.zeros((32, 32)), "bfloat16", "tile")

    tile = dict(format="bfloat16", page_bytes=source.page_bytes, pages=1)
    program = km.Program(
        "the_matrix_unit",
        kernels=HERE / "kernels",
        handlers=HERE / "kernels" / "handlers.h",
        buffers=[
            km.Buffer(index=0, **tile),
            km.Buffer(index=1, **tile),
            km.Buffer(index=16, **tile),
        ],
        compute={"math_fidelity": "HiFi4"},
        compile_args={
            "reader": source.accessor_args(),
            "writer": destination.accessor_args(),
        },
    )
    run = session.load(program)

    arguments = [source.address, 1, 2, destination.address, 0]
    run.wait(run.submit("multiply", group=(0, 1), args=arguments), timeout=10.0)
    run.stop()

    total = destination.read()
    expected = numbers[32:64] @ numbers[64:96]
    print("tile 1, first row:      ", numbers[32, :8].tolist(), "...")
    print("tile 2, first column:   ", numbers[64:72, 0].tolist(), "...")
    print("product, first number:  ", int(total[0, 0]))
    assert (total == expected).all()
    print("every number of the product is the matrix product of tile 1 and tile 2")

// What the kernels and Python agree on: the handler, and the numbers of the three
// buffers in L1: the two inputs of the product and its output.
#pragma once

#include <stdint.h>

constexpr uint32_t kHandlerMultiply = 1;
constexpr uint32_t kFirstInput = 0;
constexpr uint32_t kSecondInput = 1;
constexpr uint32_t kOutput = 16;

// The compute kernel: multiplies the two input tiles as matrices and puts the product
// in the output buffer.
#include "api/compute/cb_api.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/matmul.h"
#include "api/compute/pack.h"
#include "api/compute/reg_api.h"
#include "komorebi/device/km_worker.hpp"
#include "handlers.h"

void kernel_main() {
    km::Compute compute;
    compute_kernel_hw_startup<SrcOrder::Reverse>(kFirstInput, kSecondInput, kOutput);
    matmul_init(kFirstInput, kSecondInput);
    compute.run([&](const km_command_t& command) {
        if (km::handler(command) != kHandlerMultiply) {
            return;
        }
        if (
            !compute.wait_front(kFirstInput, 1) || !compute.wait_front(kSecondInput, 1)
        ) {
            return;
        }
        tile_regs_acquire();
        matmul_tiles(kFirstInput, kSecondInput, 0, 0, 0);
        tile_regs_commit();
        tile_regs_wait();
        if (compute.reserve_back(kOutput, 1)) {
            pack_tile(0, kOutput);
            cb_push_back(kOutput, 1);
        }
        tile_regs_release();
        cb_pop_front(kFirstInput, 1);
        cb_pop_front(kSecondInput, 1);
    });
}

reader.cpp and writer.cpp are those of lesson 20, with the handler renamed.

tile 1, first row:       [0, 2, 1, 0, 2, 2, 0, 1] ...
tile 2, first column:    [2, 0, 2, 0, 1, 0, 2, 1] ...
product, first number:   32
every number of the product is the matrix product of tile 1 and tile 2

Common mistakes¶

Expecting A x B to equal B x A. Swap the tile numbers in the command (2, 1 for 1, 2) and the product is a different tile.
Omitting SrcOrder::Reverse in the start-up call. The matrix operation needs it.
Numbers that grow past what "bfloat16" keeps. A product of larger whole numbers is rounded, and the numpy check fails; compare with a tolerance for such data.

Next: lesson 22, A whole map puts a group of cores to work on a whole tensor.