Part 4. Cores talking to cores · lesson 26 of 42

26. Reduce

A reduce combines one tile from each core of a group into a single tile on one core, here by adding them, and Komorebi has two ways to route the tiles. This lesson relies on lesson 25, A port, lesson 11, A core group and lesson 20, Computing on a tile.

Why a reduce

When a job is shared among cores (lesson 22, A whole map), each core ends up with a part of the answer, and often the answer is the sum of the parts: a total, a dot product, the last step of a matrix product. The parts have to meet in one place and be added. That is a reduce: every core sends its part, and one core, rank 0, ends with the sum.

Two ways to route it:

Form	Who sends to whom	Steps	Slots in the arrival buffer
owner	every core but rank 0 sends straight to rank 0	1, rank 0 adds size - 1 tiles	size - 1
tree	in round 0 odd ranks send to the even rank before them; in round 1 ranks 2, 6, ... send; and so on	log2(size) rounds, each core adds at most one tile per round	one per round, at most 7

The owner is the simplest and does all the adding on one core. The tree spreads the adding: no core adds more than once a round, and the whole takes 3 rounds for 8 cores and 7 for 128. Both use the ports of lesson 25, A port; the sums are the same.

The tree form for 8 cores: three rounds. Each round halves the number of cores that hold a part.

The general form

In reader.cpp, writer.cpp and the compute kernel, the group and the form:

const km::Group group = km::group_of(command, NAME.me());
const km::Arrival arrival = {BUFFER, PAGES, 1};

Reader, after reading the core's own tile:

km::reduce_arrive(reader, command, arrival, group, FORM);

Compute:

ok = km::reduce(compute, arrival, group, FORM, [&](uint32_t slot) {
    ...  // add the tile in slot `slot` of the arrival buffer
});

Writer:

if (group.rank != 0) {
    return km::reduce_emit(writer, command, arrival, group, FORM, OUTPUT, 1);
}
km::reduce_result(writer, OUTPUT, 1);

• km::group_of(command, core) is the command's group as an object: group.rank and group.size (lesson 11, A core group).
• FORM is km::kReduceOwner (0) or km::kReduceTree (1).
• reduce_arrive makes reader.cpp wait at every port a core has to wait at, as in lesson 25, A port, according to the form.
• reduce_emit makes every core but rank 0 send its sum to the core the form says, and reduce_result makes rank 0 wait for the total.
• The function given to km::reduce is the one thing that is yours: how to combine. It is called for each tile that arrives, in rank order, whatever order they came in, so the sum is the same on every run.

One line at a time

reader.run([&](const km_command_t& command) {
    if (km::handler(command) == kHandlerReduce) {
        const km::Group group = km::group_of(command, reader.me());
        const uint32_t form = command.a4;
        const auto source = TensorAccessor(source_arguments, command.a0);
        if (!reader.reserve(kOwnTile, 1)) {
            return;
        }
        noc_async_read_page(
            command.a1 + group.rank, source, get_write_ptr(kOwnTile)
        );
        noc_async_read_barrier();
        cb_push_back(kOwnTile, 1);
        km::reduce_arrive(reader, command, arrival, group, form);
    }
});

Each core reads the tile of its rank (lesson 12, Partitions, one tile each), then waits for what the others send.

compute.run([&](const km_command_t& command) {
    if (km::handler(command) != kHandlerReduce) {
        return;
    }
    const km::Group group = km::group_of(command, compute.me());
    const uint32_t form = command.a4;
    tile_regs_acquire();
    bool ok = compute.wait_front(kOwnTile, 1);
    if (ok) {
        copy_init(kOwnTile);
        copy_tile(kOwnTile, 0, 0);
        cb_pop_front(kOwnTile, 1);
    }
    ok = ok && km::reduce(compute, arrival, group, form, [&](uint32_t slot) {
        copy_init(kArrival);
        copy_tile(kArrival, km::slot_page(arrival, slot), 1);
        add_binary_tile_init();
        add_binary_tile(0, 1, 0);
    });
    tile_regs_commit();
    tile_regs_wait();
    if (ok && compute.reserve_back(kOutput, 1)) {
        pack_tile(0, kOutput);
        cb_push_back(kOutput, 1);
    }
    tile_regs_release();
});

The compute kernel copies its own tile into register 0. The function given to reduce copies the arrived tile into register 1 and adds it to register 0. When the reduce is done, register 0 holds this core's sum (the total, on rank 0) and pack_tile puts it in the output buffer.

writer.run([&](const km_command_t& command) {
    if (km::handler(command) == kHandlerReduce) {
        const km::Group group = km::group_of(command, writer.me());
        const uint32_t form = command.a4;
        if (group.rank != 0) {
            return (
                km::reduce_emit(writer, command, arrival, group, form, kOutput, 1)
            );
        }
        if (!km::reduce_result(writer, kOutput, 1)) {
            return false;
        }
        const auto destination = TensorAccessor(destination_arguments, command.a2);
        noc_async_write_page(command.a3, destination, get_read_ptr(kOutput));
        noc_async_write_barrier();
        cb_pop_front(kOutput, 1);
    }
    return true;
});

writer.cpp of every core but rank 0 sends the output tile on; rank 0 writes it to the destination tensor.

for form in (0, 1):
    arguments = [source.address, 0, destination.address, form, form]
    handle = run.submit("reduce", group=(0, CORES), args=arguments, exclusive=True)
    microseconds = run.wait(handle, timeout=10.0).microseconds
    print(
        f"form {form}: {CORES} tiles added into one in "
        f"{microseconds:5.1f} microseconds"
    )

The same command twice, with the group of 8 cores, once with each form. The two forms write to different tiles of the destination (the argument after its address), so both can be checked at the end.

The whole program

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

"""Reduce: eight cores each hold a tile; the total arrives at core 0, in two ways."""

from pathlib import Path

import numpy as np

import komorebi as km

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

with km.Session() as session:
    numbers = np.random.default_rng(1).integers(-8, 9, (CORES * 32, 32))
    source = session.tensor(numbers, "bfloat16", "tile")
    destination = session.tensor(np.zeros((2 * 32, 32)), "bfloat16", "tile")

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

    for form in (0, 1):
        arguments = [source.address, 0, destination.address, form, form]
        handle = run.submit("reduce", group=(0, CORES), args=arguments, exclusive=True)
        microseconds = run.wait(handle, timeout=10.0).microseconds
        print(
            f"form {form}: {CORES} tiles added into one in "
            f"{microseconds:5.1f} microseconds"
        )
    run.stop()

    total = destination.read()
    expected = numbers.reshape(CORES, 32, 32).sum(axis=0)
    assert (total[:32] == expected).all() and (total[32:] == expected).all()
    print("both forms: the destination tile is the sum of the 8 tiles")

// What the kernels and Python agree on: the handler, and the three buffers in L1: the
// core's own tile, the arrival buffer for the tiles of the other cores, and the output.
#pragma once

#include <stdint.h>

constexpr uint32_t kHandlerReduce = 1;
constexpr uint32_t kOwnTile = 0;
constexpr uint32_t kArrival = 1;
constexpr uint32_t kArrivalPages = 8;
constexpr uint32_t kOutput = 16;

// The reader kernel: reads the core's own tile (tile number rank), then waits for the
// tiles that the other cores send to it.
#include "api/dataflow/dataflow_api.h"
#include "komorebi/device/km_port.hpp"
#include "handlers.h"

void kernel_main() {
    constexpr auto source_arguments = TensorAccessorArgs<0>();
    km::Reader reader;
    const km::Arrival arrival = {kArrival, kArrivalPages, 1};
    reader.run([&](const km_command_t& command) {
        if (km::handler(command) == kHandlerReduce) {
            const km::Group group = km::group_of(command, reader.me());
            const uint32_t form = command.a4;
            const auto source = TensorAccessor(source_arguments, command.a0);
            if (!reader.reserve(kOwnTile, 1)) {
                return;
            }
            noc_async_read_page(
                command.a1 + group.rank, source, get_write_ptr(kOwnTile)
            );
            noc_async_read_barrier();
            cb_push_back(kOwnTile, 1);
            km::reduce_arrive(reader, command, arrival, group, form);
        }
    });
}

// The compute kernel: starts from the core's own tile and adds each tile that arrives.
#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_sfpu.h"
#include "api/compute/pack.h"
#include "api/compute/reg_api.h"
#include "api/compute/tile_move_copy.h"
#include "komorebi/device/km_port.hpp"
#include "handlers.h"

void kernel_main() {
    km::Compute compute;
    compute_kernel_hw_startup(kOwnTile, kOutput);
    const km::Arrival arrival = {kArrival, kArrivalPages, 1};
    compute.run([&](const km_command_t& command) {
        if (km::handler(command) != kHandlerReduce) {
            return;
        }
        const km::Group group = km::group_of(command, compute.me());
        const uint32_t form = command.a4;
        tile_regs_acquire();
        bool ok = compute.wait_front(kOwnTile, 1);
        if (ok) {
            copy_init(kOwnTile);
            copy_tile(kOwnTile, 0, 0);
            cb_pop_front(kOwnTile, 1);
        }
        ok = ok && km::reduce(compute, arrival, group, form, [&](uint32_t slot) {
            copy_init(kArrival);
            copy_tile(kArrival, km::slot_page(arrival, slot), 1);
            add_binary_tile_init();
            add_binary_tile(0, 1, 0);
        });
        tile_regs_commit();
        tile_regs_wait();
        if (ok && compute.reserve_back(kOutput, 1)) {
            pack_tile(0, kOutput);
            cb_push_back(kOutput, 1);
        }
        tile_regs_release();
    });
}

// The writer kernel: a core other than rank 0 sends its sum to the core it reports to;
// rank 0 writes the total to DRAM.
#include "api/dataflow/dataflow_api.h"
#include "komorebi/device/km_port.hpp"
#include "handlers.h"

void kernel_main() {
    constexpr auto destination_arguments = TensorAccessorArgs<0>();
    km::Writer writer;
    const km::Arrival arrival = {kArrival, kArrivalPages, 1};
    writer.run([&](const km_command_t& command) {
        if (km::handler(command) == kHandlerReduce) {
            const km::Group group = km::group_of(command, writer.me());
            const uint32_t form = command.a4;
            if (group.rank != 0) {
                return (
                    km::reduce_emit(writer, command, arrival, group, form, kOutput, 1)
                );
            }
            if (!km::reduce_result(writer, kOutput, 1)) {
                return false;
            }
            const auto destination = TensorAccessor(destination_arguments, command.a2);
            noc_async_write_page(command.a3, destination, get_read_ptr(kOutput));
            noc_async_write_barrier();
            cb_pop_front(kOutput, 1);
        }
        return true;
    });
}

form 0: 8 tiles added into one in   5.1 microseconds
form 1: 8 tiles added into one in   4.6 microseconds
both forms: the destination tile is the sum of the 8 tiles

Both commands put the sum of the 8 tiles in the destination, exactly: the script compares every number with numpy. With 8 cores there is little to choose between the forms; the tree pays off with many cores, where one core would have to add a hundred tiles.

Common mistakes

• An arrival buffer with too few pages. The owner form needs size - 1 slots, so a group of 8 needs at least 7 pages; with fewer, reader.cpp gives up.
• Forgetting exclusive=True (lesson 25, A port).
• A combine function that depends on the order of the tiles, such as a subtraction. The order is fixed (rank order), but the owner and the tree combine in different groupings, so only an operation that does not care, like addition of whole numbers, gives the same result from both.

Next: lesson 27, Broadcast sends one tile from one core to many.