Part 4. Cores talking to cores · lesson 29 of 42
29. Forward and consume¶
A core that receives a routed message can send it on to another core, which is a forward, or end it, which is a consume; a message lives until some core consumes it. This lesson relies on lesson 28, A routed message.
Why forward¶
In lesson 28, A routed message a message went from the core that made it to one other core and ended there. Often the work on a message is not done in one place: a ray bounces from one part of a scene to the next, a record is passed to the next stage. The core that holds the next part is known only after this part has been worked on. So the receiving core looks at the message, changes it if it likes, and forwards it to the core the new data names. The message keeps going until some core decides it is done, and consumes it.
The general form¶
Inside the while loop of lesson 28, A routed message, for each message of the packet:
message[WORD] = packet[index][WORD];
...
if (DONE) {
messages.consume();
} else {
messages.forward(DESTINATION, message);
}
messageis a local array of numbers. The message that arrived is read frompacket[index]; the one that goes on is the local array, so the core can change it first.- Every message that arrives ends in exactly one of two ways:
consume()orforward(...). A core that does neither leaves the message in limbo and the run does not end. messages.forward(DESTINATION, message)sendsmessageon to the core of rankDESTINATION. A forward to the core itself works the same way.DONEis the program's rule for when a message ends. Here it is a count of hops left, carried in the message.- A forward does not change the core's quota, and a
consumegives one unit of quota back to the core where the message ends (lesson 30, Full buffers). - The run ends when every message that was sent has been consumed. Nobody counts them in advance: each core reports how many it has sent and consumed, and the dispatcher ends the run when, twice in a row, the two totals are equal and every core is idle.
One line at a time¶
compute.run([&](const km_command_t& command) {
if (km::handler(command) != kHandlerWander) {
return;
}
const uint32_t message_count = command.a0;
const uint32_t hops = command.a1;
const uint32_t notes_address = command.a2;
km::Messages messages(compute, command);
uint32_t message[kMessageWords] = {};
uint32_t visits = 0;
uint32_t ended_here = 0;
if (messages.seeding()) {
if (messages.rank() == 0) {
for (uint32_t id = 0; id < message_count; id++) {
message[kHopsLeft] = hops;
message[kValue] = id + 1;
messages.send(id % messages.count(), message);
}
}
messages.seeded();
}
while (km::Packet packet = messages.next()) {
for (uint32_t index = 0; index < packet.count; index++) {
message[kHopsLeft] = packet[index][kHopsLeft];
message[kValue] = (packet[index][kValue] * 5 + 1) % 97;
visits++;
if (message[kHopsLeft] == 0) {
ended_here++;
messages.consume();
} else {
message[kHopsLeft] = message[kHopsLeft] - 1;
messages.forward(message[kValue] % messages.count(), message);
}
}
messages.done(packet);
}
UNPACK(({
volatile tt_l1_ptr uint32_t* notes = km::l1(notes_address);
notes[0] = visits;
notes[1] = ended_here;
}));
});
Rank 0 seeds 6 messages, message id to rank id % 6, each with 4 hops and the value
id + 1. At each visit the core makes a new value, (value * 5 + 1) % 97, and counts the
visit. If the message has no hops left, it ends here: the core counts it and consumes it.
Otherwise the hops left drop by one, and the message is forwarded to the rank that the new
value names, value % 6. At the end, each core stores its visit count and its count of
messages that ended there in L1.
arguments = [MESSAGES, HOPS, notes.address]
handle = run.submit("wander", group=(0, GROUP), args=arguments, exclusive=True)
run.wait(handle, timeout=10.0)
run.stop()
The command carries the number of messages, the hops and the address of the tensor.
visits = [0] * RANKS
ended_here = [0] * RANKS
for identifier in range(MESSAGES):
rank, value, hops_left = identifier % RANKS, identifier + 1, HOPS
while True:
visits[rank] += 1
value = (value * 5 + 1) % 97
if hops_left == 0:
ended_here[rank] += 1
break
hops_left -= 1
rank = value % RANKS
The check: the same rules in a few lines of Python, followed message by message. The cores and the replay must agree on every count.
The whole program¶
"""Forward and consume: messages hop from core to core, each hop chosen by its value."""
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
RANKS = GROUP - 2
MESSAGES = 6
HOPS = 4
with km.Session() as session:
notes = session.tensor(
np.zeros((120, 4), np.int32), "int32", shard_grid=session.grid
)
program = km.Program(
"forward_and_consume",
kernels={
"reader": RUNTIME / "messages_reader.cpp",
"compute": KERNELS / "compute.cpp",
"writer": RUNTIME / "messages_writer.cpp",
},
handlers=KERNELS / "handlers.h",
).dynamic("wander", message_bytes=64, max_live=4096, queue_messages=4096)
run = session.load(program)
arguments = [MESSAGES, HOPS, notes.address]
handle = run.submit("wander", group=(0, GROUP), args=arguments, exclusive=True)
run.wait(handle, timeout=10.0)
run.stop()
kept = notes.read()[:RANKS, :2]
for rank, (visits, ended_here) in enumerate(kept.tolist()):
print(f"rank {rank}: {visits} visits, {ended_here} messages ended here")
visits = [0] * RANKS
ended_here = [0] * RANKS
for identifier in range(MESSAGES):
rank, value, hops_left = identifier % RANKS, identifier + 1, HOPS
while True:
visits[rank] += 1
value = (value * 5 + 1) % 97
if hops_left == 0:
ended_here[rank] += 1
break
hops_left -= 1
rank = value % RANKS
assert kept.tolist() == [[v, e] for v, e in zip(visits, ended_here)]
print(
f"{sum(visits)} visits in all, {sum(ended_here)} messages ended, as the replay says"
)
// What the kernels and Python agree on: the dynamic handler, and the words of a
// message.
#pragma once
#include <stdint.h>
constexpr uint32_t kHandlerWander = 1;
// A message is 16 words of 4 bytes; word 0 belongs to Komorebi.
constexpr uint32_t kMessageWords = 16;
constexpr uint32_t kHopsLeft = 1;
constexpr uint32_t kValue = 2;
// The compute kernel of a dynamic handler: a message visits a core, changes its value,
// and is forwarded to the core that the new value names, until its hops are used up.
#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) != kHandlerWander) {
return;
}
const uint32_t message_count = command.a0;
const uint32_t hops = command.a1;
const uint32_t notes_address = command.a2;
km::Messages messages(compute, command);
uint32_t message[kMessageWords] = {};
uint32_t visits = 0;
uint32_t ended_here = 0;
if (messages.seeding()) {
if (messages.rank() == 0) {
for (uint32_t id = 0; id < message_count; id++) {
message[kHopsLeft] = hops;
message[kValue] = id + 1;
messages.send(id % messages.count(), message);
}
}
messages.seeded();
}
while (km::Packet packet = messages.next()) {
for (uint32_t index = 0; index < packet.count; index++) {
message[kHopsLeft] = packet[index][kHopsLeft];
message[kValue] = (packet[index][kValue] * 5 + 1) % 97;
visits++;
if (message[kHopsLeft] == 0) {
ended_here++;
messages.consume();
} else {
message[kHopsLeft] = message[kHopsLeft] - 1;
messages.forward(message[kValue] % messages.count(), message);
}
}
messages.done(packet);
}
UNPACK(({
volatile tt_l1_ptr uint32_t* notes = km::l1(notes_address);
notes[0] = visits;
notes[1] = ended_here;
}));
});
}
rank 0: 4 visits, 1 messages ended here
rank 1: 5 visits, 2 messages ended here
rank 2: 6 visits, 2 messages ended here
rank 3: 5 visits, 0 messages ended here
rank 4: 4 visits, 0 messages ended here
rank 5: 6 visits, 1 messages ended here
30 visits in all, 6 messages ended, as the replay says
6 messages with 4 hops are 5 visits each: 30 visits, spread over the cores by the data. Core 3 and core 4 are visited but no message ends there; at core 2 two messages end. The last line is the check against the replay.
Common mistakes¶
- Neither forwarding nor consuming a message. The message is lost from the run's point of view, the totals never match, and the card stops the command at its deadline (lesson 10, When it goes wrong).
- A rule that never ends a message, such as a hop count that is not decreased. The run does not end, and the card stops the command at its deadline.
- Forwarding the address
packet[index]itself instead of your own array. Copy the words you want intomessagefirst, as the program does.
Next: lesson 30, Full buffers is about a core with too many messages to send.