code/tt_diffusion_planner/ttaw/ops/kernels/segment_reduce_dm.cpp

// SPDX-License-Identifier: Apache-2.0
// K1 segment_reduce (ttaw.ops.segment, PLAN.md 1.3): per-segment MAX over the rows of a tensor sorted by segment.
//
//   out[s, :] = max(x[off[s] : off[s + 1], :])      for s in [0, NUM_SEG); empty segments -> FILL_BITS
//
// x: [1, 1, NUM_ROWS, C] interleaved DRAM, TILE (IN_TILE = 1) or ROW_MAJOR, bf16 or fp32; off: uint32 ROW_MAJOR
// [1, >= NUM_SEG + 1] interleaved DRAM (one page; CSR offsets, non-decreasing); out: [1, 1, NUM_SEG, C] ROW_MAJOR
// interleaved DRAM, bf16 (round to nearest even) or fp32. Exact: the maximum is taken on integer keys of the
// float bit patterns (key = bits ^ ((bits >> 31) & 0x7fffffff), signed order == float order, -0 < +0), so no float
// arithmetic happens on the RISC-V (no F extension) and bf16 / fp32 inputs come out bit-exact (fp32 -> bf16 is a
// monotonic rounding, so max-then-round == round-then-max). NaN is not supported (not produced by the models).
//
// Work split (deterministic, data-dependent, no per-core runtime args): NUM_WORKERS workers = SLOTS data-movement
// RISC-Vs per core on a GW-wide core rectangle starting at (X0, Y0); worker w = core_index * SLOTS + SLOT. The cost
// of segment s is off[s] + s (its rows plus one output row; strictly increasing in s), T = off[S] + S. Worker w owns
// the segments whose cost lies in [w T / W, (w + 1) T / W): it finds them by two binary searches over off (64-byte
// probes from DRAM), then streams its contiguous row range in units of 32 rows (one tile-row in TILE mode) with
// double buffering, reduces each segment into an L1 accumulator and writes one output row per segment (a ring of
// OUT_RING staging rows). Every segment is written by exactly one worker; nothing is shared between workers.
//
// Runtime args: COMMON only [x_addr, off_addr, out_addr] (P3 rule: per-core work derives from the core coordinates;
// the arg layout is fixed by the compile-time args). Hang-protocol notes (PLAN.md 4.4): no CB is pushed or popped
// (CBs are plain L1 scratch: data 2 units, offsets window, probe block, accumulator, output ring; sizes checked by
// the host builder), no multicast, no semaphores; every DRAM read has matching 64-byte alignment of source and
// destination (Blackhole NOC_DRAM_READ_ALIGNMENT_BYTES), every read is waited on by noc_async_read_barrier (which
// also invalidates the Blackhole L1 data cache) before the data are used, staging rows are reused only after
// noc_async_writes_flushed, and the kernel ends with noc_async_write_barrier. Offsets are clamped to [0, NUM_ROWS]
// and to non-decreasing order, so a corrupt table can never make the kernel read outside x.
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t IN_TILE = get_compile_time_arg_val(0);      // 1: TILE input, 0: ROW_MAJOR input
constexpr uint32_t IN_BYTES = get_compile_time_arg_val(1);     // input element bytes: 2 (bf16) or 4 (fp32)
constexpr uint32_t OUT_BYTES = get_compile_time_arg_val(2);    // output element bytes: 2 (bf16, RNE) or 4 (fp32)
constexpr uint32_t C = get_compile_time_arg_val(3);            // channels (multiple of 32)
constexpr uint32_t NUM_ROWS = get_compile_time_arg_val(4);     // input rows (multiple of 32): the read bound
constexpr uint32_t NUM_SEG = get_compile_time_arg_val(5);      // output rows (segments)
constexpr uint32_t NUM_WORKERS = get_compile_time_arg_val(6);  // workers of the program
constexpr uint32_t SLOTS = get_compile_time_arg_val(7);        // workers per core (1 or 2)
constexpr uint32_t SLOT = get_compile_time_arg_val(8);         // this kernel's slot on its core
constexpr uint32_t X0 = get_compile_time_arg_val(9);           // core rectangle: first logical x
constexpr uint32_t Y0 = get_compile_time_arg_val(10);          //                 first logical y
constexpr uint32_t GW = get_compile_time_arg_val(11);          //                 width in cores
constexpr uint32_t CB_DATA = get_compile_time_arg_val(12);     // 2 units of input rows (+ 64 B alignment slack)
constexpr uint32_t CB_OFF = get_compile_time_arg_val(13);      // offsets window: OFF_CHUNK uint32 (+ 64 B)
constexpr uint32_t CB_PROBE = get_compile_time_arg_val(14);    // one 64-byte probe block (+ 64 B)
constexpr uint32_t CB_ACC = get_compile_time_arg_val(15);      // C int32 keys
constexpr uint32_t CB_OUT = get_compile_time_arg_val(16);      // OUT_RING output rows (+ 64 B)
constexpr uint32_t OUT_RING = get_compile_time_arg_val(17);    // staging rows in flight
constexpr uint32_t OFF_CHUNK = get_compile_time_arg_val(18);   // offsets per window (multiple of 16)
constexpr uint32_t FILL_BITS = get_compile_time_arg_val(19);   // fp32 bits written for an empty segment
constexpr uint32_t OFF_PAGE_BYTES = get_compile_time_arg_val(20);  // page (= tensor) size of off, bytes
constexpr auto x_args = TensorAccessorArgs<21>();
constexpr auto off_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
constexpr auto out_args = TensorAccessorArgs<off_args.next_compile_time_args_offset()>();

constexpr uint32_t UNIT_ROWS = 32;                              // rows per streamed unit (one tile-row)
constexpr uint32_t CT = C / 32;                                 // 32-column tiles per row
constexpr uint32_t TILE_ELEMS = 1024;
constexpr uint32_t IN_PAGE_BYTES = IN_TILE ? TILE_ELEMS * IN_BYTES : C * IN_BYTES;
constexpr uint32_t PAGES_PER_UNIT = IN_TILE ? CT : UNIT_ROWS;
constexpr uint32_t UNIT_BYTES = PAGES_PER_UNIT * IN_PAGE_BYTES;
constexpr uint32_t OUT_ROW_BYTES = C * OUT_BYTES;
constexpr uint32_t OFF_ALIGNED_BYTES = (OFF_PAGE_BYTES + 63u) & ~63u;
constexpr uint32_t NUM_UNITS = NUM_ROWS / UNIT_ROWS;

static_assert(C % 32 == 0, "C must be a multiple of 32");
static_assert(NUM_ROWS % UNIT_ROWS == 0, "NUM_ROWS must be a multiple of 32");
static_assert(IN_BYTES == 2 || IN_BYTES == 4, "bf16 or fp32 input");
static_assert(OUT_BYTES == 2 || OUT_BYTES == 4, "bf16 or fp32 output");
static_assert(OFF_CHUNK % 16 == 0 && OFF_CHUNK >= 16, "OFF_CHUNK must be a multiple of 16");
static_assert(OUT_RING >= 1, "OUT_RING >= 1");
static_assert(IN_PAGE_BYTES % 64 == 0, "input pages must be 64-byte multiples (DRAM read alignment)");
static_assert(OUT_ROW_BYTES % 16 == 0, "output rows must be 16-byte multiples (DRAM write alignment)");

#define COMPILER_BARRIER() asm volatile("" ::: "memory")

static inline uint32_t align64(uint32_t a) { return (a + 63u) & ~63u; }

// float bits <-> signed key (an involution); signed comparison of keys == float comparison (no NaN)
static inline int32_t to_key(uint32_t bits) {
    const int32_t b = static_cast<int32_t>(bits);
    return b ^ ((b >> 31) & 0x7fffffff);
}
static inline uint32_t from_key(int32_t k) { return static_cast<uint32_t>(k ^ ((k >> 31) & 0x7fffffff)); }

static inline uint16_t to_bf16_rne(uint32_t bits) {
    return static_cast<uint16_t>((bits + 0x7fffu + ((bits >> 16) & 1u)) >> 16);
}

void kernel_main() {
    const uint32_t x_addr = get_common_arg_val<uint32_t>(0);
    const uint32_t off_addr = get_common_arg_val<uint32_t>(1);
    const uint32_t out_addr = get_common_arg_val<uint32_t>(2);

    const uint32_t core = (get_absolute_logical_y() - Y0) * GW + (get_absolute_logical_x() - X0);
    const uint32_t w = core * SLOTS + SLOT;
    if (w >= NUM_WORKERS) {
        return;
    }

    const auto xa = TensorAccessor(x_args, x_addr, IN_PAGE_BYTES);
    const auto oa = TensorAccessor(off_args, off_addr, OFF_PAGE_BYTES);
    const auto ya = TensorAccessor(out_args, out_addr, OUT_ROW_BYTES);

    const uint32_t data_l1 = align64(get_write_ptr(CB_DATA));
    const uint32_t off_l1 = align64(get_write_ptr(CB_OFF));
    const uint32_t probe_l1 = align64(get_write_ptr(CB_PROBE));
    const uint32_t out_l1 = align64(get_write_ptr(CB_OUT));
    int32_t* acc = reinterpret_cast<int32_t*>(get_write_ptr(CB_ACC));

    // ---- offsets: single 64-byte probes (binary search) and a streamed window (the worker's segments)
    auto probe = [&](uint32_t s) -> uint32_t {
        const uint32_t byte = s * 4u;
        noc_async_read(oa.get_noc_addr(0, byte & ~63u), probe_l1, 64);
        noc_async_read_barrier();
        COMPILER_BARRIER();
        const uint32_t v = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(probe_l1)[(byte & 63u) >> 2];
        return v < NUM_ROWS ? v : NUM_ROWS;
    };
    uint32_t win0 = 0xffffffffu;  // first offset index held in the window (multiple of 16)
    auto off_at = [&](uint32_t s) -> uint32_t {
        if (win0 == 0xffffffffu || s < win0 || s >= win0 + OFF_CHUNK) {
            win0 = s & ~15u;
            const uint32_t start = win0 * 4u;
            uint32_t bytes = OFF_CHUNK * 4u;
            if (start + bytes > OFF_ALIGNED_BYTES) {
                bytes = OFF_ALIGNED_BYTES - start;
            }
            noc_async_read(oa.get_noc_addr(0, start), off_l1, bytes);
            noc_async_read_barrier();
            COMPILER_BARRIER();
        }
        const uint32_t v = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(off_l1)[s - win0];
        return v < NUM_ROWS ? v : NUM_ROWS;
    };

    // ---- this worker's segments [s_lo, s_hi): costs off[s] + s in [w T / W, (w + 1) T / W)
    const uint32_t total = probe(NUM_SEG) + NUM_SEG;
    auto search = [&](uint32_t target) -> uint32_t {  // smallest s in [0, NUM_SEG] with off[s] + s >= target
        uint32_t lo = 0, hi = NUM_SEG;
        while (lo < hi) {
            const uint32_t mid = (lo + hi) >> 1;
            if (probe(mid) + mid >= target) {
                hi = mid;
            } else {
                lo = mid + 1;
            }
        }
        return lo;
    };
    const uint32_t t0 = static_cast<uint32_t>((static_cast<uint64_t>(w) * total) / NUM_WORKERS);
    const uint32_t t1 = static_cast<uint32_t>((static_cast<uint64_t>(w + 1) * total) / NUM_WORKERS);
    const uint32_t s_lo = (w == 0) ? 0 : search(t0);
    const uint32_t s_hi = (w + 1 == NUM_WORKERS) ? NUM_SEG : search(t1);
    if (s_lo >= s_hi) {
        return;
    }

    // ---- input rows: units of 32 rows, two L1 slots, the next unit prefetched while the current one is reduced
    uint32_t row_last = off_at(s_lo);  // largest row this worker reads (+1), monotone over its segments
    {
        const uint32_t r_end = off_at(s_hi);
        if (r_end > row_last) {
            row_last = r_end;
        }
    }
    const uint32_t unit_last = row_last > 0 ? (row_last - 1) / UNIT_ROWS : 0;
    auto issue_unit = [&](uint32_t u, uint32_t slot) {
        const uint32_t dst = data_l1 + slot * UNIT_BYTES;
        const uint32_t first_page = IN_TILE ? u * CT : u * UNIT_ROWS;
        for (uint32_t p = 0; p < PAGES_PER_UNIT; ++p) {
            noc_async_read(xa.get_noc_addr(first_page + p), dst + p * IN_PAGE_BYTES, IN_PAGE_BYTES);
        }
    };
    uint32_t cur_unit = 0xffffffffu, cur_slot = 0;
    bool next_ready = false;  // the unit cur_unit + 1 is in flight into slot cur_slot ^ 1
    auto ensure_unit = [&](uint32_t u) {
        if (u == cur_unit) {
            return;
        }
        if (next_ready && u == cur_unit + 1) {
            noc_async_read_barrier();
            cur_slot ^= 1u;
        } else {
            noc_async_read_barrier();  // drain a stale prefetch before its slot is reused
            issue_unit(u, cur_slot);
            noc_async_read_barrier();
        }
        COMPILER_BARRIER();
        cur_unit = u;
        next_ready = false;
        if (u + 1 <= unit_last && u + 1 < NUM_UNITS) {
            issue_unit(u + 1, cur_slot ^ 1u);
            next_ready = true;
        }
    };

    auto reset_acc = [&]() {
        for (uint32_t c = 0; c < C; ++c) {
            acc[c] = static_cast<int32_t>(0x80000000u);
        }
    };
    auto accumulate_row = [&](uint32_t r) {  // acc = max(acc, x[r, :]) on keys
        const uint32_t i = r - cur_unit * UNIT_ROWS;
        const uint32_t base = data_l1 + cur_slot * UNIT_BYTES;
        if constexpr (IN_TILE) {
            const uint32_t face_row = (i >> 4) << 1;  // faces 0/1 (rows 0-15) or 2/3 (rows 16-31)
            for (uint32_t tc = 0; tc < CT; ++tc) {
                for (uint32_t h = 0; h < 2; ++h) {
                    const uint32_t e0 = tc * TILE_ELEMS + (face_row + h) * 256 + (i & 15u) * 16;
                    int32_t* a = acc + tc * 32 + h * 16;
                    if constexpr (IN_BYTES == 4) {
                        volatile tt_l1_ptr uint32_t* src = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(base) + e0;
                        for (uint32_t k = 0; k < 16; ++k) {
                            const int32_t key = to_key(src[k]);
                            a[k] = key > a[k] ? key : a[k];
                        }
                    } else {
                        volatile tt_l1_ptr uint16_t* src = reinterpret_cast<volatile tt_l1_ptr uint16_t*>(base) + e0;
                        for (uint32_t k = 0; k < 16; ++k) {
                            const int32_t key = to_key(static_cast<uint32_t>(src[k]) << 16);
                            a[k] = key > a[k] ? key : a[k];
                        }
                    }
                }
            }
        } else {
            if constexpr (IN_BYTES == 4) {
                volatile tt_l1_ptr uint32_t* src = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(base) + i * C;
                for (uint32_t c = 0; c < C; ++c) {
                    const int32_t key = to_key(src[c]);
                    acc[c] = key > acc[c] ? key : acc[c];
                }
            } else {
                volatile tt_l1_ptr uint16_t* src = reinterpret_cast<volatile tt_l1_ptr uint16_t*>(base) + i * C;
                for (uint32_t c = 0; c < C; ++c) {
                    const int32_t key = to_key(static_cast<uint32_t>(src[c]) << 16);
                    acc[c] = key > acc[c] ? key : acc[c];
                }
            }
        }
    };

    // ---- output rows: a ring of OUT_RING staging rows, reused after noc_async_writes_flushed
    uint32_t emitted = 0;
    auto emit = [&](uint32_t s, bool empty) {
        const uint32_t ring = emitted % OUT_RING;
        if (ring == 0 && emitted > 0) {
            noc_async_writes_flushed();
        }
        const uint32_t dst = out_l1 + ring * OUT_ROW_BYTES;
        if constexpr (OUT_BYTES == 4) {
            volatile tt_l1_ptr uint32_t* o = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(dst);
            for (uint32_t c = 0; c < C; ++c) {
                o[c] = empty ? FILL_BITS : from_key(acc[c]);
            }
        } else {
            volatile tt_l1_ptr uint16_t* o = reinterpret_cast<volatile tt_l1_ptr uint16_t*>(dst);
            const uint16_t fill = to_bf16_rne(FILL_BITS);
            for (uint32_t c = 0; c < C; ++c) {
                o[c] = empty ? fill : to_bf16_rne(from_key(acc[c]));
            }
        }
        COMPILER_BARRIER();
        noc_async_write(dst, ya.get_noc_addr(s), OUT_ROW_BYTES);
        ++emitted;
    };

    // ---- the reduction
    uint32_t seg_begin = off_at(s_lo);
    for (uint32_t s = s_lo; s < s_hi; ++s) {
        uint32_t seg_end = off_at(s + 1);
        if (seg_end < seg_begin) {
            seg_end = seg_begin;  // non-decreasing offsets (a corrupt table reads nothing extra)
        }
        if (seg_end == seg_begin) {
            emit(s, true);
        } else {
            reset_acc();
            for (uint32_t r = seg_begin; r < seg_end; ++r) {
                ensure_unit(r / UNIT_ROWS);
                accumulate_row(r);
            }
            emit(s, false);
        }
        seg_begin = seg_end;
    }
    noc_async_read_barrier();  // a prefetch past the last row may still be in flight
    noc_async_write_barrier();
}