// SPDX-License-Identifier: Apache-2.0
// MK-T compute (sceneseg-p150 OPT_PLAN 3.4, the D19 seg-tail module megakernel, tt/mkt.py).
//
// Phase 1 (decode_layer_9), per tile column (the reader's (n + 2) x 12 activation tiles, A[(r * 3 + dx) * 4 + k]):
//   y[i][c] = sum over taps (dy, dx) and K tiles k of A[i + dy][dx][k] @ W9[(dy * 3 + dx) * 4 + k][c]
//             matmul_block(rt = n rows, ct = 2 column tiles, row stride kt = 12), bf16 x bf16, fp32 DEST
//   y = y + b9[c]       one more matmul_block: the ones tile x [b_hi; b_lo] (the fp32 bias as two bf16 terms, 16
//                       significant bits; 1.0 x bf16 is exact), accumulated in DEST after the conv K loop. No
//                       unpack-to-DEST copy: an fp32 bias copy_tile after the K loop hung in DEST full-sync mode
//                       (logs/sceneseg/opt_mk/HANG_1.txt)
//   y = gelu(y)         gelu_tile<false> (ttnn.gelu, exact erf form; EPI_ACT 2)
//   y = bf16(y)         typecast_tile<Float32, Float16_b> (EPI_OUT_BF16), packed to CB 16 (the writer's D store)
// Phase 2 (margins), per tile column (A[(r * 3 + dx) * 2 + k] from the D store):
//   m[i] = sum A[i + dy][dx][k] @ WM[(dy * 3 + dx) * 2 + k] (rt = n, ct = 1, row stride 6) + ones x [bm_hi; bm_lo]
//   -> fp32 margin tiles to CB 17 (the writer decides the classes, K8).
// DEST full sync (8 fp32 tiles): phase 1 uses tiles 0 .. 2n - 1 (<= 6), phase 2 0 .. n - 1.
// CT args: [0] H, [1] W, [2] NC, [3] GX, [4] RB, [5] RE.
#include <cstdint>

#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/gelu.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/matmul.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"

using namespace ckernel;

#ifndef MKT_DBG
#define MKT_DBG 0
#endif

constexpr uint32_t W = get_compile_time_arg_val(1);
constexpr uint32_t NC = get_compile_time_arg_val(2);
constexpr uint32_t GX = get_compile_time_arg_val(3);
constexpr uint32_t RB = get_compile_time_arg_val(4);
constexpr uint32_t RE = get_compile_time_arg_val(5);
constexpr uint32_t XT = W / 32;
constexpr uint32_t NRMAX = RB + (RE ? 1 : 0) + 2;
constexpr uint32_t cb_a = 0, cb_w9 = 1, cb_b9 = 2, cb_wm = 3, cb_bm = 4, cb_d9 = 16, cb_m = 17, cb_one = 10;

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    const uint32_t n = RB + (core < RE ? 1 : 0);
    compute_kernel_hw_startup(cb_a, cb_w9, cb_d9);
    cb_wait_front(cb_w9, 72);
    cb_wait_front(cb_b9, 2);
    cb_wait_front(cb_wm, 18);
    cb_wait_front(cb_bm, 1);
    cb_wait_front(cb_one, 3);

    // ---------------- phase 1
    const uint32_t no = 2 * n;
    for (uint32_t x = 0; x < XT; ++x) {
        cb_wait_front(cb_a, NRMAX * 12);         // the reader's fixed push size (rows nr .. NRMAX - 1 unused)
        tile_regs_acquire();
        matmul_block_init(cb_a, cb_w9, 0, 2, n, 12);
        for (uint32_t t = 0; t < 9; ++t) {
            for (uint32_t k = 0; k < 4; ++k) {
                const uint32_t kt = t * 4 + k;
                // tap t = dy * 3 + dx: A tile of row i is (i + dy) * 12 + dx * 4 + k = i * 12 + kt
                matmul_block(cb_a, cb_w9, kt, kt * 2, 0, 0, 2, n, 12);
            }
        }
        // + bias: ones tile (columns 0 / 1 = 1.0, one per output row: row stride 1) x [b_hi; b_lo]
        matmul_block_init(cb_one, cb_b9, 0, 2, n, 1);
        matmul_block(cb_one, cb_b9, 0, 0, 0, 0, 2, n, 1);
        gelu_tile_init<false>();
        for (uint32_t j = 0; j < no; ++j) {
            gelu_tile<false>(j);
        }
        typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
        for (uint32_t j = 0; j < no; ++j) {
            typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(j);
        }
        tile_regs_commit();
        cb_reserve_back(cb_d9, no);
        tile_regs_wait();
        for (uint32_t j = 0; j < no; ++j) {
            pack_tile(j, cb_d9);
        }
        tile_regs_release();
        cb_push_back(cb_d9, no);
        cb_pop_front(cb_a, NRMAX * 12);
    }

#if MKT_DBG != 1
    // ---------------- phase 2
    pack_reconfig_data_format(cb_m);
    for (uint32_t x = 0; x < XT; ++x) {
        cb_wait_front(cb_a, NRMAX * 6);
        tile_regs_acquire();
        matmul_block_init(cb_a, cb_wm, 0, 1, n, 6);
        for (uint32_t kt = 0; kt < 18; ++kt) {
            matmul_block(cb_a, cb_wm, kt, kt, 0, 0, 1, n, 6);
        }
        matmul_block_init(cb_one, cb_bm, 0, 1, n, 1);
        matmul_block(cb_one, cb_bm, 0, 0, 0, 0, 1, n, 1);
        tile_regs_commit();
        cb_reserve_back(cb_m, n);
        tile_regs_wait();
        for (uint32_t i = 0; i < n; ++i) {
            pack_tile(i, cb_m);
        }
        tile_regs_release();
        cb_push_back(cb_m, n);
        cb_pop_front(cb_a, NRMAX * 6);
    }
#endif
    cb_pop_front(cb_w9, 72);
    cb_pop_front(cb_b9, 2);
    cb_pop_front(cb_wm, 18);
    cb_pop_front(cb_bm, 1);
    cb_pop_front(cb_one, 3);
}