Part 6. Practice · lesson 40 of 42
40. A whole program¶
A whole program has the same few steps whatever it computes, and each step is something an earlier lesson taught. This lesson is one program, the sum of 48 tiles, read from start to end with the lesson of every part. It relies on all of Parts 1 to 4.
Why one more program¶
Every program so far showed one idea, and used the others only as far as it had to. A real program is the ideas side by side, and the question becomes where each one goes. The answer is the same each time. The data is made in Python. The cores' work is written once, in three kernels. The program is loaded once. Commands carry the numbers that change. The result is read after the cores have stopped, and checked.
The plan¶
| Step | In this program | Lesson |
|---|---|---|
| Make the data | 48 tiles of small integers in numpy | lesson 15, Tiles and tensors |
| Put it on the card | session.tensor(...) for the source, and one for each result |
lesson 16, Reading from DRAM |
| Describe the cores' work | three kernels, a handler sum_tiles, buffers |
lesson 8, Handlers, lesson 18, The three kernels |
| Share the tiles | each core takes its part of the 48 | lesson 12, Partitions |
| Move tiles through L1 | reader.cpp pushes, the compute kernel pops |
lesson 19, Circular buffers |
| Compute | each core adds its tiles into one | lesson 20, Computing on a tile |
| Combine the cores | the sums of the cores are added into one | lesson 26, Reduce |
| Write the answer | writer.cpp of the leader writes the total |
lesson 17, Writing to DRAM, lesson 13, The leader |
| Send the command | submit with the arguments, then wait |
lesson 7, Arguments, lesson 5, Commands |
| Check | read the result after stop and compare with numpy |
lesson 6, A result |
One part at a time¶
numbers = np.random.default_rng(1).integers(-3, 4, (TILES * 32, 32))
source = session.tensor(numbers, "bfloat16", "tile")
destinations = [
session.tensor(np.zeros((32, 32)), "bfloat16", "tile") for _ in range(2)
]
Python makes 48 tiles of integers from -3 to 3, and puts them on the card as one tensor of
bfloat16 tiles. It makes two tensors for the answer, one for each of the two commands,
before the program is loaded: a tensor cannot be written while the kernels run
(lesson 22, A whole map).
tile = dict(format="bfloat16", page_bytes=source.page_bytes)
program = km.Program(
"a_whole_program",
kernels=HERE / "kernels",
handlers=HERE / "kernels" / "handlers.h",
buffers=[
km.Buffer(index=0, pages=2, **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": destinations[0].accessor_args(),
},
)
The program: the folder of the kernels, the header with the handler, three buffers, and the
arithmetic setting of lesson 20, Computing on a tile. The first buffer holds the core's own tiles on their
way from reader.cpp to the compute kernel, two at a time. The second receives the sums that
other cores send (lesson 25, A port). The third holds the output.
reader.run([&](const km_command_t& command) {
if (km::handler(command) == kHandlerSumTiles) {
const km::Group group = km::group_of(command, reader.me());
const uint32_t tile_count = command.a5;
const km::Partition share = (
km::partition(tile_count, group.rank, group.size)
);
const auto source = TensorAccessor(source_arguments, command.a0);
for (uint32_t index = 0; index < share.len; index++) {
if (!reader.reserve(kOwnTiles, 1)) {
return;
}
noc_async_read_page(
command.a1 + share.start + index, source, get_write_ptr(kOwnTiles)
);
noc_async_read_barrier();
cb_push_back(kOwnTiles, 1);
}
km::reduce_arrive(reader, command, arrival, group, 0);
}
});
reader.cpp. Its core takes its share of the 48 tiles with km::partition, reads them one by
one into the first buffer, and then takes its part in the reduction.
compute.run([&](const km_command_t& command) {
if (km::handler(command) != kHandlerSumTiles) {
return;
}
const km::Group group = km::group_of(command, compute.me());
const km::Partition share = km::partition(command.a5, group.rank, group.size);
tile_regs_acquire();
bool ok = true;
for (uint32_t index = 0; ok && index < share.len; index++) {
ok = compute.wait_front(kOwnTiles, 1);
if (ok) {
copy_init(kOwnTiles);
copy_tile(kOwnTiles, 0, index == 0 ? 0 : 1);
if (index > 0) {
add_binary_tile_init();
add_binary_tile(0, 1, 0);
}
cb_pop_front(kOwnTiles, 1);
}
}
ok = ok && km::reduce(compute, arrival, group, 0, [&](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 takes the same share, and adds its tiles into one: the first tile goes
into register 0, and each later tile into register 1 and is added to 0. Then km::reduce
adds the sums that the other cores send. The result goes to the output buffer.
writer.run([&](const km_command_t& command) {
if (km::handler(command) == kHandlerSumTiles) {
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. A core that is not the leader sends its sum to the core it reports to; the leader
writes the total into the destination tensor.
times = []
for cores, destination in zip((1, 8), destinations):
arguments = [source.address, 0, destination.address, 0, 0, TILES]
handle = run.submit(
"sum_tiles", group=(0, cores), args=arguments, exclusive=True
)
times.append(run.wait(handle, timeout=10.0).microseconds)
run.stop()
expected = numbers.reshape(TILES, 32, 32).sum(axis=0)
for cores, microseconds, destination in zip((1, 8), times, destinations):
print(f"{cores} cores: {TILES} tiles added into one in {microseconds:5.1f} us")
assert (destination.read() == expected).all()
The same command twice, for 1 core and for 8. The arguments are the source and destination addresses, the first tile of each, 0 for the unused form, and the number of tiles. When both commands are done, the program stops the cores and checks each destination against numpy.
The whole program¶
"""A whole program: the sum of 48 tiles by 1 core and by 8, checked in numpy."""
from pathlib import Path
import numpy as np
import komorebi as km
HERE = Path(__file__).resolve().parent
TILES = 48
with km.Session() as session:
numbers = np.random.default_rng(1).integers(-3, 4, (TILES * 32, 32))
source = session.tensor(numbers, "bfloat16", "tile")
destinations = [
session.tensor(np.zeros((32, 32)), "bfloat16", "tile") for _ in range(2)
]
tile = dict(format="bfloat16", page_bytes=source.page_bytes)
program = km.Program(
"a_whole_program",
kernels=HERE / "kernels",
handlers=HERE / "kernels" / "handlers.h",
buffers=[
km.Buffer(index=0, pages=2, **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": destinations[0].accessor_args(),
},
)
run = session.load(program)
times = []
for cores, destination in zip((1, 8), destinations):
arguments = [source.address, 0, destination.address, 0, 0, TILES]
handle = run.submit(
"sum_tiles", group=(0, cores), args=arguments, exclusive=True
)
times.append(run.wait(handle, timeout=10.0).microseconds)
run.stop()
expected = numbers.reshape(TILES, 32, 32).sum(axis=0)
for cores, microseconds, destination in zip((1, 8), times, destinations):
print(f"{cores} cores: {TILES} tiles added into one in {microseconds:5.1f} us")
assert (destination.read() == expected).all()
print("both sums are the numpy sum of the 48 tiles")
// What the kernels and Python agree on: the handler, and three buffers in L1: the
// core's own tiles, the arrival buffer for other cores' sums, and the output.
#pragma once
#include <stdint.h>
constexpr uint32_t kHandlerSumTiles = 1;
constexpr uint32_t kOwnTiles = 0;
constexpr uint32_t kArrival = 1;
constexpr uint32_t kArrivalPages = 8;
constexpr uint32_t kOutput = 16;
// The reader kernel: reads the tiles of the core's share of the source, one after the
// other, then takes part in the reduction.
#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) == kHandlerSumTiles) {
const km::Group group = km::group_of(command, reader.me());
const uint32_t tile_count = command.a5;
const km::Partition share = (
km::partition(tile_count, group.rank, group.size)
);
const auto source = TensorAccessor(source_arguments, command.a0);
for (uint32_t index = 0; index < share.len; index++) {
if (!reader.reserve(kOwnTiles, 1)) {
return;
}
noc_async_read_page(
command.a1 + share.start + index, source, get_write_ptr(kOwnTiles)
);
noc_async_read_barrier();
cb_push_back(kOwnTiles, 1);
}
km::reduce_arrive(reader, command, arrival, group, 0);
}
});
}
// The compute kernel: adds the tiles of the core's share into one tile, then adds the
// sums of the other cores to it.
#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(kOwnTiles, kOutput);
const km::Arrival arrival = {kArrival, kArrivalPages, 1};
compute.run([&](const km_command_t& command) {
if (km::handler(command) != kHandlerSumTiles) {
return;
}
const km::Group group = km::group_of(command, compute.me());
const km::Partition share = km::partition(command.a5, group.rank, group.size);
tile_regs_acquire();
bool ok = true;
for (uint32_t index = 0; ok && index < share.len; index++) {
ok = compute.wait_front(kOwnTiles, 1);
if (ok) {
copy_init(kOwnTiles);
copy_tile(kOwnTiles, 0, index == 0 ? 0 : 1);
if (index > 0) {
add_binary_tile_init();
add_binary_tile(0, 1, 0);
}
cb_pop_front(kOwnTiles, 1);
}
}
ok = ok && km::reduce(compute, arrival, group, 0, [&](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) == kHandlerSumTiles) {
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;
});
}
Run it from the folder tutorial/40_a_whole_program:
1 cores: 48 tiles added into one in 25.7 us
8 cores: 48 tiles added into one in 6.9 us
both sums are the numpy sum of the 48 tiles
Both sums equal numpy's. Eight cores took 6.9 microseconds against 25.7 for one: 3.7 times as fast with 8 times the cores, since the eight sums still have to be combined.
Common mistakes¶
- Writing a tensor after
load. The write waits for the cores to be idle, and they never are: after 60 seconds it raisesRuntimeErrorwithTimeout (60000 ms) waiting for physical cores to finish. Make every tensor first. - Reading a tensor before
stop. The same wait, the same error, after the same 60 seconds. Read it afterstop, as every program of this tutorial does.
Next: lesson 41, Your own kernel starts from this program and changes it.