Part 3. Data · lesson 16 of 42

16. Reading from DRAM¶

A data-movement kernel can copy one tile of a tensor from DRAM into the core's L1, where the core can use it. Which of the two data-movement kernels does it is the program's choice (lesson 4, Hello, core); in this lesson it is reader.cpp. This lesson relies on lesson 15, Tiles and tensors, lesson 7, Arguments and lesson 3, The card.

Why the core must copy¶

A core works only on what is in its own L1 (lesson 3, The card). Data in DRAM is out of its reach until it is copied in. A data-movement CPU does the copy, over the NoC: it asks for tile number so-and-so of a tensor and says where in L1 the tile is to land.

A copy takes time and is not finished when the call returns. The kernel has to wait for it, and then has to tell the others that the tile has arrived.

The general form¶

The kernel needs three things: a place in L1 for the tile, a way to name the tensor, and the copy itself.

In Python, before load:

program = km.Program(
    NAME,
    kernels=FOLDER,
    handlers=HEADER,
    buffers=[km.Buffer(index=INDEX, format=DTYPE, page_bytes=BYTES, pages=PAGES)],
    compile_args={"reader": TENSOR.accessor_args()},
)

In reader.cpp:

constexpr auto ARGUMENTS = TensorAccessorArgs<0>();
...
const auto tensor = TensorAccessor(ARGUMENTS, ADDRESS);
reader.reserve(INDEX, PAGES);
noc_async_read_page(TILE, tensor, get_write_ptr(INDEX));
noc_async_read_barrier();
cb_push_back(INDEX, PAGES);

km.Buffer(...) sets aside room in the L1 of every core. INDEX is its number, 0 to 31, and DTYPE and BYTES say what one tile of the tensor is: "uint32" and 4096 bytes for the tiles of the last lesson. PAGES is how many tiles fit; this lesson uses one. The cb_ in the C++ names below belongs to the kind of buffer that lesson 19, Circular buffers explains. For now a buffer is a numbered room in L1.
compile_args hands the kernel what it needs to find the tensor's pages. In the kernel, TensorAccessorArgs<0>() receives it, and TensorAccessor(ARGUMENTS, ADDRESS) is an object that turns a tile number into a place in DRAM. ADDRESS is where the tensor starts. Python has it as tensor.address; it travels in a command argument.
reader.reserve(INDEX, PAGES) waits until the buffer has room for PAGES tiles.
noc_async_read_page(TILE, tensor, where) starts a copy of tile number TILE of the tensor to the L1 address where. get_write_ptr(INDEX) is the address of the room in the buffer. The call returns at once; the copy goes on.
noc_async_read_barrier() waits until every copy that was started is finished.
cb_push_back(INDEX, PAGES) says "these tiles are now in the buffer". Whoever waits for them is let through.

One tile goes from DRAM to the L1 of core 0; `writer.cpp` sends a number of it to Python.

One line at a time¶

The tensor first, in Python:

tile_numbers = np.arange(4, dtype=np.uint32).reshape(4, 1, 1)
positions = np.arange(32 * 32, dtype=np.uint32).reshape(1, 32, 32)
source = session.tensor(
    (tile_numbers * 1000 + positions).reshape(128, 32), "uint32", "tile"
)

Four tiles; the number at place p of tile t is 1000 * t + p, so that a number tells where it came from. (positions is 0 to 1023, and the numpy shapes line up tile numbers and positions.)

program = km.Program(
    "reading_from_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()},
)

The program names the buffer and gives reader.cpp what it needs to find source. The buffer holds one tile of 4096 bytes in the format of the tensor.

reader.run([&](const km_command_t& command) {
    if (km::handler(command) == kHandlerFetch) {
        const uint32_t source_address = command.a0;
        const uint32_t tile_number = command.a1;
        const auto source = TensorAccessor(source_arguments, source_address);
        if (!reader.reserve(kTileBuffer, 1)) {
            return;
        }
        noc_async_read_page(tile_number, source, get_write_ptr(kTileBuffer));
        noc_async_read_barrier();
        cb_push_back(kTileBuffer, 1);
    }
});

The command's arguments say which tensor and which tile. reader.cpp makes the accessor, waits for room, starts the copy, waits for it to finish, and announces the tile.

writer.run([&](const km_command_t& command) {
    if (km::handler(command) == kHandlerFetch) {
        const uint32_t position = command.a2;
        if (!writer.wait_front(kTileBuffer, 1)) {
            return false;
        }
        const auto tile = (
            reinterpret_cast<volatile tt_l1_ptr uint32_t*>(
                get_read_ptr(kTileBuffer)
            )
        );
        writer.result(tile[position]);
        cb_pop_front(kTileBuffer, 1);
    }
    return true;
});

writer.cpp is the one that can send a result (lesson 6, A result), so it takes the number out of the tile. wait_front waits until reader.cpp has announced the tile; get_read_ptr is the address of the tile in L1; the reinterpret_cast says "read it as an array of uint32_t"; tile[position] is the number at that place. cb_pop_front frees the room for the next command.

for tile, position in [(2, 0), (2, 5), (3, 0)]:
    handle = run.submit(
        "fetch", group=(0, 1), args=[source.address, tile, position], result=True
    )
    number = run.wait(handle, timeout=10.0).result
    print(f"tile {tile}, number {position}: {number}")
    assert number == tile * 1000 + position

Three commands, each fetching one tile and asking for one number of it.

The whole program¶

reading_from_dram.pykernels/handlers.hkernels/reader.cppkernels/writer.cpp

"""Reading from DRAM: a core copies a tile into its L1; the writer reads a number."""

from pathlib import Path

import numpy as np

import komorebi as km

HERE = Path(__file__).resolve().parent

with km.Session() as session:
    tile_numbers = np.arange(4, dtype=np.uint32).reshape(4, 1, 1)
    positions = np.arange(32 * 32, dtype=np.uint32).reshape(1, 32, 32)
    source = session.tensor(
        (tile_numbers * 1000 + positions).reshape(128, 32), "uint32", "tile"
    )

    program = km.Program(
        "reading_from_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()},
    )
    run = session.load(program)

    for tile, position in [(2, 0), (2, 5), (3, 0)]:
        handle = run.submit(
            "fetch", group=(0, 1), args=[source.address, tile, position], result=True
        )
        number = run.wait(handle, timeout=10.0).result
        print(f"tile {tile}, number {position}: {number}")
        assert number == tile * 1000 + position
    run.stop()

// What the kernels and Python agree on: the handler, and the number of the buffer in L1
// in which the reader puts the tile.
#pragma once

#include <stdint.h>

constexpr uint32_t kHandlerFetch = 1;
constexpr uint32_t kTileBuffer = 0;

// 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) == kHandlerFetch) {
            const uint32_t source_address = command.a0;
            const uint32_t tile_number = command.a1;
            const auto source = TensorAccessor(source_arguments, source_address);
            if (!reader.reserve(kTileBuffer, 1)) {
                return;
            }
            noc_async_read_page(tile_number, source, get_write_ptr(kTileBuffer));
            noc_async_read_barrier();
            cb_push_back(kTileBuffer, 1);
        }
    });
}

// The writer kernel: waits for the tile in L1 and sends one number of it back.
#include "api/dataflow/dataflow_api.h"
#include "komorebi/device/km_worker.hpp"
#include "handlers.h"

void kernel_main() {
    km::Writer writer;
    writer.run([&](const km_command_t& command) {
        if (km::handler(command) == kHandlerFetch) {
            const uint32_t position = command.a2;
            if (!writer.wait_front(kTileBuffer, 1)) {
                return false;
            }
            const auto tile = (
                reinterpret_cast<volatile tt_l1_ptr uint32_t*>(
                    get_read_ptr(kTileBuffer)
                )
            );
            writer.result(tile[position]);
            cb_pop_front(kTileBuffer, 1);
        }
        return true;
    });
}

The compute file is the empty one.

tile 2, number 0: 2000
tile 2, number 5: 2005
tile 3, number 0: 3000

Each number is 1000 * tile + position: the tile that was asked for arrived, and its number at that place. The script checks it. It asks only for positions below 16 because that is where the row-by-row order of the host array and the order of the numbers in the tile agree; beyond them a tile keeps its numbers in four blocks of 16 x 16.

Common mistakes¶

Starting the copy and not waiting. noc_async_read_page returns at once, and without noc_async_read_barrier() the tile may not be in L1 yet when it is used.
Typing page_bytes by hand. It must be the size of one tile of the tensor; write source.page_bytes instead of 4096 and it cannot go wrong when the dtype changes.
Forgetting cb_pop_front in writer.cpp. The first command works. The second finds the buffer still full, reader.cpp waits for room that never comes, and the command ends with CommandTimeout (lesson 10, When it goes wrong).
Taking a number of the tile at a place that is not in the first row of the first 16 columns, and expecting the host array's order. The numbers inside a tile are kept in four blocks of 16 x 16.

Next: lesson 17, Writing to DRAM sends a tile the other way.