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

25. A port

A port is a counter on the receiving core: the sending core writes a tile into the receiver's L1 and then adds 1 to the port, and the receiver waits until the port says the tile is there. This lesson relies on lesson 24, Why cores talk, lesson 11, A core group and lesson 19, Circular buffers.

Why a port

The sender can write straight into another core's L1 over the NoC, but the receiver has to know when the write has finished, or it would use half a tile. The sender therefore does two things in order: it writes the tile, and, when the write is done, it adds 1 to a counter in the receiver's L1. The receiver waits for the counter, and only then uses the tile. The counter is the port; every core has eight of them, numbered 0 to 7.

All three kernels of the two cores take part: on the sender reader.cpp fetches the tile and writer.cpp sends it; on the receiver reader.cpp waits at the port and writer.cpp uses the tile.

A tile goes from DRAM to the outgoing buffer of core 1, over the NoC to the arrival buffer of core 0, and a number of it goes to Python.

The general form

In the kernels, with km::Arrival arrival = {BUFFER, PAGES, 1}; made before run, and #include "komorebi/device/km_port.hpp":

Sender's writer.cpp:

km::emit(writer, command, arrival, OUTGOING, 1, RECEIVER, PORT, SLOT);

Receiver's reader.cpp:

km::arrive(reader, command, arrival, PORT, 1);

Receiver, to use the tile, after reader.cpp:

writer.wait_front(BUFFER, PAGES, km::kBehindPort);
... get_read_ptr(BUFFER) ...
cb_pop_front(BUFFER, PAGES);

In Python:

run.submit(KIND, group=(FIRST, COUNT), args=[...], exclusive=True)

• BUFFER is the number of the arrival buffer, a buffer of PAGES tiles that every core has. It is declared in km.Buffer like any other. A tile sent to a core lands in one of its pages.
• OUTGOING is the buffer in which the sender has the tile to send.
• RECEIVER is the number of the receiving core, as the core numbers of lesson 3.
• PORT is which of the 8 counters the sender increments and the receiver waits for: 0 to 7.
• SLOT is which page of the arrival buffer the tile lands in: 0 to PAGES - 1. Two tiles sent to the same core by different senders use different slots.
• emit copies the tile, waits until the write is done, increments the port, waits until that is done, and pops the outgoing tile.
• arrive waits until the port counts 1, clears it, and tells the other kernels that the arrival buffer is full.
• exclusive=True makes the dispatcher hold back a command whose group overlaps one still running, so that a tile of one command can never land in the slots of the next.

One line at a time

reader.run([&](const km_command_t& command) {
    if (km::handler(command) == kHandlerSend) {
        const uint32_t rank = reader.me() - km::first(command);
        if (rank == 1) {
            const auto source = TensorAccessor(source_arguments, command.a0);
            if (!reader.reserve(kOutgoing, 1)) {
                return;
            }
            noc_async_read_page(command.a1, source, get_write_ptr(kOutgoing));
            noc_async_read_barrier();
            cb_push_back(kOutgoing, 1);
        } else {
            km::arrive(reader, command, arrival, 0, 1);
        }
    }
});

Both cores run this reader.cpp. The one with rank 1 is the sender: it reads the tile into the outgoing buffer, as in lesson 16, Reading from DRAM. The one with rank 0 waits at port 0 for 1 arrival.

writer.run([&](const km_command_t& command) {
    if (km::handler(command) == kHandlerSend) {
        const uint32_t rank = writer.me() - km::first(command);
        if (rank == 1) {
            const uint32_t receiver = km::first(command);
            return km::emit(writer, command, arrival, kOutgoing, 1, receiver, 0, 0);
        }
        if (!writer.wait_front(kArrival, 1, km::kBehindPort)) {
            return false;
        }
        const auto tile = km::l1(get_read_ptr(kArrival));
        writer.result(tile[command.a2]);
        cb_pop_front(kArrival, 1);
    }
    return true;
});

The sender's writer.cpp calls emit: the tile goes to receiver, the first core of the group, port 0, slot 0. The receiver's writer.cpp waits until the arrival buffer is full, reads number a2 of the tile, sets it as the result (lesson 13, The leader), and frees the buffer. The receiver is rank 0 on purpose: it is the leader, so its result reaches Python.

for tile, position in [(1, 0), (3, 5)]:
    arguments = [source.address, tile, position]
    handle = run.submit(
        "send", group=(0, 2), args=arguments, exclusive=True, result=True
    )
    number = run.wait(handle, timeout=10.0).result
    print(
        f"core 1 sent tile {tile}; core 0 holds number {position} of it: {number}"
    )
    assert number == 1000 + tile * 1024 + position

Two commands to the group of cores 0 and 1. The first arguments are the source tensor, which tile core 1 reads, and which number of it core 0 reports.

The whole program

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

"""A port: core 1 sends a tile from DRAM to core 0, which reports a number of it."""

from pathlib import Path

import numpy as np

import komorebi as km

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

with km.Session() as session:
    numbers = np.arange(4 * 32 * 32, dtype=np.uint32).reshape(4 * 32, 32) + 1000
    source = session.tensor(numbers, "uint32", "tile")

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

    for tile, position in [(1, 0), (3, 5)]:
        arguments = [source.address, tile, position]
        handle = run.submit(
            "send", group=(0, 2), args=arguments, exclusive=True, result=True
        )
        number = run.wait(handle, timeout=10.0).result
        print(
            f"core 1 sent tile {tile}; core 0 holds number {position} of it: {number}"
        )
        assert number == 1000 + tile * 1024 + position
    run.stop()

// What the kernels and Python agree on: the handler, and two buffers in L1: the
// outgoing tile of the sender, and the arrival buffer of the receiver.
#pragma once

#include <stdint.h>

constexpr uint32_t kHandlerSend = 1;
constexpr uint32_t kOutgoing = 0;
constexpr uint32_t kArrival = 1;

// The reader kernel: the sending core (rank 1) reads a tile from DRAM; the receiving
// core (rank 0) waits for the tile to arrive.
#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, 1, 1};
    reader.run([&](const km_command_t& command) {
        if (km::handler(command) == kHandlerSend) {
            const uint32_t rank = reader.me() - km::first(command);
            if (rank == 1) {
                const auto source = TensorAccessor(source_arguments, command.a0);
                if (!reader.reserve(kOutgoing, 1)) {
                    return;
                }
                noc_async_read_page(command.a1, source, get_write_ptr(kOutgoing));
                noc_async_read_barrier();
                cb_push_back(kOutgoing, 1);
            } else {
                km::arrive(reader, command, arrival, 0, 1);
            }
        }
    });
}

// The writer kernel: the sending core (rank 1) sends the tile to the receiving core
// (rank 0), which returns a number of it as the result.
#include "api/dataflow/dataflow_api.h"
#include "komorebi/device/km_port.hpp"
#include "handlers.h"

void kernel_main() {
    km::Writer writer;
    const km::Arrival arrival = {kArrival, 1, 1};
    writer.run([&](const km_command_t& command) {
        if (km::handler(command) == kHandlerSend) {
            const uint32_t rank = writer.me() - km::first(command);
            if (rank == 1) {
                const uint32_t receiver = km::first(command);
                return km::emit(writer, command, arrival, kOutgoing, 1, receiver, 0, 0);
            }
            if (!writer.wait_front(kArrival, 1, km::kBehindPort)) {
                return false;
            }
            const auto tile = km::l1(get_read_ptr(kArrival));
            writer.result(tile[command.a2]);
            cb_pop_front(kArrival, 1);
        }
        return true;
    });
}

The compute file is the empty one.

core 1 sent tile 1; core 0 holds number 0 of it: 2024
core 1 sent tile 3; core 0 holds number 5 of it: 4077

The numbers are 1000 plus their place in the whole tensor, so tile 1 starts at 2024 and tile 3 number 5 is 4077. Core 0 never read the source tensor; the tile came from core 1.

Common mistakes

• Leaving out exclusive=True. The documentation of Komorebi requires it for every command that uses ports.
• A sender that sends to one port while the receiver waits at another. The receiver never sees its arrival, and the card stops the command at its deadline (lesson 10, When it goes wrong).
• An arrival buffer with fewer slots than the ones used. The call gives up and the command ends the same way.
• A receiver that is not the leader and sets the result: Python receives the leader's (lesson 13, The leader).

Next: lesson 26, Reduce sends from many cores to one.