Skip to content

Part 3. Data · lesson 23 of 42

23. Both movers

A core has two data-movement CPUs. So far one of them read, with reader.cpp, and the other wrote, with writer.cpp. Here both take part in one stream of tiles: a pattern says which of the two moves each block, and each moves its blocks through a buffer of its own. This lesson relies on lesson 16, Reading from DRAM, lesson 17, Writing to DRAM and lesson 19, Circular buffers.

Why both

Each data-movement CPU sends its requests over its own NoC: the first over NoC 0, the second over NoC 1 (lesson 4, Hello, core). One CPU that reads tile after tile keeps one NoC busy and leaves the other idle, and on one core it is the CPU that sets the pace, not DRAM. Two CPUs that each read every other block keep two reads in flight on two NoCs.

reader.cpp and writer.cpp are roles: a CPU loaded with reader.cpp reads because your code tells it to, not because it can do nothing else. When one stream is the work of a command, both CPUs can carry it. In this part they are called the movers: mover A is the CPU of reader.cpp, on NoC 0, and mover B the CPU of writer.cpp, on NoC 1.

The general form

A stream is cut into blocks of a few tiles. A pattern is a word of the letters A and B: block k goes to the mover of letter k, counted round the word. "AB" gives the movers every other block; "AAAB" gives mover A three blocks of every four.

km.Pattern(LETTERS).word  # a command argument
km.BufferPair(INDEX_A, INDEX_B, DTYPE, page_bytes=BYTES, pages=PAGES)
km.Program(NAME, kernels=FOLDER, compile_args={"movement": ARGS}, ...)

km::Movement movement;
const km::Stream stream(DIRECTION, TENSOR, FIRST_TILE, TILES, BLOCK_TILES, 1,
                        {INDEX_A, INDEX_B}, km::Pattern::of(WORD));
movement.move(stream);  // this mover's blocks of the stream; false: it gave up
  • FOLDER holds movement.cpp and compute.cpp. The one movement.cpp is built for both data-movement CPUs; km::Movement is a km::Reader on mover A and a km::Writer on mover B, so the loop of lesson 4, Hello, core runs as before.
  • A buffer pair is two circular buffers, one for each mover: its halves. km.BufferPair declares both with one format and page size, PAGES pages each.
  • DIRECTION is km::kMoveRead (tiles from DRAM into the mover's half) or km::kMoveWrite (tiles from its half to DRAM). TILES tiles from FIRST_TILE make the stream, BLOCK_TILES to a block; the last block may be shorter.
  • Both movers call move with the same stream. Each moves only the blocks the pattern gives it, waits for room or data in its own half, and returns when its blocks are done.
  • A half is filled by one CPU and emptied by one CPU for as long as the program is loaded, as every buffer is (lesson 19, Circular buffers).

One line at a time

movement.run([&](const km_command_t& command) {
    if (km::handler(command) == kHandlerCopy) {
        const uint32_t source_address = command.a0;
        const uint32_t first_tile = command.a1;
        const uint32_t destination_address = command.a2;
        const uint32_t first_destination_tile = command.a3;
        const uint32_t tile_count = command.a4;
        const km::Pattern pattern = km::Pattern::of(command.a5);
        const km::Stream read(
            km::kMoveRead,
            TensorAccessor(arguments, source_address),
            first_tile,
            tile_count,
            kBlockTiles,
            1,
            halves,
            pattern
        );
        const km::Stream write(
            km::kMoveWrite,
            TensorAccessor(arguments, destination_address),
            first_destination_tile,
            tile_count,
            kBlockTiles,
            1,
            halves,
            pattern
        );
        return movement.move(read, write);
    }
    return true;
});

This is the copy of lesson 19, Circular buffers, with the loop of reader.cpp and the loop of writer.cpp replaced by two streams. read brings the tiles from the source into the halves, write sends them on from the halves to the destination. move takes both: each mover reads one of its blocks into its half, then writes it on from there, so no compute kernel stands between them and each half is filled and emptied by its own mover.

program = km.Program(
    "both_movers",
    kernels=HERE / "kernels",  # movement.cpp and compute.cpp
    handlers=HERE / "kernels" / "handlers.h",
    buffers=[km.BufferPair(0, 1, "uint32", page_bytes=source.page_bytes, pages=16)],
    compile_args={"movement": source.accessor_args()},
)

The pair holds 16 pages in each half: two blocks of 8 tiles, so that a mover can read its next block while it writes the last one. "movement" gives the same tensor arguments to both movers.

for round_number, letters in enumerate(("A", "AB")):
    pattern = km.Pattern(letters)
    first = round_number * TILES
    arguments = [source.address, 0, destination.address, first, TILES, pattern.word]
    handle = run.submit("copy", group=(0, 1), args=arguments)
    microseconds = run.wait(handle, timeout=10.0).microseconds
    print(
        f"pattern {letters:2}: {TILES} tiles copied in "
        f"{microseconds:6.1f} microseconds"
    )

The same command twice, on one core: with pattern "A" mover A copies every block alone, with "AB" the movers take every other block.

The whole program

"""Both movers: one core copies 256 tiles with one data-movement CPU, then with both."""

from pathlib import Path

import numpy as np

import komorebi as km

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

with km.Session() as session:
    numbers = np.arange(TILES * 32 * 32, dtype=np.uint32).reshape(TILES * 32, 32)
    source = session.tensor(numbers, "uint32", "tile")
    destination = session.tensor(
        np.zeros((2 * TILES * 32, 32), np.uint32), "uint32", "tile"
    )

    program = km.Program(
        "both_movers",
        kernels=HERE / "kernels",  # movement.cpp and compute.cpp
        handlers=HERE / "kernels" / "handlers.h",
        buffers=[km.BufferPair(0, 1, "uint32", page_bytes=source.page_bytes, pages=16)],
        compile_args={"movement": source.accessor_args()},
    )
    run = session.load(program)

    for round_number, letters in enumerate(("A", "AB")):
        pattern = km.Pattern(letters)
        first = round_number * TILES
        arguments = [source.address, 0, destination.address, first, TILES, pattern.word]
        handle = run.submit("copy", group=(0, 1), args=arguments)
        microseconds = run.wait(handle, timeout=10.0).microseconds
        print(
            f"pattern {letters:2}: {TILES} tiles copied in "
            f"{microseconds:6.1f} microseconds"
        )
    run.stop()

    (source_pages,) = session.read(source)
    (destination_pages,) = session.read(destination)
    for round_number in range(2):
        copied = destination_pages[round_number * TILES : (round_number + 1) * TILES]
        assert (copied == source_pages).all()
    print("both copies hold every tile of the source")

// What the kernels and Python agree on: the handler, the two halves of the buffer pair
// (one for each data-movement CPU), and how many tiles make a block.
#pragma once

#include <stdint.h>

constexpr uint32_t kHandlerCopy = 1;
constexpr uint32_t kHalfA = 0;
constexpr uint32_t kHalfB = 1;
constexpr uint32_t kBlockTiles = 8;

// The data-movement kernel of both CPUs: each copies the blocks the pattern gives it,
// reading them into its own half of the pair and writing them on from there.
#include "api/dataflow/dataflow_api.h"
#include "komorebi/device/km_movement.hpp"
#include "handlers.h"

void kernel_main() {
    constexpr auto arguments = TensorAccessorArgs<0>();  // source and destination alike
    constexpr km::BufferPair halves = {kHalfA, kHalfB};
    km::Movement movement;
    movement.run([&](const km_command_t& command) {
        if (km::handler(command) == kHandlerCopy) {
            const uint32_t source_address = command.a0;
            const uint32_t first_tile = command.a1;
            const uint32_t destination_address = command.a2;
            const uint32_t first_destination_tile = command.a3;
            const uint32_t tile_count = command.a4;
            const km::Pattern pattern = km::Pattern::of(command.a5);
            const km::Stream read(
                km::kMoveRead,
                TensorAccessor(arguments, source_address),
                first_tile,
                tile_count,
                kBlockTiles,
                1,
                halves,
                pattern
            );
            const km::Stream write(
                km::kMoveWrite,
                TensorAccessor(arguments, destination_address),
                first_destination_tile,
                tile_count,
                kBlockTiles,
                1,
                halves,
                pattern
            );
            return movement.move(read, write);
        }
        return true;
    });
}

// The compute kernel. It has nothing to do: the tiles go from one half to DRAM.
#include "komorebi/device/km_worker.hpp"

void kernel_main() {
    km::Compute compute;
    compute.run([&](const km_command_t& command) {});
}

pattern A : 256 tiles copied in   39.6 microseconds
pattern AB: 256 tiles copied in   21.6 microseconds
both copies hold every tile of the source

One mover copied the 256 tiles in 39.6 microseconds, both in 21.6: nearly twice as fast, as two CPUs each did half of the reads and half of the writes on two NoCs. The last line is the program's check: both copies hold every tile of the source.

Common mistakes

  • Expecting both movers to be faster everywhere. On one core they are. On many cores the two NoCs reach DRAM at different rates, and the best pattern follows them: on the whole card "AAAB" reads fastest (Movement).
  • Giving a mover's stream to a buffer the compute kernel fills in another command. A half keeps the CPU that fills it and the CPU that empties it for the whole load; move gives up if they change.
  • A buffer that is not a whole number of blocks. Each block is one item of BLOCK_TILES pages, and a half must hold a whole number of them.

A compute kernel that takes blocks from a pair takes them in order, block k from the half km::block_half(pair, pattern, k); the Movement reference page has the rules.

That is the end of Part 3. A core can now read tiles from DRAM, work on them with the compute kernel and the matrix unit, write them back with one or both of its data-movement CPUs, and a group can do that for a whole tensor. Part 4 lets cores pass tiles to each other, without a trip to DRAM.

Next: lesson 24, Why cores talk.