Part 6. Practice · lesson 41 of 42
41. Your own kernel¶
You write a kernel of your own by changing a kernel that works, one thing at a time, with the right answer worked out in numpy before the card runs. This lesson changes the program of lesson 40, A whole program so that it sums the squares of the tiles. It relies on lesson 40, A whole program, lesson 39, Finding a bug and lesson 38, Measuring.
Why change, and not write¶
A new kernel written from nothing has many places to be wrong at once: the buffers, the partition, the reduction, the arithmetic. If it does not work, nothing says which. A program that works, changed in one place, has one place to look. And a check that is made before the change tells you the change did what you meant.
The method¶
- Copy the folder of a program that works. Run it, to be sure that it works here.
- Decide the one change, and say what the answer will be after it. Compute that answer in numpy, from the same input.
- Make the change. Change the handler's name if the work is no longer what the name says.
- Run. The program compares the card's answer with numpy's and says where they differ.
- If they differ, go back to what you changed. If you want to change something else, go to step 2.
Two rules keep the steps short. Change only one thing between two runs. And check with
numbers that cannot hide a mistake: small whole numbers, whose sums bfloat16 holds
exactly, and not random fractions that you can only compare "about".
One line at a time¶
The change is in the compute kernel. In lesson 40, A whole program each tile was added to the sum as it was. Now each tile is first multiplied by itself:
compute.run([&](const km_command_t& command) {
if (km::handler(command) != kHandlerSumSquares) {
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) {
const uint32_t into = index == 0 ? 0 : 1;
copy_init(kOwnTiles);
copy_tile(kOwnTiles, 0, into);
mul_binary_tile_init();
mul_binary_tile(into, into, into);
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 three lines that are new are const uint32_t into = ..., mul_binary_tile_init() and
mul_binary_tile(into, into, into). into is the register the tile goes to: 0 for the first
tile, 1 for the others. mul_binary_tile(a, b, c) multiplies the tiles in registers a and
b, number by number, into register c; with the same register for all three, it squares the
tile where it lies. The addition that follows is the old one. reader.cpp, writer.cpp, and the
reduction between cores are unchanged; only the handler's name is new,
sum_squares.
times = []
for cores, destination in zip((1, 8), destinations):
arguments = [source.address, 0, destination.address, 0, 0, TILES]
handle = run.submit(
"sum_squares", group=(0, cores), args=arguments, exclusive=True
)
times.append(run.wait(handle, timeout=10.0).microseconds)
run.stop()
expected = (numbers**2).reshape(TILES, 32, 32).sum(axis=0)
for cores, microseconds, destination in zip((1, 8), times, destinations):
print(
f"{cores} cores: squares of {TILES} tiles added in {microseconds:5.1f} us"
)
total = destination.read()
wrong = total != expected
if wrong.any():
print(
f"{wrong.sum()} of {wrong.size} numbers differ from numpy; the first:"
)
print(f" the card made {total[wrong][0]}, numpy made {expected[wrong][0]}")
raise SystemExit(1)
On the Python side, two things changed: the name of the handler, and the expected answer,
which is now the sum of the squares. The input is the integers -2 to 2, so every square is at
most 4 and the total of 48 of them is at most 192, which bfloat16 holds exactly.
The check also changed: when the answers differ, the program says how many numbers differ and shows the first. That is the line to read when step 4 fails.
The whole program¶
"""Your own kernel: the sum of the squares of 48 tiles, 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(-2, 3, (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(
"your_own_kernel",
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_squares", group=(0, cores), args=arguments, exclusive=True
)
times.append(run.wait(handle, timeout=10.0).microseconds)
run.stop()
expected = (numbers**2).reshape(TILES, 32, 32).sum(axis=0)
for cores, microseconds, destination in zip((1, 8), times, destinations):
print(
f"{cores} cores: squares of {TILES} tiles added in {microseconds:5.1f} us"
)
total = destination.read()
wrong = total != expected
if wrong.any():
print(
f"{wrong.sum()} of {wrong.size} numbers differ from numpy; the first:"
)
print(f" the card made {total[wrong][0]}, numpy made {expected[wrong][0]}")
raise SystemExit(1)
print("both sums are the numpy sum of the squares of the 48 tiles")
// The compute kernel: squares each tile of the core's share and adds the squares 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) != kHandlerSumSquares) {
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) {
const uint32_t into = index == 0 ? 0 : 1;
copy_init(kOwnTiles);
copy_tile(kOwnTiles, 0, into);
mul_binary_tile_init();
mul_binary_tile(into, into, into);
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 other kernels are those of lesson 40, A whole program, with the handler renamed. Run it from
the folder tutorial/41_your_own_kernel:
1 cores: squares of 48 tiles added in 25.8 us
8 cores: squares of 48 tiles added in 7.1 us
both sums are the numpy sum of the squares of the 48 tiles
It worked, and the times are those of the sum: squaring costs the cores almost nothing next to reading the tiles.
When the check fails¶
To see what a failed check looks like, the same program was run with the three lines of
mul_binary_tile taken out of the kernel, so that it only adds, and numpy still expected
the squares. The program ended with:
1 cores: squares of 48 tiles added in 25.6 us
1024 of 1024 numbers differ from numpy; the first:
the card made -17.0, numpy made 103
All 1024 numbers differ, and the card's first number is negative, which a sum of squares can never be. So the squares are missing, and the place to look is the change you just made. A wrong answer in every number usually points at the work done on every tile; a wrong answer in a few numbers points more often at the places where the cores meet.
Common mistakes¶
- Changing two things at once. If the answer is wrong, you do not know which.
- Working out the expected answer after seeing the card's. It is then easy to believe the card's.
- Checking with random fractions. A difference in the last digit looks like an error or like rounding, and you cannot tell which. Use small whole numbers.
- Keeping the old handler name for new work.
sum_tilesthat sums squares misleads the next person who reads it, and you.
Next: lesson 42, Every primitive lists every call that this tutorial used, with its lesson.