Skip to content

Part 2. Many cores · lesson 12 of 42

12. Partitions

A partition is the share of a job that one core of a group takes: the items of the job are cut into one contiguous piece for each rank, and km::partition says where the piece of a rank starts and how long it is. This lesson relies on lesson 11, A core group and lesson 7, Arguments.

Why partitions

A group of cores has the same kernel code and the same command, so the cores have to find different work in it by the one thing that differs, the rank (lesson 11, A core group). Take a job of 10 items, ten numbers or ten rows of a table, to be done by 4 cores. The cores need rules that every core can follow alone, without asking the others, and that together give each item to exactly one core. The rule: cut the items into 4 pieces, in order, as equal as they can be. 10 does not divide by 4, so the first two ranks take one item more.

Ten items shared by four ranks: three, three, two, two. Colours show which rank takes which item.

The general form

const km::Partition share = km::partition(ITEMS, RANK, COUNT);
  • ITEMS is how many items the whole job has.
  • RANK is the rank of the core that asks, from 0 to COUNT - 1.
  • COUNT is the number of cores in the group. In a kernel it is km::count(command).
  • share.start is the number of the first item of this rank's piece, counted from 0.
  • share.len is how many items the piece holds. The pieces of the ranks follow one another with no gap and no overlap, and the first ITEMS % COUNT ranks have one item more than the others.

A kernel then does its work for the items from share.start up to, not including, share.start + share.len.

One line at a time

writer.run([&](const km_command_t& command) {
    if (km::handler(command) == kHandlerShare) {
        const uint32_t items = command.a0;
        const uint32_t rank = command.a1;
        const km::Partition share = km::partition(items, rank, km::count(command));
        const uint32_t first_item = share.start;
        const uint32_t item_count = share.len;
        writer.result(first_item << 16 | item_count);
    }
    return true;
});

writer.cpp reads the number of items and a rank from the arguments. Normally the rank is the one the core has itself (writer.me() - km::first(command), from the last lesson), but a Python program sees only the answer of the group's first core, which has rank 0 (lesson 13, The leader explains why). To look at every rank's piece from Python, this program sends the rank as an argument, and the first core works out the piece of that rank with the same function. The result word has two parts packed into one 32-bit number: first_item << 16 moves the first item to the upper 16 bits, and | item_count puts the count in the lower 16.

for rank in range(CORES):
    handle = run.submit("share", group=(0, CORES), args=[ITEMS, rank], result=True)
    word = run.wait(handle, timeout=10.0).result
    first_item, item_count = word >> 16, word & 0xFFFF
    covered += range(first_item, first_item + item_count)
    print(f"rank {rank} takes {item_count} items, from item {first_item}")

Four commands to the group of four cores, one for each rank. Python takes the word apart again: word >> 16 is the first item, word & 0xFFFF the count. The loop prints the piece of each rank and collects the items in covered.

The whole program

"""Partitions: ten items shared by four cores, one contiguous piece each."""

from pathlib import Path

import komorebi as km

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

program = km.Program(
    "partitions",
    kernels=HERE / "kernels",
    handlers=HERE / "kernels" / "handlers.h",
)

ITEMS = 10
CORES = 4

with km.Session() as session:
    run = session.load(program)
    covered = []
    for rank in range(CORES):
        handle = run.submit("share", group=(0, CORES), args=[ITEMS, rank], result=True)
        word = run.wait(handle, timeout=10.0).result
        first_item, item_count = word >> 16, word & 0xFFFF
        covered += range(first_item, first_item + item_count)
        print(f"rank {rank} takes {item_count} items, from item {first_item}")
    run.stop()

assert covered == list(range(ITEMS))
print("together they cover items 0 to 9, each exactly once")

// The handlers of this program: the numbers that the kernels and Python agree on.
#pragma once

#include <stdint.h>

constexpr uint32_t kHandlerShare = 1;

// The writer kernel: works out which items the core of a given rank takes, and sends
// the first item and the number of items back in one word.
#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) == kHandlerShare) {
            const uint32_t items = command.a0;
            const uint32_t rank = command.a1;
            const km::Partition share = km::partition(items, rank, km::count(command));
            const uint32_t first_item = share.start;
            const uint32_t item_count = share.len;
            writer.result(first_item << 16 | item_count);
        }
        return true;
    });
}

rank 0 takes 3 items, from item 0
rank 1 takes 3 items, from item 3
rank 2 takes 2 items, from item 6
rank 3 takes 2 items, from item 8
together they cover items 0 to 9, each exactly once

The first two ranks take three items, the other two take two, and the last line is the script's own check: the pieces, put one after the other, are exactly the items 0 to 9.

Common mistakes

  • Using ITEMS / COUNT for the piece. With 10 items and 4 cores that is 2 for each, and two items are left over. km::partition places the remainder.
  • More cores than items. With 3 items and 4 cores the last rank gets share.len == 0, and a kernel that does not check for it works on an item that does not exist.
  • Using a core's number where the rank belongs. A group that starts at core 40 has no rank 40; km::partition(ITEMS, 40, COUNT) is a piece that does not exist.

Next: lesson 13, The leader is about the one core that speaks for the whole group.