Part 4. Cores talking to cores · lesson 27 of 42
27. Broadcast¶
A broadcast sends the same tile from one core to every other core of a row (or a column) with a single write over the NoC. This lesson relies on lesson 25, A port, lesson 26, Reduce and lesson 19, Circular buffers.
Why a broadcast¶
Sometimes every core needs the same data: the same weights, the same row of a matrix. If each core read it from DRAM, it would be read from DRAM as many times as there are cores. The NoC can do better: the sender reads a tile once and writes it to a whole rectangle of cores in one go, a multicast.
A port (lesson 25, A port) goes to one receiver. A broadcast goes to many, and it handles a whole series of tiles, called blocks, so that the sender can read the next block while the receivers work on the last. The receivers tell the sender when they have finished with a block; the sender writes into a slot of the receivers' buffer again only when all have. The number of blocks that may be in flight at once is the window.
The general form¶
Reader of every core of the broadcast, in the function given to run:
km::Broadcast NAME(reader, command, FIRST_WORD);
if (!NAME.ok() || !NAME.included()) {
return;
}
for (uint32_t block = 0; block < COUNT; block++) {
if (!NAME.sender()) {
if (!NAME.arrive(block)) { return; }
continue;
}
const uint32_t address = NAME.reserve(block);
... read the block into address ...
if (!NAME.cast(block)) { return; }
}
NAME.finish();
Compute kernel of every core:
km::broadcast_receive(compute, BUFFER, PAGES);
... use the block, in BUFFER ...
km::broadcast_done(BUFFER, PAGES);
...
km::broadcast_end(compute, BUFFER, PAGES, WINDOW, COUNT);
FIRST_WORDis where, among the program's arguments for this kernel, the eleven numbers of the broadcast start; the host makes them.ok()says that the numbers match the command, andincluded()that this core is in the broadcast. Both are checks; a core that fails either has nothing to do.sender()is true on rank 0, the core where the blocks start. Every other corearrives: waits until blockblockhas been written to its buffer.reserve(block)gives the sender the address of the slot for this block, after waiting until every receiver has finished with the block that was there before.cast(block)waits until the reads have landed, writes the block to every receiver with a multicast, and marks it valid.finish()waits until the compute kernel has taken every block, and ends the broadcast.broadcast_receivewaits in the compute kernel until the block is in the buffer.broadcast_donesays it has finished with it, which is what lets the sender go on, andbroadcast_endruns after the last block.- A broadcast runs in one row inside one panel of the card (lesson 3, The card), or in one column.
One line at a time¶
The host first works out the eleven numbers for each core:
nocs = session.nocs()
group = layout.Group(0, CORES)
words = layout.broadcast_args(
nocs.grid, nocs.panels, group, (0, 0, CORES, 1), "row", 0,
nocs.noc0, nocs.noc_xy, (0, 0, WINDOW, 1),
) # fmt: skip
fields = km.abi.struct_fields(km.abi.headers().device[-1], "BroadcastArgs")
def arguments_of(core):
"""The program's arguments of a core: the reader's, the writer's, compute's."""
if core not in words:
return [0] * len(fields), [], [WINDOW]
return [words[core][field] for field in fields], [], [WINDOW]
session.nocs() is the card's map of the NoC. layout.broadcast_args(...) computes, for
each core of the group, the numbers km::Broadcast needs: its rank, where the sender
is and the corners of the rectangle of receivers on the NoC. The arguments are
(0, 0, 6, 1), the rectangle (first column, first row, columns, rows), "row", the NoC the
reader uses (0), and (0, 0, 2, 1): broadcast number 0, buffer 0, a window of 2 blocks,
1 page a block. arguments_of(core) returns the three argument lists of a core, in the
order reader, writer, compute. Those of reader.cpp are these eleven numbers.
reader.run([&](const km_command_t& command) {
if (km::handler(command) != kHandlerBroadcast) {
return;
}
km::Broadcast row(reader, command, 0);
if (!row.ok() || !row.included()) {
return;
}
const auto source = TensorAccessor(source_arguments, command.a0);
const uint32_t block_count = command.a1;
for (uint32_t block = 0; block < block_count; block++) {
if (!row.sender()) {
if (!row.arrive(block)) {
return;
}
continue;
}
const uint32_t address = row.reserve(block);
if (address == 0) {
return;
}
noc_async_read_page(block, source, address);
if (!row.cast(block)) {
return;
}
}
row.finish();
});
reader.cpp is the general form: the sender reads each block of the source into the slot
that reserve gives it and casts it; the others arrive.
compute.run([&](const km_command_t& command) {
if (km::handler(command) != kHandlerBroadcast) {
return;
}
const uint32_t block_count = command.a1;
tile_regs_acquire();
fill_tile_init();
fill_tile(0, 0.0f);
bool ok = true;
for (uint32_t block = 0; ok && block < block_count; block++) {
ok = km::broadcast_receive(compute, kRowBuffer, 1);
if (ok) {
copy_init(kRowBuffer);
copy_tile(kRowBuffer, 0, 1);
add_binary_tile_init();
add_binary_tile(0, 1, 0);
km::broadcast_done(kRowBuffer, 1);
}
}
ok = ok && km::broadcast_end(compute, kRowBuffer, 1, window, block_count);
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();
});
Every core, the sender too, adds the blocks it receives into one tile and packs the sum.
arguments = [source.address, BLOCKS, destination.address]
handle = run.submit("broadcast", group=(0, CORES), args=arguments, exclusive=True)
event = run.wait(handle, timeout=10.0)
print(f"{BLOCKS} blocks to {CORES} cores in {event.microseconds:.1f} microseconds")
One command, exclusive=True, to the six cores.
The whole program¶
"""Broadcast: core 0 reads four tiles and sends each to the other cores of its row."""
from pathlib import Path
import numpy as np
import komorebi as km
from komorebi import layout
HERE = Path(__file__).resolve().parent
CORES = 6 # the left panel of row 0
BLOCKS = 4
WINDOW = 2
with km.Session() as session:
numbers = np.random.default_rng(1).integers(-8, 9, (BLOCKS * 32, 32))
source = session.tensor(numbers, "bfloat16", "tile")
destination = session.tensor(np.zeros((CORES * 32, 32)), "bfloat16", "tile")
nocs = session.nocs()
group = layout.Group(0, CORES)
words = layout.broadcast_args(
nocs.grid, nocs.panels, group, (0, 0, CORES, 1), "row", 0,
nocs.noc0, nocs.noc_xy, (0, 0, WINDOW, 1),
) # fmt: skip
fields = km.abi.struct_fields(km.abi.headers().device[-1], "BroadcastArgs")
def arguments_of(core):
"""The program's arguments of a core: the reader's, the writer's, compute's."""
if core not in words:
return [0] * len(fields), [], [WINDOW]
return [words[core][field] for field in fields], [], [WINDOW]
tile = dict(format="bfloat16", page_bytes=source.page_bytes)
program = km.Program(
"broadcast",
kernels=HERE / "kernels",
handlers=HERE / "kernels" / "handlers.h",
buffers=[
km.Buffer(index=0, pages=WINDOW, **tile),
km.Buffer(index=16, pages=1, **tile),
],
args=arguments_of,
compute={"math_fidelity": "HiFi4", "fp32_dest_acc_en": True},
compile_args={
"reader": source.accessor_args(),
"writer": destination.accessor_args(),
},
)
run = session.load(program)
arguments = [source.address, BLOCKS, destination.address]
handle = run.submit("broadcast", group=(0, CORES), args=arguments, exclusive=True)
event = run.wait(handle, timeout=10.0)
print(f"{BLOCKS} blocks to {CORES} cores in {event.microseconds:.1f} microseconds")
run.stop()
total = destination.read().reshape(CORES, 32, 32)
expected = numbers.reshape(BLOCKS, 32, 32).sum(axis=0)
assert all((total[core] == expected).all() for core in range(CORES))
print(f"each of the {CORES} cores holds the sum of the {BLOCKS} blocks")
// What the kernels and Python agree on: the handler, and the buffers in L1: the
// broadcast buffer, in which every core of the row receives the blocks, and the output.
#pragma once
#include <stdint.h>
constexpr uint32_t kHandlerBroadcast = 1;
constexpr uint32_t kRowBuffer = 0;
constexpr uint32_t kOutput = 16;
// The reader kernel: the sender of the row reads each block from DRAM and multicasts
// it to the row; every other core waits for each block.
#include "api/dataflow/dataflow_api.h"
#include "komorebi/device/km_broadcast.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) != kHandlerBroadcast) {
return;
}
km::Broadcast row(reader, command, 0);
if (!row.ok() || !row.included()) {
return;
}
const auto source = TensorAccessor(source_arguments, command.a0);
const uint32_t block_count = command.a1;
for (uint32_t block = 0; block < block_count; block++) {
if (!row.sender()) {
if (!row.arrive(block)) {
return;
}
continue;
}
const uint32_t address = row.reserve(block);
if (address == 0) {
return;
}
noc_async_read_page(block, source, address);
if (!row.cast(block)) {
return;
}
}
row.finish();
});
}
// The compute kernel: adds the blocks that arrive into one tile.
#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/eltwise_unary/fill.h"
#include "api/compute/pack.h"
#include "api/compute/reg_api.h"
#include "api/compute/tile_move_copy.h"
#include "komorebi/device/km_broadcast.hpp"
#include "handlers.h"
void kernel_main() {
km::Compute compute;
const uint32_t window = compute.arg(0);
compute_kernel_hw_startup(kRowBuffer, kOutput);
compute.run([&](const km_command_t& command) {
if (km::handler(command) != kHandlerBroadcast) {
return;
}
const uint32_t block_count = command.a1;
tile_regs_acquire();
fill_tile_init();
fill_tile(0, 0.0f);
bool ok = true;
for (uint32_t block = 0; ok && block < block_count; block++) {
ok = km::broadcast_receive(compute, kRowBuffer, 1);
if (ok) {
copy_init(kRowBuffer);
copy_tile(kRowBuffer, 0, 1);
add_binary_tile_init();
add_binary_tile(0, 1, 0);
km::broadcast_done(kRowBuffer, 1);
}
}
ok = ok && km::broadcast_end(compute, kRowBuffer, 1, window, block_count);
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: writes the core's sum to the tile of the destination that is its
// rank.
#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) == kHandlerBroadcast) {
const uint32_t rank = writer.me() - km::first(command);
const auto destination = TensorAccessor(destination_arguments, command.a2);
if (!writer.wait_front(kOutput, 1)) {
return false;
}
noc_async_write_page(rank, destination, get_read_ptr(kOutput));
noc_async_write_barrier();
cb_pop_front(kOutput, 1);
}
return true;
});
}
Only core 0 read the source. The script checks, for each of the six cores, that the tile it wrote to the destination is the sum of the four blocks of the source: every core received every block.
Common mistakes¶
- A group that is not the rectangle the host's numbers were made for. The numbers name the
group (
km::Broadcast::ok()is false otherwise), and the core does nothing. - A broadcast across two panels of the card, or across a row that is not in one panel. The
host's
broadcast_argsrefuses a row whose cores are not next to each other on the NoC. - Forgetting
broadcast_donein the compute kernel. The sender never gets the markers that free a slot, and gives up waiting: the card stops the command at its deadline. - A buffer that does not hold
WINDOW * PAGESpages. The buffer of a broadcast is used only through these calls and must have exactly that many.
Next: lesson 28, A routed message sends to a core that the data picks.