Part 3. Data · lesson 18 of 42

18. The three kernels¶

The data-movement kernel reader.cpp, the compute function and the data-movement kernel writer.cpp are three separate programs that run at the same time on one core, each in a loop of its own, and they meet only in the buffers in L1. This lesson relies on lesson 16, Reading from DRAM and lesson 17, Writing to DRAM.

Why three¶

The core has five CPUs (lesson 3, The card) and each can run only its own code. If one CPU had to read a tile, work on it and write it out, it would do the three in turn, and while it waited for DRAM the arithmetic would stand still. With a CPU for each job, reader.cpp can be fetching tile 3 while the arithmetic works on tile 2 and writer.cpp sends out tile 1.

The price is that the three have to be told when to go: writer.cpp must not write a tile reader.cpp has not finished reading. That is what the buffer's four calls do. The reserve and cb_push_back of reader.cpp and the wait_front and cb_pop_front of writer.cpp are the handshake between them (lesson 16, Reading from DRAM, lesson 17, Writing to DRAM). None of the three knows the others' code; each knows only the buffer.

So far the compute kernel did nothing. lesson 20, Computing on a tile puts it between reader.cpp and writer.cpp. In this lesson only reader.cpp and writer.cpp work, and the new thing is that each has a loop.

The general form¶

In reader.cpp, and in writer.cpp in the same way:

for (uint32_t index = 0; index < COUNT; index++) {
    STEP
}

COUNT is how many tiles this command moves: here an argument of the command.
STEP in reader.cpp is reserve, noc_async_read_page, noc_async_read_barrier, cb_push_back; in writer.cpp wait_front, noc_async_write_page, noc_async_write_barrier, cb_pop_front.
The two loops have the same COUNT. Each tile that reader.cpp pushes is a tile that writer.cpp pops. If the counts differ, one of them waits for ever.
The loops run at the same time. reader.cpp's second iteration can start while writer.cpp is still writing the first tile, as soon as the buffer is free.

Two loops, one buffer. `reader.cpp` pushes a tile when the buffer is free, `writer.cpp` waits for one.

One line at a time¶

reader.run([&](const km_command_t& command) {
    if (km::handler(command) == kHandlerCopy) {
        const uint32_t source_address = command.a0;
        const uint32_t first_tile = command.a1;
        const uint32_t tile_count = command.a4;
        const auto source = TensorAccessor(source_arguments, source_address);
        for (uint32_t index = 0; index < tile_count; index++) {
            if (!reader.reserve(kTileBuffer, 1)) {
                return;
            }
            noc_async_read_page(
                first_tile + index, source, get_write_ptr(kTileBuffer)
            );
            noc_async_read_barrier();
            cb_push_back(kTileBuffer, 1);
        }
    }
});

Compared to lesson 16, Reading from DRAM, the body is inside a for loop, and the tile read is first_tile + index. The first tile and the count are arguments 1 and 4 of the command.

writer.run([&](const km_command_t& command) {
    if (km::handler(command) == kHandlerCopy) {
        const uint32_t destination_address = command.a2;
        const uint32_t first_tile = command.a3;
        const uint32_t tile_count = command.a4;
        const auto destination = (
            TensorAccessor(destination_arguments, destination_address)
        );
        for (uint32_t index = 0; index < tile_count; index++) {
            if (!writer.wait_front(kTileBuffer, 1)) {
                return false;
            }
            noc_async_write_page(
                first_tile + index, destination, get_read_ptr(kTileBuffer)
            );
            noc_async_write_barrier();
            cb_pop_front(kTileBuffer, 1);
        }
    }
    return true;
});

writer.cpp's loop is the same shape: the destination's first tile is argument 3 and the count is argument 4, the same count as reader.cpp's.

one_by_one = sum(copy(tile, 1) for tile in range(TILES))
all_at_once = copy(0, TILES)
print(f"{TILES} commands of 1 tile: {one_by_one:6.1f} microseconds in all")
print(f"1 command of {TILES} tiles:  {all_at_once:6.1f} microseconds")

copy(first_tile, tile_count) is a small function of the program that sends the command and returns the microseconds the command took, from the event's microseconds. The first line sums eight commands of one tile each. The second sends one command of eight tiles.

The whole program¶

three_kernels.pykernels/handlers.hkernels/reader.cppkernels/writer.cpp

"""Three kernels: the reader and the writer each loop over the tiles of one command."""

from pathlib import Path

import numpy as np

import komorebi as km

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

with km.Session() as session:
    numbers = np.arange(TILES * 32 * 32, dtype=np.uint32).reshape(TILES * 32, 32)
    source = session.tensor(numbers, "uint32", "tile")
    destination = session.tensor(np.zeros_like(numbers), "uint32", "tile")

    program = km.Program(
        "three_kernels",
        kernels=HERE / "kernels",
        handlers=HERE / "kernels" / "handlers.h",
        buffers=[
            km.Buffer(index=0, format="uint32", page_bytes=source.page_bytes, pages=1)
        ],
        compile_args={
            "reader": source.accessor_args(),
            "writer": destination.accessor_args(),
        },
    )
    run = session.load(program)

    def copy(first_tile, tile_count):
        """Copies tile_count tiles in one command; returns its microseconds."""
        handle = run.submit(
            "copy",
            group=(0, 1),
            args=[
                source.address,
                first_tile,
                destination.address,
                first_tile,
                tile_count,
            ],
        )
        return run.wait(handle, timeout=10.0).microseconds

    one_by_one = sum(copy(tile, 1) for tile in range(TILES))
    all_at_once = copy(0, TILES)
    print(f"{TILES} commands of 1 tile: {one_by_one:6.1f} microseconds in all")
    print(f"1 command of {TILES} tiles:  {all_at_once:6.1f} microseconds")
    run.stop()

    (source_pages,) = session.read(source)
    (destination_pages,) = session.read(destination)
    assert (destination_pages == source_pages).all()
    print(f"all {TILES} tiles of the destination are the tiles of the source")

// What the kernels and Python agree on: the handler, and the number of the buffer in L1
// in which the tile waits between the reader and the writer.
#pragma once

#include <stdint.h>

constexpr uint32_t kHandlerCopy = 1;
constexpr uint32_t kTileBuffer = 0;

// The reader kernel: copies tile_count tiles from DRAM, one after the other, through
// the buffer in L1.
#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) == kHandlerCopy) {
            const uint32_t source_address = command.a0;
            const uint32_t first_tile = command.a1;
            const uint32_t tile_count = command.a4;
            const auto source = TensorAccessor(source_arguments, source_address);
            for (uint32_t index = 0; index < tile_count; index++) {
                if (!reader.reserve(kTileBuffer, 1)) {
                    return;
                }
                noc_async_read_page(
                    first_tile + index, source, get_write_ptr(kTileBuffer)
                );
                noc_async_read_barrier();
                cb_push_back(kTileBuffer, 1);
            }
        }
    });
}

// The writer kernel: waits for the tile in L1 and writes tile_count tiles to another
// tensor, one after the other.
#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) == kHandlerCopy) {
            const uint32_t destination_address = command.a2;
            const uint32_t first_tile = command.a3;
            const uint32_t tile_count = command.a4;
            const auto destination = (
                TensorAccessor(destination_arguments, destination_address)
            );
            for (uint32_t index = 0; index < tile_count; index++) {
                if (!writer.wait_front(kTileBuffer, 1)) {
                    return false;
                }
                noc_async_write_page(
                    first_tile + index, destination, get_read_ptr(kTileBuffer)
                );
                noc_async_write_barrier();
                cb_pop_front(kTileBuffer, 1);
            }
        }
        return true;
    });
}

The compute file is the empty one.

8 commands of 1 tile:   19.1 microseconds in all
1 command of 8 tiles:     8.2 microseconds
all 8 tiles of the destination are the tiles of the source

All eight tiles arrive, in both ways, which the script checks. The times differ a little on every run. The one command of eight tiles took less than half the time of the eight commands of one tile: the tiles took the same work either way, and each command costs a trip from the dispatcher to the core and back, which the eight commands paid eight times.

Common mistakes¶

Different counts in reader.cpp and writer.cpp. If reader.cpp reads 8 and writer.cpp writes 4, the command ends with reader.cpp waiting for room: CommandTimeout (lesson 10, When it goes wrong).
Expecting reader.cpp to finish all its tiles before writer.cpp starts. With a buffer of one tile they take turns, tile by tile; the next lesson is about how to let them overlap.

Next: lesson 19, Circular buffers gives the buffer more than one tile of room.