code/tt_meteor/tt/kernels/pack_rows.cpp

// SPDX-License-Identifier: Apache-2.0
// PACK (meteor-p150 OPT round 4 R4.2, tt/pack.py; OPT_PLAN #8 MK7 output tail): the frame's fp32 outputs into the
// packed readback rows in ONE data-movement program on one RISC / NOC (reads and writes on the same NOC: no two-NOC
// small DRAM writes, SYS-1419). Replaces the stock per-output untilize / pad / RM reshape / concat / slice chain.
//
// Entry e (an output, logical [R, C] rows of fp32) lands as its R * C elements in row-major order at element offset
// O_e of the packed stream (O_e a multiple of 1024: whole 4 KB packed rows); the packed stream is cut into NS segment
// tensors of SEG_PAGES rows of 1024 fp32 (page p -> segment p / SEG_PAGES). The tail of the last row of an entry is
// padding (never read by the host).
// Work unit = 32 rows of one entry (u = U_e + g: rows [32 g, min(R, 32 g + 32))):
//   KIND 0 (TILE fp32, interleaved, padded width Wt tiles): the Wt tiles of tile row g are read, the 32 x C valid
//          elements reordered from the tile faces (face f = (row / 16) * 2 + (col % 32) / 16, row-major 16 x 16)
//          into the staging buffer by the RISC;
//   KIND 1 (ROW_MAJOR fp32, C * 4 % 64 == 0): the 32 row pages are read straight into the staging buffer.
// Then the 32 * C * 4 B (a multiple of 128 B; the last unit rounded up to 64 B) go out in pieces split at the 4 KB
// packed-row boundaries.
// CM entries (channel-major, OPT round 4 R4.3: the host's NHWC -> NCHW transpose done here; TILE only, R % 32 == 0,
// HW % 16 == 0): row r = n * HW + hw, column c lands at element (n * C + c) * HW + hw, so the unit is reordered into
// [C][32] and each channel's 32 values go out as one 128 B piece (two 64 B pieces when the unit crosses a plane). Reads of unit k + 1 are issued before unit k is reordered (two input / staging buffers).
// Work split (no per-core RT args, probe P3): core i = y * GX + x owns units [BND[i], BND[i + 1]) (CT table, balanced
// by the bytes moved on the host).
// CT args: [0] NC, [1] GX, [2] NE, [3] SEG_PAGES, [4] NS, [5] WMAX (tiles per unit buffer), [6] STG (staging bytes),
//          [7 .. 7 + NC] BND, then per entry 8 words (KIND, R, C, Wt, U_e, O_e, CM, HW), then the TensorAccessorArgs of the
//          TILE inputs, of the ROW_MAJOR inputs (one set per kind: every input is DRAM interleaved) and of the
//          segments. Common RT args: NE input addresses,
//          then NS segment addresses.
// CB 0: 2 pages of WMAX x 4 KB (input buffers); CB 1: 2 pages of STG bytes (staging buffers); both reserved once and
// used as plain L1 scratch, never pushed. (The first bring-up reserved 2 x WMAX pages of the 2-page CB 0 and hung in
// cb_reserve_back: OPT_REPORT "Known hangs" R4.2.)
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t NC = get_compile_time_arg_val(0);
constexpr uint32_t GX = get_compile_time_arg_val(1);
constexpr uint32_t NE = get_compile_time_arg_val(2);
constexpr uint32_t SEG_PAGES = get_compile_time_arg_val(3);
constexpr uint32_t NS = get_compile_time_arg_val(4);
constexpr uint32_t WMAX = get_compile_time_arg_val(5);
constexpr uint32_t STG = get_compile_time_arg_val(6);
constexpr uint32_t BND0 = 7;
constexpr uint32_t META = BND0 + NC + 1;
constexpr uint32_t MW = 8;        // meta words per entry
constexpr auto tile_args = TensorAccessorArgs<META + MW * NE>();
constexpr auto rm_args = TensorAccessorArgs<tile_args.next_compile_time_args_offset()>();
constexpr auto out_args = TensorAccessorArgs<rm_args.next_compile_time_args_offset()>();
constexpr uint32_t TB = 4096;     // fp32 tile bytes == packed row bytes
constexpr uint32_t ROWB = 4096;

struct Unit {
    uint32_t e, g, kind, rows, c, wt, o, cm, hw;
};

FORCE_INLINE Unit unit_of(uint32_t u) {
    uint32_t e = 0;
    while (e + 1 < NE && kernel_compile_time_args[META + MW * (e + 1) + 4] <= u) {
        ++e;
    }
    const uint32_t m = META + MW * e;
    Unit r;
    r.e = e;
    r.kind = kernel_compile_time_args[m];
    const uint32_t R = kernel_compile_time_args[m + 1];
    r.c = kernel_compile_time_args[m + 2];
    r.wt = kernel_compile_time_args[m + 3];
    r.g = u - kernel_compile_time_args[m + 4];
    r.o = kernel_compile_time_args[m + 5];
    r.cm = kernel_compile_time_args[m + 6];
    r.hw = kernel_compile_time_args[m + 7];
    const uint32_t r0 = r.g * 32;
    r.rows = (R - r0 < 32) ? R - r0 : 32;
    return r;
}

FORCE_INLINE void issue_reads(const Unit& un, uint32_t buf) {
    const uint32_t addr = get_common_arg_val<uint32_t>(un.e);
    if (un.kind == 0) {
        const auto src = TensorAccessor(tile_args, addr, TB);
        for (uint32_t k = 0; k < un.wt; ++k) {
            noc_async_read(src.get_noc_addr(un.g * un.wt + k), buf + k * TB, TB);
        }
    } else {
        const uint32_t pb = un.c * 4;
        const auto src = TensorAccessor(rm_args, addr, pb);
        for (uint32_t i = 0; i < un.rows; ++i) {
            noc_async_read(src.get_noc_addr(un.g * 32 + i), buf + i * pb, pb);
        }
    }
}

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    const uint32_t u0 = kernel_compile_time_args[BND0 + core];
    const uint32_t u1 = kernel_compile_time_args[BND0 + core + 1];
    if (u0 >= u1) {
        return;
    }
    cb_reserve_back(0, 2);                              // CB 0: 2 pages of WMAX x 4 KB (the two input buffers)
    const uint32_t inb = get_write_ptr(0);
    cb_reserve_back(1, 2);                              // CB 1: 2 pages of STG bytes (the two staging buffers)
    const uint32_t stb = get_write_ptr(1);

    Unit cur = unit_of(u0);
    issue_reads(cur, inb);
    uint32_t slot = 0;
    for (uint32_t u = u0; u < u1; ++u) {
        noc_async_read_barrier();                       // unit u is in buffer `slot`
        // the other input buffer and staging buffer `slot` were write sources of earlier units: their writes must
        // have left L1 before they are refilled
        noc_async_writes_flushed();
        Unit nxt = cur;
        if (u + 1 < u1) {
            nxt = unit_of(u + 1);
            issue_reads(nxt, inb + (slot ^ 1) * WMAX * TB);
        }
        const uint32_t ib = inb + slot * WMAX * TB;
        const uint32_t sb = stb + slot * STG;
        const uint32_t n = cur.rows * cur.c;            // elements of this unit
        uint32_t src = ib;
        if (cur.cm) {
            // [C][32] from the tile faces: column c, rows i: tile c / 32, face (i / 16) * 2 + (c % 32) / 16, word
            // (i % 16) * 16 + c % 16
            volatile tt_l1_ptr uint32_t* dst = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(sb);
            for (uint32_t c = 0; c < cur.c; ++c) {
                const uint32_t base = ib + (c >> 5) * TB + ((((c >> 4) & 1) * 256) + (c & 15)) * 4;
                volatile tt_l1_ptr uint32_t* s0 = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(base);
                volatile tt_l1_ptr uint32_t* s1 = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(base + 512 * 4);
                for (uint32_t i = 0; i < 16; ++i) {
                    dst[i] = s0[i * 16];
                    dst[16 + i] = s1[i * 16];
                }
                dst += 32;
            }
            const uint32_t r0 = cur.g * 32;
            const uint32_t n0 = r0 / cur.hw;
            const uint32_t hw0 = r0 - n0 * cur.hw;
            const uint32_t k = (cur.hw - hw0 < 32) ? cur.hw - hw0 : 32;   // rows in plane n0 (16 or 32)
            for (uint32_t c = 0; c < cur.c; ++c) {
                const uint32_t sc = sb + c * 128;
                for (uint32_t part = 0; part < 2; ++part) {
                    const uint32_t rows = part == 0 ? k : 32 - k;
                    if (rows == 0) {
                        continue;
                    }
                    uint32_t src2 = sc + (part == 0 ? 0 : k * 4);
                    uint32_t off = (cur.o + ((n0 + part) * cur.c + c) * cur.hw + (part == 0 ? hw0 : 0)) * 4;
                    uint32_t bytes = rows * 4;
                    while (bytes) {
                        const uint32_t page = off / ROWB;
                        const uint32_t inpg = off - page * ROWB;
                        const uint32_t len = (ROWB - inpg < bytes) ? ROWB - inpg : bytes;
                        const uint32_t sg = page / SEG_PAGES;
                        const auto dst2 = TensorAccessor(out_args, get_common_arg_val<uint32_t>(NE + sg), ROWB);
                        noc_async_write(src2, dst2.get_noc_addr(page - sg * SEG_PAGES, inpg), len);
                        src2 += len;
                        off += len;
                        bytes -= len;
                    }
                }
            }
            cur = nxt;
            slot ^= 1;
            continue;
        }
        if (cur.kind == 0) {
            volatile tt_l1_ptr uint32_t* dst = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(sb);
            for (uint32_t i = 0; i < cur.rows; ++i) {
                const uint32_t fr = (i >> 4) * 2;           // face row
                const uint32_t ro = (i & 15) * 16;          // row offset inside the face
                for (uint32_t c0 = 0; c0 < cur.c; c0 += 16) {
                    const uint32_t t = c0 >> 5;
                    const uint32_t f = fr + ((c0 >> 4) & 1);
                    volatile tt_l1_ptr uint32_t* s =
                        reinterpret_cast<volatile tt_l1_ptr uint32_t*>(ib + t * TB + (f * 256 + ro) * 4);
                    const uint32_t m = (cur.c - c0 < 16) ? cur.c - c0 : 16;
                    for (uint32_t j = 0; j < m; ++j) {
                        dst[j] = s[j];
                    }
                    dst += m;
                }
            }
            src = sb;
        }
        // out: n * 4 bytes (rounded up to 64 B) at byte offset (O_e + g * 32 * C) * 4, split at 4 KB rows
        uint32_t bytes = (n * 4 + 63) & ~63u;
        uint32_t off = (cur.o + cur.g * 32 * cur.c) * 4;
        while (bytes) {
            const uint32_t page = off / ROWB;
            const uint32_t inpg = off - page * ROWB;
            const uint32_t len = (ROWB - inpg < bytes) ? ROWB - inpg : bytes;
            const uint32_t s = page / SEG_PAGES;
            const auto dst = TensorAccessor(out_args, get_common_arg_val<uint32_t>(NE + s), ROWB);
            noc_async_write(src, dst.get_noc_addr(page - s * SEG_PAGES, inpg), len);
            src += len;
            off += len;
            bytes -= len;
        }
        cur = nxt;
        slot ^= 1;
    }
    noc_async_write_barrier();
}