Part 3. Data · lesson 17 of 42
17. Writing to DRAM¶
A data-movement kernel can copy a tile from the core's L1 to a tile of a tensor in DRAM,
and that is how a program puts something where Python can read it. In this lesson it is
writer.cpp that does it, and reader.cpp that reads. This lesson relies on
lesson 16, Reading from DRAM and lesson 15, Tiles and tensors.
Why a kernel writes¶
A result that is a single number can travel back in the result word (lesson 6, A result). A result that is a tile cannot: it is 1024 numbers. Its way out is DRAM: a data-movement kernel copies it to a tensor that Python made, and Python reads the tensor. With a read and a write, a program can now change data, and the smallest program that does is a copy: read a tile in, write the tile out somewhere else.
The general form¶
In writer.cpp:
constexpr auto ARGUMENTS = TensorAccessorArgs<0>();
...
const auto tensor = TensorAccessor(ARGUMENTS, ADDRESS);
writer.wait_front(INDEX, 1);
noc_async_write_page(TILE, tensor, get_read_ptr(INDEX));
noc_async_write_barrier();
cb_pop_front(INDEX, 1);
In Python, one more entry of compile_args:
ARGUMENTSandTensorAccessor(...)are as in lesson 16, Reading from DRAM, here for the destination.ADDRESSisdestination.address, passed as a command argument.writer.wait_front(INDEX, 1)waits until a tile is in the buffer.noc_async_write_page(TILE, tensor, where)starts a copy of the L1 data atwhereinto tile numberTILEof the tensor.get_read_ptr(INDEX)is the address of the tile in the buffer.noc_async_write_barrier()waits until every write that was started is finished.cb_pop_front(INDEX, 1)says "I am done with the tile", and frees the room for the reader.
One line at a time¶
The tensor to write into, made of zeros: two tiles. Python makes it before load, just
like the source.
writer.run([&](const km_command_t& command) {
if (km::handler(command) == kHandlerCopy) {
const uint32_t destination_address = command.a2;
const uint32_t destination_tile = command.a3;
const auto destination = (
TensorAccessor(destination_arguments, destination_address)
);
if (!writer.wait_front(kTileBuffer, 1)) {
return false;
}
noc_async_write_page(
destination_tile, destination, get_read_ptr(kTileBuffer)
);
noc_async_write_barrier();
cb_pop_front(kTileBuffer, 1);
}
return true;
});
writer.cpp's half of the copy. The command says where the destination is (argument 2) and
which tile of it to write (argument 3). Once the tile is in the buffer, writer.cpp copies
it out, waits until the write is done, and frees the buffer. reader.cpp is that of
lesson 16, Reading from DRAM: arguments 0 and 1 are the source and its tile.
for source_tile, destination_tile in [(2, 0), (3, 1)]:
handle = run.submit(
"copy",
group=(0, 1),
args=[source.address, source_tile, destination.address, destination_tile],
)
run.wait(handle, timeout=10.0)
Two commands: tile 2 of the source to tile 0 of the destination, and tile 3 to tile 1. Each one waits before the next: the buffer holds one tile.
(source_pages,) = session.read(source)
(destination_pages,) = session.read(destination)
first = destination_pages.view(np.uint32)[:, 0].tolist()
print(f"the first number of each destination tile: {first}")
assert (destination_pages[0] == source_pages[2]).all()
assert (destination_pages[1] == source_pages[3]).all()
print("tiles 2 and 3 of the source are tiles 0 and 1 of the destination")
After stop, Python reads the raw pages of both tensors with session.read, prints the
first number of each destination tile, and checks that each page is the same bytes as
the source page it was copied from.
The whole program¶
"""Writing to DRAM: a core copies tiles from one tensor to another."""
from pathlib import Path
import numpy as np
import komorebi as km
HERE = Path(__file__).resolve().parent
with km.Session() as session:
numbers = np.arange(4 * 32 * 32, dtype=np.uint32).reshape(4 * 32, 32) + 1000
source = session.tensor(numbers, "uint32", "tile")
destination = session.tensor(np.zeros((64, 32), np.uint32), "uint32", "tile")
program = km.Program(
"writing_to_dram",
kernels=HERE / "kernels",
handlers=HERE / "kernels" / "handlers.h",
buffers=[
km.Buffer(index=0, format="uint32", page_bytes=source.page_bytes, pages=1)
],
compile_args={
"reader": source.accessor_args(),
"writer": destination.accessor_args(),
},
)
run = session.load(program)
for source_tile, destination_tile in [(2, 0), (3, 1)]:
handle = run.submit(
"copy",
group=(0, 1),
args=[source.address, source_tile, destination.address, destination_tile],
)
run.wait(handle, timeout=10.0)
run.stop()
(source_pages,) = session.read(source)
(destination_pages,) = session.read(destination)
first = destination_pages.view(np.uint32)[:, 0].tolist()
print(f"the first number of each destination tile: {first}")
assert (destination_pages[0] == source_pages[2]).all()
assert (destination_pages[1] == source_pages[3]).all()
print("tiles 2 and 3 of the source are tiles 0 and 1 of the destination")
// The reader kernel: copies one tile from DRAM into the buffer in L1.
#include "api/dataflow/dataflow_api.h"
#include "komorebi/device/km_worker.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) == kHandlerCopy) {
const uint32_t source_address = command.a0;
const uint32_t source_tile = command.a1;
const auto source = TensorAccessor(source_arguments, source_address);
if (!reader.reserve(kTileBuffer, 1)) {
return;
}
noc_async_read_page(source_tile, source, get_write_ptr(kTileBuffer));
noc_async_read_barrier();
cb_push_back(kTileBuffer, 1);
}
});
}
// The writer kernel: waits for the tile in L1 and writes it to a tile of another
// tensor.
#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) == kHandlerCopy) {
const uint32_t destination_address = command.a2;
const uint32_t destination_tile = command.a3;
const auto destination = (
TensorAccessor(destination_arguments, destination_address)
);
if (!writer.wait_front(kTileBuffer, 1)) {
return false;
}
noc_async_write_page(
destination_tile, destination, get_read_ptr(kTileBuffer)
);
noc_async_write_barrier();
cb_pop_front(kTileBuffer, 1);
}
return true;
});
}
The compute file is the empty one.
the first number of each destination tile: [3048, 4072]
tiles 2 and 3 of the source are tiles 0 and 1 of the destination
The numbers of the source count up from 1000, row by row, so tile 2 starts at 1000 + 2 x 1024 = 3048, and tile 3 at 4072. The destination began as zeros; now its first tile starts with 3048 and its second with 4072, and the script has checked that every byte is the same.
Common mistakes¶
- Reading the tensor too early.
session.readdoes not wait for the cores. Read after the commands have been waited for (lesson 5, Commands). - Writing a tile number the destination does not have. The tensor has two tiles: tile 2 is past its end.
- Forgetting
noc_async_write_barrier(). The command may be reported done before the tile is in DRAM.
Next: lesson 18, The three kernels looks at the three kernels as a whole.