Part 3. Data · lesson 22 of 42

22. A whole map

A map applies the same operation to every tile of a tensor. Here the cores of a group each take their share of the tiles (lesson 12, Partitions) and run the whole read, compute and write path of Part 3 on it, at the same time. This lesson relies on lesson 11, A core group, lesson 12, Partitions, lesson 20, Computing on a tile and lesson 19, Circular buffers.

Why a map

Until now one core worked on a handful of tiles. A real tensor has thousands. Each tile is independent of the others, so the tiles can be shared among the cores of a group and all the cores work at once: a group of 12 does the work of one core in about a twelfth of the time. This is the pattern most programs on the card are made of.

The general form

In every one of the three kernels, inside the function you give to run:

const uint32_t rank = name.me() - km::first(command);
const km::Partition share = km::partition(TILES, rank, km::count(command));
for (uint32_t index = 0; index < share.len; index++) {
    STEP for tile share.start + index
}

• TILES is how many tiles the tensor has. It travels as a command argument.
• The three lines are lessons 11 and 12 together. Every core works out the same share for itself, with no message to any other core.
• STEP is reader.cpp's, the compute kernel's or writer.cpp's step: lesson 18, The three kernels and lesson 20, Computing on a tile.
• The three kernels run the same loop count, share.len, so that every tile reader.cpp pushes is one the compute kernel and writer.cpp take.
• A command's group chooses how many cores share the tiles. The kernels do not change.

A group of 4 cores and 60 tiles: each core's share is 15 tiles, and they work on their own part of the tensor at the same time.

One line at a time

reader.run([&](const km_command_t& command) {
    if (km::handler(command) == kHandlerMap) {
        const uint32_t source_address = command.a0;
        const uint32_t tile_count = command.a1;
        const uint32_t rank = reader.me() - km::first(command);
        const km::Partition share = (
            km::partition(tile_count, rank, km::count(command))
        );
        const auto source = TensorAccessor(source_arguments, source_address);
        for (uint32_t index = 0; index < share.len; index++) {
            const uint32_t tile = share.start + index;
            if (
                !reader.reserve(kFirstInput, 1) || !reader.reserve(kSecondInput, 1)
            ) {
                return;
            }
            noc_async_read_page(tile, source, get_write_ptr(kFirstInput));
            noc_async_read_page(
                tile_count + tile, source, get_write_ptr(kSecondInput)
            );
            noc_async_read_barrier();
            cb_push_back(kFirstInput, 1);
            cb_push_back(kSecondInput, 1);
        }
    }
});

reader.cpp works out its share, then for each tile of it reads the tile from the first half of the source and the tile tile_count further on, from the second half, into the two input buffers.

compute.run([&](const km_command_t& command) {
    if (km::handler(command) != kHandlerMap) {
        return;
    }
    const uint32_t rank = compute.me() - km::first(command);
    const km::Partition share = km::partition(command.a1, rank, km::count(command));
    for (uint32_t index = 0; index < share.len; index++) {
        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 has the same loop around the body of lesson 20, Computing on a tile. It does not need the tile numbers, only how many tiles: share.len.

writer.run([&](const km_command_t& command) {
    if (km::handler(command) == kHandlerMap) {
        const uint32_t destination_address = command.a2;
        const uint32_t first_tile = command.a3;
        const uint32_t rank = writer.me() - km::first(command);
        const km::Partition share = (
            km::partition(command.a1, rank, km::count(command))
        );
        const auto destination = (
            TensorAccessor(destination_arguments, destination_address)
        );
        for (uint32_t index = 0; index < share.len; index++) {
            if (!writer.wait_front(kOutput, 1)) {
                return false;
            }
            noc_async_write_page(
                first_tile + share.start + index, destination, get_read_ptr(kOutput)
            );
            noc_async_write_barrier();
            cb_pop_front(kOutput, 1);
        }
    }
    return true;
});

writer.cpp writes the tile first_tile + share.start + index of the destination. The first_tile is an argument that shifts the output, so the program can keep several results in one tensor.

for round_number, count in enumerate((1, 4, 12)):
    arguments = [source.address, TILES, destination.address, round_number * TILES]
    handle = run.submit("map", group=(0, count), args=arguments)
    microseconds = run.wait(handle, timeout=10.0).microseconds
    print(
        f"group of {count:2}: {TILES} tiles added in "
        f"{microseconds:5.1f} microseconds"
    )

The same command with the group of 1, 4 and then 12 cores. Each writes its result to a different third of the destination: a tensor write needs the cores to be idle, and a loaded program never is, so the program keeps the three results apart and checks them at the end.

The whole program

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

"""A whole map: a group of cores adds two tensors, each core its share of the tiles."""

from pathlib import Path

import numpy as np

import komorebi as km

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

with km.Session() as session:
    rng = np.random.default_rng(1)
    first = rng.integers(-8, 9, (TILES * 32, 32))
    second = rng.integers(-8, 9, (TILES * 32, 32))
    source = session.tensor(np.concatenate([first, second]), "bfloat16", "tile")
    destination = session.tensor(np.zeros((3 * TILES * 32, 32)), "bfloat16", "tile")

    tile = dict(format="bfloat16", page_bytes=source.page_bytes, pages=2)
    program = km.Program(
        "a_whole_map",
        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)

    for round_number, count in enumerate((1, 4, 12)):
        arguments = [source.address, TILES, destination.address, round_number * TILES]
        handle = run.submit("map", group=(0, count), args=arguments)
        microseconds = run.wait(handle, timeout=10.0).microseconds
        print(
            f"group of {count:2}: {TILES} tiles added in "
            f"{microseconds:5.1f} microseconds"
        )
    run.stop()

    total = destination.read()
    for round_number in range(3):
        rows = slice(round_number * TILES * 32, (round_number + 1) * TILES * 32)
        assert (total[rows] == first + second).all()
    print("all three results are the sum of the two tensors")

// 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 kHandlerMap = 1;
constexpr uint32_t kFirstInput = 0;
constexpr uint32_t kSecondInput = 1;
constexpr uint32_t kOutput = 16;

// The reader kernel: for each tile of this core's share, copies tile i of the first
// half of the source and tile i of the second half into the two input 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) == kHandlerMap) {
            const uint32_t source_address = command.a0;
            const uint32_t tile_count = command.a1;
            const uint32_t rank = reader.me() - km::first(command);
            const km::Partition share = (
                km::partition(tile_count, rank, km::count(command))
            );
            const auto source = TensorAccessor(source_arguments, source_address);
            for (uint32_t index = 0; index < share.len; index++) {
                const uint32_t tile = share.start + index;
                if (
                    !reader.reserve(kFirstInput, 1) || !reader.reserve(kSecondInput, 1)
                ) {
                    return;
                }
                noc_async_read_page(tile, source, get_write_ptr(kFirstInput));
                noc_async_read_page(
                    tile_count + tile, source, get_write_ptr(kSecondInput)
                );
                noc_async_read_barrier();
                cb_push_back(kFirstInput, 1);
                cb_push_back(kSecondInput, 1);
            }
        }
    });
}

// The compute kernel: adds, for each tile of this core's share, the two input tiles.
#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) != kHandlerMap) {
            return;
        }
        const uint32_t rank = compute.me() - km::first(command);
        const km::Partition share = km::partition(command.a1, rank, km::count(command));
        for (uint32_t index = 0; index < share.len; index++) {
            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 each tile of this core's share 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) == kHandlerMap) {
            const uint32_t destination_address = command.a2;
            const uint32_t first_tile = command.a3;
            const uint32_t rank = writer.me() - km::first(command);
            const km::Partition share = (
                km::partition(command.a1, rank, km::count(command))
            );
            const auto destination = (
                TensorAccessor(destination_arguments, destination_address)
            );
            for (uint32_t index = 0; index < share.len; index++) {
                if (!writer.wait_front(kOutput, 1)) {
                    return false;
                }
                noc_async_write_page(
                    first_tile + share.start + index, destination, get_read_ptr(kOutput)
                );
                noc_async_write_barrier();
                cb_pop_front(kOutput, 1);
            }
        }
        return true;
    });
}

group of  1: 60 tiles added in  38.5 microseconds
group of  4: 60 tiles added in  11.3 microseconds
group of 12: 60 tiles added in   5.6 microseconds
all three results are the sum of the two tensors

Sixty tiles took 38 microseconds on one core, 11 on four and 5.6 on twelve: about 3.4 and 7 times faster. Not 4 and 12 times: every command also pays a fixed cost to reach the cores and to hear that they are done. The last line is the program's check: all three results are the sum of the two source tensors.

Common mistakes

• Different loop counts in the three kernels. If reader.cpp's share and writer.cpp's share differ, one waits for a tile that never comes, and the command ends with CommandTimeout.
• Writing the tensor while the program is loaded. destination.write(...) waits for every core to be idle, and never returns: it fails after a minute. Make tensors before load.
• More cores than tiles. The last cores get share.len == 0 and their loops run zero times; the command still finishes.

A core can now read tiles from DRAM, work on them with the compute kernel and the matrix unit, write them back, and a group can do that for a whole tensor. One more lesson in this part lets both data-movement CPUs of a core carry one stream.

Next: lesson 23, Both movers.