Part 3. Data · lesson 20 of 42

20. Computing on a tile

The compute kernel takes tiles from buffers, lets the core's arithmetic units work on them, and puts the result tile in another buffer, in three steps: unpack, one operation, pack. This lesson relies on lesson 19, Circular buffers, lesson 18, The three kernels and lesson 15, Tiles and tensors.

Why three steps

The arithmetic units do not read L1 themselves. Of the three CPUs that run the compute kernel (lesson 3, The card), the first, unpack, moves tiles from buffers into the units' working registers. The second, math, runs the operation on what is there. The third, pack, moves the result from the registers back into a buffer. You write one function for all three; each CPU runs the part that is its own. Between them they pass the tile through the destination register, a set of working tiles inside the unit, numbered from 0.

Two tiles in, one tile out: unpack and pack touch the buffers; math works on the register between them.

The general form

compute_kernel_hw_startup(FIRST, SECOND, OUTPUT);
add_init(FIRST, SECOND);
compute.run([&](const km_command_t& command) {
    compute.wait_front(FIRST, 1);
    compute.wait_front(SECOND, 1);
    tile_regs_acquire();
    add_tiles(FIRST, SECOND, 0, 0, 0);
    tile_regs_commit();
    tile_regs_wait();
    compute.reserve_back(OUTPUT, 1);
    pack_tile(0, OUTPUT);
    cb_push_back(OUTPUT, 1);
    tile_regs_release();
    cb_pop_front(FIRST, 1);
    cb_pop_front(SECOND, 1);
});

• FIRST, SECOND and OUTPUT are the indexes of the three buffers: two to read, one to write.
• compute_kernel_hw_startup and add_init are set-up, done once, before run. They tell the units which buffers and which operation to expect.
• wait_front waits until reader.cpp has pushed a tile into the buffer, as writer.cpp did in lesson 18, The three kernels.
• tile_regs_acquire() takes the destination register for this kernel. tile_regs_commit() says "math is done"; tile_regs_wait() is where pack waits for it. tile_regs_release() gives the register back.
• add_tiles(FIRST, SECOND, 0, 0, 0) adds the first tile of buffer FIRST to the first tile of SECOND and leaves the sum in register 0. The last number is the register; the two before it are which tile of each buffer.
• compute.reserve_back and cb_push_back are reader.cpp's reserve and push on the output side; pack_tile(0, OUTPUT) copies register 0 into the tile just reserved.
• Every tile taken with wait_front is released with cb_pop_front.

One line at a time

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

The numbers are whole numbers from -8 to 8, in four tiles, stored as "bfloat16", a 2-byte number with a fraction. Whole numbers this small are exact in it, so the sum can be checked exactly. The destination is one tile of zeros.

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

Three buffers of one tile: 0 and 1 for the two inputs, 16 for the output. dict(...) holds what the three have in common.

reader.run([&](const km_command_t& command) {
    if (km::handler(command) == kHandlerAdd) {
        const uint32_t source_address = command.a0;
        const uint32_t first_tile = command.a1;
        const uint32_t second_tile = command.a2;
        const auto source = TensorAccessor(source_arguments, source_address);
        if (!reader.reserve(kFirstInput, 1) || !reader.reserve(kSecondInput, 1)) {
            return;
        }
        noc_async_read_page(first_tile, source, get_write_ptr(kFirstInput));
        noc_async_read_page(second_tile, source, get_write_ptr(kSecondInput));
        noc_async_read_barrier();
        cb_push_back(kFirstInput, 1);
        cb_push_back(kSecondInput, 1);
    }
});

reader.cpp puts one tile of the source in each of the two input buffers.

compute.run([&](const km_command_t& command) {
    if (km::handler(command) != kHandlerAdd) {
        return;
    }
    if (
        !compute.wait_front(kFirstInput, 1) || !compute.wait_front(kSecondInput, 1)
    ) {
        return;
    }
    tile_regs_acquire();
    add_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 compute kernel is the general form with the checks: if a wait gives up, it returns.

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

One command: add tile 1 and tile 2 of the source, write the result as tile 0 of the destination. writer.cpp of this lesson is that of lesson 17, Writing to DRAM, reading from buffer 16.

The whole program

computing_on_a_tile.pykernels/handlers.hkernels/reader.cppkernels/compute.cppkernels/writer.cpp

"""Computing on a tile: the compute kernel adds two tiles."""

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(-8, 9, (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(
        "computing_on_a_tile",
        kernels=HERE / "kernels",
        handlers=HERE / "kernels" / "handlers.h",
        buffers=[
            km.Buffer(index=0, **tile),
            km.Buffer(index=1, **tile),
            km.Buffer(index=16, **tile),
        ],
        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("add", 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 row:      ", numbers[64, :8].tolist(), "...")
    print("destination, first row: ", total[0, :8].astype(int).tolist(), "...")
    assert (total == expected).all()
    print("every number of the destination is the sum 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 addition and its output.
#pragma once

#include <stdint.h>

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

// The reader kernel: copies two tiles of the source tensor into two buffers.
#include "api/dataflow/dataflow_api.h"
#include "komorebi/device/km_worker.hpp"
#include "handlers.h"

void kernel_main() {
    constexpr auto source_arguments = TensorAccessorArgs<0>();
    km::Reader reader;
    reader.run([&](const km_command_t& command) {
        if (km::handler(command) == kHandlerAdd) {
            const uint32_t source_address = command.a0;
            const uint32_t first_tile = command.a1;
            const uint32_t second_tile = command.a2;
            const auto source = TensorAccessor(source_arguments, source_address);
            if (!reader.reserve(kFirstInput, 1) || !reader.reserve(kSecondInput, 1)) {
                return;
            }
            noc_async_read_page(first_tile, source, get_write_ptr(kFirstInput));
            noc_async_read_page(second_tile, source, get_write_ptr(kSecondInput));
            noc_async_read_barrier();
            cb_push_back(kFirstInput, 1);
            cb_push_back(kSecondInput, 1);
        }
    });
}

// The compute kernel: adds the two input tiles and puts the sum 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/eltwise_binary.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(kFirstInput, kSecondInput, kOutput);
    add_init(kFirstInput, kSecondInput);
    compute.run([&](const km_command_t& command) {
        if (km::handler(command) != kHandlerAdd) {
            return;
        }
        if (
            !compute.wait_front(kFirstInput, 1) || !compute.wait_front(kSecondInput, 1)
        ) {
            return;
        }
        tile_regs_acquire();
        add_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 writer kernel: writes the tile of the output buffer to the destination tensor.
#include "api/dataflow/dataflow_api.h"
#include "komorebi/device/km_worker.hpp"
#include "handlers.h"

void kernel_main() {
    constexpr auto destination_arguments = TensorAccessorArgs<0>();
    km::Writer writer;
    writer.run([&](const km_command_t& command) {
        if (km::handler(command) == kHandlerAdd) {
            const uint32_t destination_address = command.a3;
            const uint32_t destination_tile = command.a4;
            const auto destination = (
                TensorAccessor(destination_arguments, destination_address)
            );
            if (!writer.wait_front(kOutput, 1)) {
                return false;
            }
            noc_async_write_page(destination_tile, destination, get_read_ptr(kOutput));
            noc_async_write_barrier();
            cb_pop_front(kOutput, 1);
        }
        return true;
    });
}

tile 1, first row:       [-6, 4, 0, -3, 6, 3, -3, 1] ...
tile 2, first row:       [3, 4, 5, -6, 1, -2, -4, 6] ...
destination, first row:  [-3, 8, 5, -9, 7, 1, -7, 7] ...
every number of the destination is the sum of tile 1 and tile 2

The first row of the two source tiles, and of the destination: -6 + 3 is -3, 4 + 4 is 8, 0 + 5 is 5. The script checks all 1024 numbers against numpy.

Common mistakes

• Forgetting cb_pop_front on an input. The next command finds the buffer full, as in lesson 16, Reading from DRAM.
• Changing the order of the register calls. Keep the order of the general form: acquire, operation, commit, wait, pack, release.
• Expecting a fraction from "bfloat16" to be exact. It keeps about three decimal digits; whole numbers up to 256 are exact.

Next: lesson 21, The matrix unit replaces the addition with a product of tiles.