// SPDX-License-Identifier: Apache-2.0
// msda_k (tt/msda_k.py, OPT_PLAN #1 / #2 / #10): the data-movement half of the fixed-point MSDA kernel. The same
// source runs on both data-movement RISC-Vs; PART 0 (NCRISC, "reader") owns the points [0, P/2), PART 1 (BRISC,
// "writer") the points [P/2, P) and also writes the output tiles. Each RISC has its own CBs (cb_geo, cb_in, cb_loc),
// so there is no cross-RISC handshake.
//
// Work item w (contiguous range per core, core index from the logical coordinates: common RT args only, probe P3):
//   j = w / NB (head block), b = w % NB (32-query block); value map m = j * SUB + b / (NB / SUB);
//   loc  int32 TILE [1, H, NB * 32, 32], tile j * NB + b: row r, column 2p / 2p + 1 = x / y of point p as
//        round(pixel * 2^16) (pixel = (g + 1) size / 2 - 0.5, align_corners=False);
//   attn fp32 TILE  [1, H, NB * 32, 32], tile j * NB + b: the softmax weight of point p in column p;
//   out  bf16 TILE  [1, 1, NB * 32, H * 32], tile b * H + j.
// value: VSRC 0 = (N, h, w, 32) ROW_MAJOR, page m * h * w + y * w + x; VSRC 1 = the head split [B, heads, rows, 32]
// TILE as nlp_create_qkv_heads writes it (map m = j * SUB + s: slice j, rows s * h * w ..; TPH tiles per slice), the
// stick of a pixel = row r of its tile: two 32-byte face rows (faces 2 (r / 16) and + 1), no untilize needed.
// Per item and RISC: cb_geo <- 3 fp32 tiles (fx, fy, attn with the other RISC's point columns zeroed) and
// cb_in <- 2P bf16 tiles (step k = (p - p0) * 4 + corner; corners NW, NE, SW, SE): row r of step k is the 32-channel
// value stick of that corner, read as two 32-byte halves straight into the tile's face rows; corners outside the map
// (and padding rows) read the zeros page. The compute kernel turns geo into the 4 corner weights on the SFPU.
#include <stdint.h>

#include "api/dataflow/dataflow_api.h"
#include "ttnn/cpp/ttnn/kernel_lib/reduce_helpers_dataflow.hpp"

constexpr uint32_t PART = get_compile_time_arg_val(1);
constexpr uint32_t cb_in = get_compile_time_arg_val(2);
constexpr uint32_t cb_geo = get_compile_time_arg_val(3);
constexpr uint32_t cb_loc = get_compile_time_arg_val(4);
constexpr uint32_t cb_e = get_compile_time_arg_val(5);
constexpr uint32_t cb_out = get_compile_time_arg_val(6);
constexpr int32_t HV = get_compile_time_arg_val(7);
constexpr int32_t WV = get_compile_time_arg_val(8);
constexpr uint32_t NQ = get_compile_time_arg_val(9);
constexpr uint32_t NB = get_compile_time_arg_val(10);
constexpr uint32_t H = get_compile_time_arg_val(11);
constexpr uint32_t SUB = get_compile_time_arg_val(12);
constexpr uint32_t P = get_compile_time_arg_val(13);
constexpr uint32_t NWORK = get_compile_time_arg_val(14);
constexpr uint32_t NC = get_compile_time_arg_val(15);
constexpr uint32_t GX = get_compile_time_arg_val(16);
constexpr uint32_t WRITES_OUT = get_compile_time_arg_val(17);
constexpr uint32_t MAKES_E = get_compile_time_arg_val(18);
// 0: ROW_MAJOR maps (64-byte pages); 1: TILE heads; 2 (BEVFORMER_MSDA_PAIRS): pixel-pair maps [N, h, w + 1, 64]
// (tt/pairmap_k.py: 128-byte page (m, y, i) = the sticks of x = i - 1 and x = i): ONE read per point and row lands
// the NW + NE (or SW + SE) sticks as a row of a [32 x 64] row-major staging block (tilized by the compute kernel)
constexpr uint32_t VSRC = get_compile_time_arg_val(19);
constexpr uint32_t TPH = get_compile_time_arg_val(20);    // VSRC 1: tiles per head slice (SUB maps of HV x WV rows)
// SKIP: items strided over the cores (w = core, core + NC, ..: padding blocks spread evenly) and points with all four
// corners outside the map (or rows past NQ) issue no reads: their weight is set to 0, and their tile rows keep
// earlier (finite) gathers; the step tiles are zeroed once. 0: contiguous item ranges, zeros-page reads.
constexpr uint32_t SKIP = get_compile_time_arg_val(21);
// LOC_HS 0: loc is the head split [1, H, NB * 32, 32] (tile j * NB + b, columns 2p / 2p + 1); LOC_HS > 0
// (BEVFORMER_MSDA_LOC=compact): loc is the merged [1, 1, NB * 32, H * LOC_HS] map (the sampling projection's own
// layout: head j in columns j * LOC_HS .. + 2P), LOC_TW tiles per row block: tile b * LOC_TW + j * LOC_HS / 32,
// column (j * LOC_HS) % 32 + 2p
constexpr uint32_t LOC_HS = get_compile_time_arg_val(22);
constexpr uint32_t LOC_TW = get_compile_time_arg_val(23);
// SM 1 (BEVFORMER_MSDA_SOFTMAX=kernel): attn is the compact attention-logit map [1, 1, NB * 32, H * P] (head j in
// columns j * P ..; AT_TW tiles per row block); the PART-0 RISC rebuilds the head's logit tile (columns 0..P-1 the
// logits, the rest PAD_LOGIT: what the head-block projection writes) into cb_lg for the compute kernel's softmax,
// and makes the reduce scalers and the two half masks once. Rows past NQ get logits 0 (finite; their values are 0).
constexpr uint32_t SM = get_compile_time_arg_val(24);
constexpr uint32_t AT_TW = get_compile_time_arg_val(25);
constexpr uint32_t PAD_BITS = get_compile_time_arg_val(26);   // fp32 bits of PAD_LOGIT
constexpr uint32_t cb_lg = 9, cb_scmax = 10, cb_scsum = 11, cb_msk = 12, cb_scr = 18;
constexpr auto v_args = TensorAccessorArgs<27>();
constexpr auto l_args = TensorAccessorArgs<v_args.next_compile_time_args_offset()>();
constexpr auto a_args = TensorAccessorArgs<l_args.next_compile_time_args_offset()>();
constexpr auto o_args = TensorAccessorArgs<a_args.next_compile_time_args_offset()>();

constexpr uint32_t PH = P / 2;
constexpr uint32_t P0 = PART * PH;
constexpr uint32_t NS = PH * 4;  // steps (tiles) per item and RISC
constexpr uint32_t TB16 = 2048, TB32 = 4096;
constexpr uint32_t VPAGE = VSRC == 0 ? 64 : (VSRC == 2 ? 128 : 2048);  // RM stick; pixel pair; TILE: a bf16 tile
static_assert(VSRC != 2 || !SKIP, "msda_k: pair maps and SKIP are exclusive");
constexpr uint32_t BPS = NB / SUB;
constexpr uint32_t MAP = static_cast<uint32_t>(HV * WV);
constexpr int32_t LOC_LIM = 1 << 28;  // |pixel| < 4096 after the clamp: no overflow below
static_assert(P % 2 == 0 && P <= 8, "msda_k: P even, <= 8 (x / y of a point share a face)");
static_assert(NB % SUB == 0, "msda_k: blocks split evenly over the value maps of a head");
static_assert(LOC_HS == 0 || (LOC_HS >= 2 * P && 32 % LOC_HS == 0), "msda_k: LOC_HS holds 2P columns, divides 32");

static inline uint32_t tile_elem(uint32_t r, uint32_t c) {
    return (((r >> 4) << 1) + (c >> 4)) * 256 + ((r & 15) << 4) + (c & 15);
}

// q in [0, 65536): the fp32 bit pattern of q / 65536 (exact)
static inline uint32_t q16_to_f32(uint32_t q) {
    if (q == 0) {
        return 0;
    }
    const uint32_t n = 31 - __builtin_clz(q);
    return ((127 - 16 + n) << 23) | ((q << (23 - n)) & 0x7FFFFF);
}

// one corner stick (32 channels) -> the face rows at dst (channels 0..15) and dst + 512 (16..31). pos: RM page
// (VSRC 0) or slice * TPH * 32 + row (VSRC 1).
// VSRC 3 (BEVFORMER_DEC_VALUE=fused, D20 round 3): the value is the projection itself, TILE [1, 1, SUB h w, H * 32]
// (no head split): the head block j of map row pos is tile (pos / 32) * H + j; g_hj = the current item's head
static uint32_t g_hj = 0;

template <typename A>
static inline void gather(const A& v_acc, uint32_t pos, bool ok, uint64_t zeros, uint32_t dst) {
    if (!ok) {
        noc_async_read(zeros, dst, 32);
        noc_async_read(zeros, dst + 512, 32);
        return;
    }
    uint64_t s;
    uint32_t hi;
    if constexpr (VSRC == 0) {
        s = v_acc.get_noc_addr(pos);
        hi = 32;
    } else {
        const uint32_t r = pos & 31;
        const uint32_t t = VSRC == 3 ? (pos >> 5) * H + g_hj : pos >> 5;
        s = v_acc.get_noc_addr(t, ((r >> 4) << 10) + ((r & 15) << 5));
        hi = 512;
    }
#ifndef MSDA_K_NOREAD
    noc_async_read(s, dst, 32);
    noc_async_read(s + hi, dst + 512, 32);
#endif
}

static inline void zero_l1(uint32_t addr, uint32_t bytes) {
    const uint64_t z = get_noc_addr(MEM_ZEROS_BASE);
    for (uint32_t o = 0; o < bytes; o += MEM_ZEROS_SIZE) {
        noc_async_read(z, addr + o, MEM_ZEROS_SIZE);
    }
}

void kernel_main() {
    const uint32_t ci = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (ci >= NC) {
        return;
    }
    const uint32_t w0 = SKIP ? ci : ci * NWORK / NC;
    const uint32_t w1 = SKIP ? NWORK : (ci + 1) * NWORK / NC;
    const uint32_t wstep = SKIP ? NC : 1;
    if (w0 >= w1) {
        return;
    }
    const auto v_acc = TensorAccessor(v_args, get_common_arg_val<uint32_t>(0), VPAGE);
    const auto l_acc = TensorAccessor(l_args, get_common_arg_val<uint32_t>(1), TB32);
    const auto a_acc = TensorAccessor(a_args, get_common_arg_val<uint32_t>(2), TB32);
    const auto o_acc = TensorAccessor(o_args, get_common_arg_val<uint32_t>(3), TB16);
    const uint64_t zeros = get_noc_addr(MEM_ZEROS_BASE);

    if constexpr (MAKES_E) {
        // E_p (bf16): row p all ones, so W @ E_p broadcasts column p of W over the 32 channels
        cb_reserve_back(cb_e, P);
        const uint32_t e_addr = get_write_ptr(cb_e);
        zero_l1(e_addr, P * TB16);
        noc_async_read_barrier();
        volatile tt_l1_ptr uint16_t* e = reinterpret_cast<volatile tt_l1_ptr uint16_t*>(e_addr);
        for (uint32_t p = 0; p < P; ++p) {
            for (uint32_t c = 0; c < 32; ++c) {
                e[p * (TB16 / 2) + tile_elem(p, c)] = 0x3F80;
            }
        }
        cb_push_back(cb_e, P);
    }
    if constexpr (SM && MAKES_E) {
        // the softmax's reduce scalers, as ttnn's softmax reader makes them; the half masks (fp32 1 / 0 per column:
        // RISC h owns the points [h P / 2, (h + 1) P / 2)); the logit pages pre-filled with PAD_LOGIT
        dataflow_kernel_lib::calculate_and_prepare_reduce_scaler<cb_scmax, ckernel::PoolType::MAX,
                                                                 ckernel::ReduceDim::REDUCE_ROW>();
        dataflow_kernel_lib::calculate_and_prepare_reduce_scaler<cb_scsum, ckernel::PoolType::SUM,
                                                                 ckernel::ReduceDim::REDUCE_ROW>();
        cb_reserve_back(cb_msk, 2);
        volatile tt_l1_ptr uint32_t* m = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_msk));
        for (uint32_t h = 0; h < 2; ++h) {
            for (uint32_t r = 0; r < 32; ++r) {
                for (uint32_t c = 0; c < 32; ++c) {
                    m[h * 1024 + tile_elem(r, c)] = (c >= h * PH && c < (h + 1) * PH) ? 0x3F800000u : 0u;
                }
            }
        }
        cb_push_back(cb_msk, 2);
        cb_reserve_back(cb_lg, 2);
        volatile tt_l1_ptr uint32_t* lg = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_lg));
        for (uint32_t i = 0; i < 2048; ++i) {
            lg[i] = PAD_BITS;
        }
    }
    // the geo pages: zero once (columns >= P and the right faces stay zero / finite forever)
    cb_reserve_back(cb_geo, 6);
    zero_l1(get_write_ptr(cb_geo), 6 * TB32);
    if constexpr (SKIP) {
        cb_reserve_back(cb_in, 2 * NS);
        zero_l1(get_write_ptr(cb_in), 2 * NS * TB16);
    }
    cb_reserve_back(cb_loc, 1);
    const uint32_t lbuf = get_write_ptr(cb_loc);
    noc_async_read_barrier();
    volatile tt_l1_ptr int32_t* loc = reinterpret_cast<volatile tt_l1_ptr int32_t*>(lbuf);

    uint32_t prev = 0xFFFFFFFF;
    for (uint32_t w = w0; w < w1; w += wstep) {
        const uint32_t j = w / NB;
        const uint32_t b = w - j * NB;
        // VSRC 0: page of the map's pixel 0; VSRC 1: the slice's first tile row + the map's first row in it
        const uint32_t vbase = VSRC == 0   ? (j * SUB + b / BPS) * MAP
                               : VSRC == 2 ? (j * SUB + b / BPS) * HV * (WV + 1)
                               : VSRC == 3 ? (b / BPS) * MAP
                                           : j * TPH * 32 + (b / BPS) * MAP;
        g_hj = j;
        cb_reserve_back(cb_geo, 3);
        cb_reserve_back(cb_in, NS);
        const uint32_t g_addr = get_write_ptr(cb_geo);
        const uint32_t in_base = get_write_ptr(cb_in);
        const uint32_t ltile = LOC_HS == 0 ? j * NB + b : b * LOC_TW + (j * LOC_HS) / 32;
        const uint32_t lcol = LOC_HS == 0 ? 0 : (j * LOC_HS) % 32;
        noc_async_read(l_acc.get_noc_addr(ltile), lbuf, TB32);
        if constexpr (SM) {
            if constexpr (MAKES_E) {
                noc_async_read(a_acc.get_noc_addr(b * AT_TW + (j * P) / 32), get_write_ptr(cb_scr), TB32);
            }
        } else {
            noc_async_read(a_acc.get_noc_addr(j * NB + b), g_addr + 2 * TB32, TB32);
        }
        noc_async_read_barrier();
        if constexpr (SM && MAKES_E) {
            // the head's logit tile: columns 0..P-1 <- compact columns (j P) % 32 + p; columns >= P keep PAD
            cb_reserve_back(cb_lg, 1);
            volatile tt_l1_ptr uint32_t* lg = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_lg));
            volatile tt_l1_ptr uint32_t* src = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_scr));
            const uint32_t acol = (j * P) % 32;
            for (uint32_t r = 0; r < 32; ++r) {
                const bool ok = b * 32 + r < NQ;
                for (uint32_t p = 0; p < P; ++p) {
                    lg[tile_elem(r, p)] = ok ? src[tile_elem(r, acol + p)] : 0u;
                }
            }
            cb_push_back(cb_lg, 1);
        }
        volatile tt_l1_ptr uint32_t* fx = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(g_addr);
        volatile tt_l1_ptr uint32_t* fy = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(g_addr + TB32);
        volatile tt_l1_ptr uint32_t* at = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(g_addr + 2 * TB32);
        // the other RISC's points carry zero weight in this copy (their fx / fy here are stale but finite); SM: the
        // compute kernel applies the half mask to the softmax instead
        if constexpr (!SM) {
            for (uint32_t r = 0; r < 32; ++r) {
                for (uint32_t p = 0; p < P; ++p) {
                    if (p < P0 || p >= P0 + PH) {
                        at[tile_elem(r, p)] = 0;
                    }
                }
            }
        }
        for (uint32_t r = 0; r < 32; ++r) {
            const uint32_t roff = (((r >> 4) << 1) << 9) + ((r & 15) << 5);  // face (r / 16, 0), row r % 16
            const bool row_ok = b * 32 + r < NQ;
            for (uint32_t p = P0; p < P0 + PH; ++p) {
                const uint32_t dst = in_base + (p - P0) * 4 * TB16 + roff;
                const uint32_t ew = tile_elem(r, p);
                if (!row_ok) {
                    fx[ew] = 0;
                    fy[ew] = 0;
                    if constexpr (!SM) {
                        at[ew] = 0;
                    }
                    if constexpr (SKIP) {
                        continue;
                    }
                    if constexpr (VSRC == 2) {
                        const uint32_t blk = in_base + (p - P0) * 4 * TB16 + r * 128;
                        noc_async_read(zeros, blk, 128);
                        noc_async_read(zeros, blk + 2 * TB16, 128);
                        continue;
                    }
                    for (uint32_t c = 0; c < 4; ++c) {
                        noc_async_read(zeros, dst + c * TB16, 32);
                        noc_async_read(zeros, dst + c * TB16 + 512, 32);
                    }
                    continue;
                }
                int32_t vx = loc[tile_elem(r, lcol + 2 * p)];
                int32_t vy = loc[tile_elem(r, lcol + 2 * p + 1)];
                vx = vx > LOC_LIM ? LOC_LIM : (vx < -LOC_LIM ? -LOC_LIM : vx);
                vy = vy > LOC_LIM ? LOC_LIM : (vy < -LOC_LIM ? -LOC_LIM : vy);
                const int32_t x0 = vx >> 16;
                const int32_t y0 = vy >> 16;
                fx[ew] = q16_to_f32(static_cast<uint32_t>(vx) & 0xFFFF);
                fy[ew] = q16_to_f32(static_cast<uint32_t>(vy) & 0xFFFF);
                const bool xv0 = x0 >= 0 && x0 < WV, xv1 = x0 + 1 >= 0 && x0 + 1 < WV;
                const bool yv0 = y0 >= 0 && y0 < HV, yv1 = y0 + 1 >= 0 && y0 + 1 < HV;
                if constexpr (SKIP) {
                    static_assert(!(SKIP && SM), "msda_k: SKIP and SM are exclusive");
                    if (!((yv0 || yv1) && (xv0 || xv1))) {
                        at[ew] = 0;
                        continue;
                    }
                }
                if constexpr (VSRC == 2) {
                    // staging blocks of point p: top (NW | NE) at + 0, bottom (SW | SE) at + 2 tiles; row r: 128 B
                    const uint32_t blk = in_base + (p - P0) * 4 * TB16 + r * 128;
                    const bool xp = x0 >= -1 && x0 < WV;                 // page i = x0 + 1 in [0, WV]
                    const uint32_t pg = vbase + static_cast<uint32_t>(y0 * (WV + 1) + x0 + 1);
                    if (xp && yv0) {
                        noc_async_read(v_acc.get_noc_addr(pg), blk, 128);
                    } else {
                        noc_async_read(zeros, blk, 128);
                    }
                    if (xp && yv1) {
                        noc_async_read(v_acc.get_noc_addr(pg + WV + 1), blk + 2 * TB16, 128);
                    } else {
                        noc_async_read(zeros, blk + 2 * TB16, 128);
                    }
                    continue;
                }
                const uint32_t pos = vbase + static_cast<uint32_t>(y0 * WV + x0);
                gather(v_acc, pos, yv0 && xv0, zeros, dst);
                gather(v_acc, pos + 1, yv0 && xv1, zeros, dst + TB16);
                gather(v_acc, pos + WV, yv1 && xv0, zeros, dst + 2 * TB16);
                gather(v_acc, pos + WV + 1, yv1 && xv1, zeros, dst + 3 * TB16);
            }
        }
        noc_async_read_barrier();
        cb_push_back(cb_geo, 3);
        cb_push_back(cb_in, NS);
        if constexpr (WRITES_OUT) {
            if (prev != 0xFFFFFFFF) {
                const uint32_t jp = prev / NB;
                const uint32_t bp = prev - jp * NB;
                cb_wait_front(cb_out, 1);
                noc_async_write(get_read_ptr(cb_out), o_acc.get_noc_addr(bp * H + jp), TB16);
                noc_async_writes_flushed();
                cb_pop_front(cb_out, 1);
            }
        }
        prev = w;
    }
    if constexpr (WRITES_OUT) {
        const uint32_t jp = prev / NB;
        const uint32_t bp = prev - jp * NB;
        cb_wait_front(cb_out, 1);
        noc_async_write(get_read_ptr(cb_out), o_acc.get_noc_addr(bp * H + jp), TB16);
        noc_async_write_barrier();
        cb_pop_front(cb_out, 1);
    }
}