Skip to content

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

28. A routed message

A routed message is a small block of numbers that one core sends to another core, where the sender picks the receiver while the program runs, from the data. This lesson relies on lesson 24, Why cores talk, lesson 11, A core group and lesson 12, Partitions.

Why routed messages

A port, a reduce or a broadcast has its senders and receivers fixed by the command before anything runs. Many programs cannot say that. A ray of light goes to the core that holds the part of the scene it hit next; a key to be sorted goes to the core in charge of its range. Which core that is depends on the data.

For this, Komorebi has dynamic handlers: a handler whose cores send messages, each of 64 bytes (16 numbers), to the core they choose with a single call. Nothing about who receives is decided in advance. The run ends by itself when every message has ended.

The life of the program of this lesson.

The general form

In Python:

program = km.Program(
    NAME,
    kernels={"reader": READER, "compute": COMPUTE, "writer": WRITER},
    handlers=HEADER,
).dynamic(HANDLER, message_bytes=64, max_live=LIVE, queue_messages=LIVE)
run.submit(HANDLER, group=(FIRST, COUNT), args=[...], exclusive=True)

In the compute kernel, inside the function you give to run:

km::Messages messages(compute, command);
if (messages.seeding()) {
    messages.send(DESTINATION, message);
    messages.seeded();
}
while (km::Packet packet = messages.next()) {
    for (uint32_t index = 0; index < packet.count; index++) {
        ... packet[index][WORD] ...
        messages.consume();
    }
    messages.done(packet);
}
  • READER and WRITER are Komorebi's own: messages_reader.cpp and messages_writer.cpp, in the runtime's kernel folder. Only the compute kernel is yours.
  • .dynamic(HANDLER, ...) marks the handler as dynamic. message_bytes is the size of a message. max_live is how many messages may exist at once in the whole group; queue_messages is how many each core's inbox in DRAM holds, and may not be smaller.
  • The command starts the run. It stays in flight until the run has ended.
  • The last 2 cores of the group are counter tiles: they keep the inboxes' counters and take no messages. A group of 8 has 6 cores that take messages, ranks 0 to 5.
  • message is an array of 16 numbers. Word 0 is Komorebi's; the others are yours.
  • messages.seeding() is true once, at the start, and is the place for the first sends. messages.seeded() ends it.
  • messages.send(DESTINATION, message) sends a copy of message to the core of rank DESTINATION of the group. It returns false when this core has used up its quota (lesson 30, Full buffers).
  • messages.next() waits for the next packet: up to 32 messages that arrived. It returns an empty packet when the run has ended.
  • packet[index] is the address in L1 of message number index; packet[index][WORD] is one of its numbers.
  • messages.consume() says that a message ends here. messages.done(packet) frees the packet.
  • messages.rank() is this core's rank, and messages.count() the cores that take messages.

One line at a time

program = km.Program(
    "routed_message",
    kernels={
        "reader": RUNTIME / "messages_reader.cpp",
        "compute": KERNELS / "compute.cpp",
        "writer": RUNTIME / "messages_writer.cpp",
    },
    handlers=KERNELS / "handlers.h",
).dynamic("route", message_bytes=64, max_live=4096, queue_messages=4096)

The program takes its reader and writer from the runtime and its compute kernel from this folder, and declares the handler dynamic with room for 4096 messages.

compute.run([&](const km_command_t& command) {
    if (km::handler(command) != kHandlerRoute) {
        return;
    }
    const uint32_t message_count = command.a0;
    const uint32_t notes_address = command.a1;
    km::Messages messages(compute, command);
    uint32_t message[kMessageWords] = {};
    uint32_t received = 0;
    uint32_t sum = 0;
    if (messages.seeding()) {
        if (messages.rank() == 0) {
            for (uint32_t value = 0; value < message_count; value++) {
                message[kValue] = value;
                messages.send(value * value % messages.count(), message);
            }
        }
        messages.seeded();
    }
    while (km::Packet packet = messages.next()) {
        for (uint32_t index = 0; index < packet.count; index++) {
            received++;
            sum += packet[index][kValue];
            messages.consume();
        }
        messages.done(packet);
    }
    UNPACK(({
        volatile tt_l1_ptr uint32_t* notes = km::l1(notes_address);
        notes[0] = received;
        notes[1] = sum;
    }));
});

The first part is seeding: rank 0 sends one message for each value from 0 to 11. The value goes in word 1, and the destination is value * value % messages.count(): the square of the value, modulo the 6 cores. The squares of 0 to 11 modulo 6 can only be 0, 1, 3 and 4, so cores 2 and 5 will receive nothing. The destination comes from the data: change the values and the destinations change, with no change to anything else.

The while loop is the same on every core, including rank 0. It takes packets, counts each message, adds its value, and consumes it. At the end, the compute kernel stores the count and the sum in the core's L1, where Python reads them (lesson 24, Why cores talk).

handle = run.submit(
    "route", group=(0, GROUP), args=[MESSAGES, notes.address], exclusive=True
)
run.wait(handle, timeout=10.0)
run.stop()

One command, which stays in flight until every message has been consumed. wait returns then. The command's arguments are the number of messages and the address of the L1 tensor.

The whole program

"""A routed message: rank 0 sends twelve messages; each goes where its value says."""

from pathlib import Path

import numpy as np

import komorebi as km

HERE = Path(__file__).resolve().parent
KERNELS = HERE / "kernels"
RUNTIME = HERE.parents[1] / "src" / "komorebi_native" / "kernels"
GROUP = 8  # 6 cores take messages; the last 2 cores of a dynamic group keep counters
MESSAGES = 12

with km.Session() as session:
    notes = session.tensor(
        np.zeros((120, 4), np.int32), "int32", shard_grid=session.grid
    )

    program = km.Program(
        "routed_message",
        kernels={
            "reader": RUNTIME / "messages_reader.cpp",
            "compute": KERNELS / "compute.cpp",
            "writer": RUNTIME / "messages_writer.cpp",
        },
        handlers=KERNELS / "handlers.h",
    ).dynamic("route", message_bytes=64, max_live=4096, queue_messages=4096)
    run = session.load(program)

    handle = run.submit(
        "route", group=(0, GROUP), args=[MESSAGES, notes.address], exclusive=True
    )
    run.wait(handle, timeout=10.0)
    run.stop()

    kept = notes.read()[: GROUP - 2, :2]
    for rank, (received, total) in enumerate(kept.tolist()):
        print(f"rank {rank}: received {received} messages, values add up to {total}")
    values = np.arange(MESSAGES)
    destinations = values * values % (GROUP - 2)
    for rank in range(GROUP - 2):
        assert kept[rank].tolist() == [
            int((destinations == rank).sum()),
            int(values[destinations == rank].sum()),
        ]

// What the kernels and Python agree on: the dynamic handler, and the words of a
// message.
#pragma once

#include <stdint.h>

constexpr uint32_t kHandlerRoute = 1;

// A message is 16 words of 4 bytes; word 0 belongs to Komorebi. Word 1 is the value.
constexpr uint32_t kMessageWords = 16;
constexpr uint32_t kValue = 1;

// The compute kernel of a dynamic handler: rank 0 sends 12 messages, each to the core
// its value picks; every core adds up what it receives and keeps it in L1.
#include "api/compute/common.h"
#include "komorebi/device/km_messages.hpp"
#include "komorebi/device/km_worker.hpp"
#include "handlers.h"

void kernel_main() {
    km::Compute compute;
    compute.run([&](const km_command_t& command) {
        if (km::handler(command) != kHandlerRoute) {
            return;
        }
        const uint32_t message_count = command.a0;
        const uint32_t notes_address = command.a1;
        km::Messages messages(compute, command);
        uint32_t message[kMessageWords] = {};
        uint32_t received = 0;
        uint32_t sum = 0;
        if (messages.seeding()) {
            if (messages.rank() == 0) {
                for (uint32_t value = 0; value < message_count; value++) {
                    message[kValue] = value;
                    messages.send(value * value % messages.count(), message);
                }
            }
            messages.seeded();
        }
        while (km::Packet packet = messages.next()) {
            for (uint32_t index = 0; index < packet.count; index++) {
                received++;
                sum += packet[index][kValue];
                messages.consume();
            }
            messages.done(packet);
        }
        UNPACK(({
            volatile tt_l1_ptr uint32_t* notes = km::l1(notes_address);
            notes[0] = received;
            notes[1] = sum;
        }));
    });
}

This folder holds only the compute kernel: reader.cpp and writer.cpp are the runtime's.

rank 0: received 2 messages, values add up to 6
rank 1: received 4 messages, values add up to 24
rank 2: received 0 messages, values add up to 0
rank 3: received 2 messages, values add up to 12
rank 4: received 4 messages, values add up to 24
rank 5: received 0 messages, values add up to 0

Rank 0 receives the values 0 and 6, rank 1 values 1, 5, 7 and 11 (24 in all), rank 3 values 3 and 9, and rank 4 values 2, 4, 8 and 10. The program checks every row against the same calculation done in numpy.

Common mistakes

  • A destination outside the group, such as rank 8 here. Komorebi stops the run with the reason bad dest. The ranks of the counter tiles take no message either.
  • Sending without consuming. Every message has to end with a consume somewhere, or the run never ends and the card stops the command at its deadline (lesson 10, When it goes wrong).
  • Sending more messages than the quota (lesson 30, Full buffers).
  • Forgetting exclusive=True, or sending a second dynamic command while the first is in flight. One dynamic run at a time.

Next: lesson 29, Forward and consume is about a message that goes on to several cores.