code/tt_meteor/tt/kernels/terms_epi_reader.cpp

// SPDX-License-Identifier: Apache-2.0
// TERMS_EPI reader (meteor-p150 OPT_PLAN #2, tt/epi.py): the NP fp32 term outputs of one TermsConv (interleaved TILE,
// one shape) into CBs 0 .. NP-1 and, with HAS_RES, the residual map (same shape) into CB 4, NB tiles per batch.
// Work split (no per-core RT args, probe P3): core i = y * GX + x of the first NC cores in row-major order owns the
// contiguous tile range [i * B + min(i, E), + B + (i < E)) of the T output tiles, B = T / NC, E = T % NC.
// A part may be a column range of a wider tensor (ST2, METEOR_TERMS_ST2: the stacked [W_hi | W_lo] conv's [R, 2C]
// output read as its two products): output tile t = (r, c) = (t / CT, t % CT) reads tile r * W_i + O_i + c of part i
// (W_i = CT, O_i = 0 for a plain part: tile t). The residual has the output's shape (tile t).
// CT args: [0] T, [1] NC, [2] GX, [3] NP (1..3), [4] HAS_RES, [5] NB, [6] CT (output column tiles), [7 + 2i] /
// [8 + 2i] W_i / O_i (column tiles / column-tile offset of part i), then the TensorAccessorArgs of part 0, 1, 2
// (repeats of part 0 when NP < 3) and the residual (part 0 when absent). Common RT args: [p0, p1, p2, res].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t T = get_compile_time_arg_val(0);
constexpr uint32_t NC = get_compile_time_arg_val(1);
constexpr uint32_t GX = get_compile_time_arg_val(2);
constexpr uint32_t NP = get_compile_time_arg_val(3);
constexpr uint32_t HAS_RES = get_compile_time_arg_val(4);
constexpr uint32_t NB = get_compile_time_arg_val(5);
constexpr uint32_t CT = get_compile_time_arg_val(6);
constexpr uint32_t W0 = get_compile_time_arg_val(7), O0 = get_compile_time_arg_val(8);
constexpr uint32_t W1 = get_compile_time_arg_val(9), O1 = get_compile_time_arg_val(10);
constexpr uint32_t W2 = get_compile_time_arg_val(11), O2 = get_compile_time_arg_val(12);
constexpr bool PLAIN = (W0 == CT && O0 == 0 && W1 == CT && O1 == 0 && W2 == CT && O2 == 0);
constexpr auto p0_args = TensorAccessorArgs<13>();
constexpr auto p1_args = TensorAccessorArgs<p0_args.next_compile_time_args_offset()>();
constexpr auto p2_args = TensorAccessorArgs<p1_args.next_compile_time_args_offset()>();
constexpr auto rs_args = TensorAccessorArgs<p2_args.next_compile_time_args_offset()>();
constexpr uint32_t TB = 4096;
constexpr uint32_t cb_res = 4;

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = T / NC, E = T % NC;
    const uint32_t t0 = core * B + (core < E ? core : E);
    const uint32_t n = B + (core < E ? 1 : 0);
    if (n == 0) {
        return;
    }
    const auto a0 = TensorAccessor(p0_args, get_common_arg_val<uint32_t>(0), TB);
    const auto a1 = TensorAccessor(p1_args, get_common_arg_val<uint32_t>(1), TB);
    const auto a2 = TensorAccessor(p2_args, get_common_arg_val<uint32_t>(2), TB);
    const auto ar = TensorAccessor(rs_args, get_common_arg_val<uint32_t>(3), TB);
    // (r, c) of the next tile for the column-range parts
    uint32_t r = t0 / CT, c = t0 % CT;
    for (uint32_t t = t0; t < t0 + n; t += NB) {
        const uint32_t nb = (t + NB <= t0 + n) ? NB : t0 + n - t;
        uint32_t w0 = 0, w1 = 0, w2 = 0, wr = 0;
        cb_reserve_back(0, nb);
        w0 = get_write_ptr(0);
        if constexpr (NP > 1) {
            cb_reserve_back(1, nb);
            w1 = get_write_ptr(1);
        }
        if constexpr (NP > 2) {
            cb_reserve_back(2, nb);
            w2 = get_write_ptr(2);
        }
        if constexpr (HAS_RES) {
            cb_reserve_back(cb_res, nb);
            wr = get_write_ptr(cb_res);
        }
        for (uint32_t j = 0; j < nb; ++j) {
            uint32_t i0 = t + j, i1 = t + j, i2 = t + j;
            if constexpr (!PLAIN) {
                i0 = r * W0 + O0 + c;
                i1 = r * W1 + O1 + c;
                i2 = r * W2 + O2 + c;
                if (++c == CT) {
                    c = 0;
                    ++r;
                }
            }
            noc_async_read(a0.get_noc_addr(i0), w0, TB);
            w0 += TB;
            if constexpr (NP > 1) {
                noc_async_read(a1.get_noc_addr(i1), w1, TB);
                w1 += TB;
            }
            if constexpr (NP > 2) {
                noc_async_read(a2.get_noc_addr(i2), w2, TB);
                w2 += TB;
            }
            if constexpr (HAS_RES) {
                noc_async_read(ar.get_noc_addr(t + j), wr, TB);
                wr += TB;
            }
        }
        noc_async_read_barrier();
        cb_push_back(0, nb);
        if constexpr (NP > 1) {
            cb_push_back(1, nb);
        }
        if constexpr (NP > 2) {
            cb_push_back(2, nb);
        }
        if constexpr (HAS_RES) {
            cb_push_back(cb_res, nb);
        }
    }
}