code/tt_bevformer/tt/kernels/dcn_k_reader.cpp

// SPDX-License-Identifier: Apache-2.0
// dcn_k (tt/dcn_k.py, OPT_PLAN #3): the gather half of the DCNv2 im2col kernel. The same source runs on both
// data-movement RISC-Vs: PART 0 (NCRISC) gathers the corners NW, NE and writes the geometry tiles, PART 1 (BRISC)
// gathers SW, SE and writes the output tiles. Each RISC has its own CBs.
//
// Work item w (contiguous range per core; core index from the logical coordinates, common RT args only): block
// b = w / 9 (32 output pixels, rows b * 32 .. of the [N H W] map), tap k = w % 9.
//   x    bf16 ROW_MAJOR [1, 1, N H W, C], L1 interleaved: page q = the C-channel stick of pixel q;
//   loc  int32 TILE [1, 1, N H W, 32]: columns 2k, 2k + 1 = round(pixel * 2^16) of tap k's sampling point (x, y);
//   sig  fp32 TILE [1, 1, N H W, 32]: column 18 + k = sigmoid(mask logit of tap k);
//   out  bf16 TILE [1, 1, N H W, 9 C]: the modulated im2col, tiles b * 9 CT + k * CT + t (CT = C / 32).
// Per item: cb_rm (this RISC) <- 2 x 32 rows x C bf16 row-major (corner 2 PART + i at pages i CT .. i CT + CT - 1);
// PART 0 also cb_geo <- fx, fy, mask of the item's 32 rows in column 0 (fp32 bit patterns; the other columns zero).
#include <stdint.h>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t PART = get_compile_time_arg_val(1);
constexpr uint32_t cb_rm = 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_out = get_compile_time_arg_val(5);
constexpr int32_t HV = get_compile_time_arg_val(6);
constexpr int32_t WV = get_compile_time_arg_val(7);
constexpr uint32_t NQ = get_compile_time_arg_val(8);
constexpr uint32_t NB = get_compile_time_arg_val(9);
constexpr uint32_t C = get_compile_time_arg_val(10);
constexpr uint32_t NWORK = get_compile_time_arg_val(11);
constexpr uint32_t NC = get_compile_time_arg_val(12);
constexpr uint32_t GX = get_compile_time_arg_val(13);
constexpr uint32_t WRITES_OUT = get_compile_time_arg_val(14);
constexpr uint32_t WRITES_GEO = get_compile_time_arg_val(15);
// B0: the first 32-row block of this call (BEVFORMER_DCN_IM2COL=l1 runs the im2col in row chunks whose output is an
// L1 height-sharded tensor the GEMM reads in place): items cover blocks B0 .. B0 + NB - 1, loc / sig / x are read at
// the absolute rows, the output tiles are written at the chunk's own rows
constexpr uint32_t B0 = get_compile_time_arg_val(16);
constexpr auto x_args = TensorAccessorArgs<17>();
constexpr auto l_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
constexpr auto s_args = TensorAccessorArgs<l_args.next_compile_time_args_offset()>();
constexpr auto o_args = TensorAccessorArgs<s_args.next_compile_time_args_offset()>();

constexpr uint32_t TB16 = 2048, TB32 = 4096;
#ifdef DCN_K_GLUE
// BEVFORMER_DCN_GLUE=kernel (D20 round 2): the glue runs in the compute kernel. `loc` is the raw offset conv output
// (fp32 TILE [NHW, 32], read in place, e.g. the conv's L1 shards), `sig` the B table (fp32 TILE [NHW, 32]: base * S +
// O) and `sc` the S tile (fp32 TILE, row 0 replicated: tile 0 of the full scale map). PART 0 reads S once and each
// block's off / B tiles into cb 9 / 10, the compute kernel returns loc (int32) to both RISCs (cb 6 / 7) and sig
// (fp32) to PART 0 (cb 8): the same LLK steps as glue_k (bit-identical to the separate program).
constexpr uint32_t cb_goff = 9, cb_gb = 10, cb_gs = 11, cb_sig0 = 8;
constexpr auto sc_args = TensorAccessorArgs<o_args.next_compile_time_args_offset()>();
#endif
constexpr uint32_t CT = C / 32;
constexpr uint32_t ROWB = C * 2;
constexpr uint32_t HW = static_cast<uint32_t>(HV * WV);
constexpr int32_t LOC_LIM = 1 << 28;
static_assert(C % 32 == 0 && ROWB % MEM_ZEROS_SIZE == 0, "dcn_k: C a multiple of 256");

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

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);
}

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 = ci * NWORK / NC;
    const uint32_t w1 = (ci + 1) * NWORK / NC;
    if (w0 == w1) {
        return;
    }
#ifdef DCN_K_PAIRS
    // BEVFORMER_DCN_PAIRS (D20 round 2): x is the pixel-pair map [N HV (WV + 1), 2 C] (dcnpair_k): page (n, y, i) =
    // the sticks of x = i - 1 and x = i; one 2C-wide read per row brings both corners of this RISC (NW + NE or
    // SW + SE) into cb_rm as one row of a 32 x 2C row-major block (the compute kernel tilizes it 2 CT tiles wide)
    const auto x_acc = TensorAccessor(x_args, get_common_arg_val<uint32_t>(0), 2 * ROWB);
#else
    const auto x_acc = TensorAccessor(x_args, get_common_arg_val<uint32_t>(0), ROWB);
#endif
    const auto l_acc = TensorAccessor(l_args, get_common_arg_val<uint32_t>(1), TB32);
    const auto s_acc = TensorAccessor(s_args, get_common_arg_val<uint32_t>(2), TB32);
    const auto o_acc = TensorAccessor(o_args, get_common_arg_val<uint32_t>(3), TB16);

    if constexpr (WRITES_GEO) {
        cb_reserve_back(cb_geo, 6);
        zero_l1(get_write_ptr(cb_geo), 6 * TB32);
    }
#ifdef DCN_K_GLUE
    if constexpr (PART == 0) {
        const auto sc_acc = TensorAccessor(sc_args, get_common_arg_val<uint32_t>(4), TB32);
        cb_reserve_back(cb_gs, 1);
        noc_async_read(sc_acc.get_noc_addr(0), get_write_ptr(cb_gs), TB32);
        noc_async_read_barrier();
        cb_push_back(cb_gs, 1);
    }
    volatile tt_l1_ptr int32_t* loc = nullptr;
    volatile tt_l1_ptr uint32_t* sig = nullptr;
#else
    cb_reserve_back(cb_loc, 2);
    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);
    volatile tt_l1_ptr uint32_t* sig = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(lbuf + TB32);
#endif
    uint32_t cur_b = 0xFFFFFFFF;
    for (uint32_t w = w0; w < w1; ++w) {
        const uint32_t b = w / 9 + B0;                                // absolute block
        const uint32_t k = w - (b - B0) * 9;
        if (b != cur_b) {
#ifdef DCN_K_GLUE
            if (cur_b != 0xFFFFFFFF) {                                // the previous block's loc / sig pages
                cb_pop_front(cb_loc, 1);
                if constexpr (WRITES_GEO) {
                    cb_pop_front(cb_sig0, 1);
                }
            }
            if constexpr (PART == 0) {                                // this block's glue inputs
                cb_reserve_back(cb_goff, 1);
                cb_reserve_back(cb_gb, 1);
                noc_async_read(l_acc.get_noc_addr(b), get_write_ptr(cb_goff), TB32);
                noc_async_read(s_acc.get_noc_addr(b), get_write_ptr(cb_gb), TB32);
                noc_async_read_barrier();
                cb_push_back(cb_goff, 1);
                cb_push_back(cb_gb, 1);
            }
            cb_wait_front(cb_loc, 1);
            loc = reinterpret_cast<volatile tt_l1_ptr int32_t*>(get_read_ptr(cb_loc));
            if constexpr (WRITES_GEO) {
                cb_wait_front(cb_sig0, 1);
                sig = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_read_ptr(cb_sig0));
            }
#else
            noc_async_read(l_acc.get_noc_addr(b), lbuf, TB32);
            if constexpr (WRITES_GEO) {
                noc_async_read(s_acc.get_noc_addr(b), lbuf + TB32, TB32);
            }
            noc_async_read_barrier();
#endif
            cur_b = b;
        }
#ifdef DCN_K_NEAREST
        // BEVFORMER_DCN_SAMPLE=nearest (D20 round 3): one corner per sampling point (the nearest pixel, round half
        // up); the items alternate between the RISCs (item w0 + 2i on PART 0, w0 + 2i + 1 on PART 1), each gathering
        // its item's 32 sticks into its own cb_rm (CT pages); PART 0 writes the mask tile (column 0) of every item
        {
            const bool mine = ((w - w0) & 1) == PART;
            uint32_t rm = 0;
            if (mine) {
                cb_reserve_back(cb_rm, CT);
                rm = get_write_ptr(cb_rm);
            }
            uint32_t g_addr = 0;
            if constexpr (WRITES_GEO) {
                cb_reserve_back(cb_geo, 1);
                g_addr = get_write_ptr(cb_geo);
            }
            volatile tt_l1_ptr uint32_t* ma = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(g_addr);
            if (mine || WRITES_GEO) {
                for (uint32_t r = 0; r < 32; ++r) {
                    const uint32_t q = b * 32 + r;
                    const uint32_t e = tile_elem(r, 0);
                    if (q >= NQ) {
                        if constexpr (WRITES_GEO) {
                            ma[e] = 0;
                        }
                        if (mine) {
                            zero_l1(rm + r * ROWB, ROWB);
                        }
                        continue;
                    }
                    if constexpr (WRITES_GEO) {
                        ma[e] = sig[tile_elem(r, 18 + k)];
                    }
                    if (!mine) {
                        continue;
                    }
                    const uint32_t el = tile_elem(r, 2 * k);
                    int32_t vx = loc[el];
                    int32_t vy = loc[el + 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 + 0x8000) >> 16;
                    const int32_t y0 = (vy + 0x8000) >> 16;
                    if (y0 >= 0 && y0 < HV && x0 >= 0 && x0 < WV) {
                        const uint32_t pos = (q / HW) * HW + static_cast<uint32_t>(y0 * WV + x0);
                        noc_async_read(x_acc.get_noc_addr(pos), rm + r * ROWB, ROWB);
                    } else {
                        zero_l1(rm + r * ROWB, ROWB);
                    }
                }
            }
            noc_async_read_barrier();
            if constexpr (WRITES_GEO) {
                cb_push_back(cb_geo, 1);
            }
            if (mine) {
                cb_push_back(cb_rm, CT);
            }
        }
#else
        cb_reserve_back(cb_rm, 2 * CT);
        const uint32_t rm = get_write_ptr(cb_rm);
        uint32_t g_addr = 0;
        if constexpr (WRITES_GEO) {
            cb_reserve_back(cb_geo, 3);
            g_addr = get_write_ptr(cb_geo);
        }
        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* ma = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(g_addr + 2 * TB32);
        for (uint32_t r = 0; r < 32; ++r) {
            const uint32_t q = b * 32 + r;
            const uint32_t dst0 = rm + r * ROWB;
            const uint32_t dst1 = dst0 + CT * TB16;
#ifdef DCN_K_MM
            constexpr bool DIAG = true;   // BEVFORMER_DCN_SUM=mm: fx / fy / mask on the diagonal (diag(w_c) tiles)
#else
            constexpr bool DIAG = false;
#endif
            if (q >= NQ) {
                if constexpr (WRITES_GEO) {
                    const uint32_t e = tile_elem(r, DIAG ? r : 0);
                    fx[e] = 0;
                    fy[e] = 0;
                    ma[e] = 0;
                }
#ifdef DCN_K_PAIRS
                zero_l1(rm + r * 2 * ROWB, 2 * ROWB);
#else
                zero_l1(dst0, ROWB);
                zero_l1(dst1, ROWB);
#endif
                continue;
            }
            const uint32_t el = tile_elem(r, 2 * k);
            int32_t vx = loc[el];
            int32_t vy = loc[el + 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);
#ifdef DCN_K_NEAREST_EMU
            // accuracy probe of DCN_K_NEAREST on the bilinear path: zero fractions -> weight 1 on NW only
            vx = (vx + 0x8000) & ~static_cast<int32_t>(0xFFFF);
            vy = (vy + 0x8000) & ~static_cast<int32_t>(0xFFFF);
#endif
            const int32_t x0 = vx >> 16;
            const int32_t y0 = (vy >> 16) + static_cast<int32_t>(PART);  // PART 1: the lower row (SW, SE)
#ifndef DCN_K_NOGEO
            if constexpr (WRITES_GEO) {
                const uint32_t e = tile_elem(r, DIAG ? r : 0);
                fx[e] = q16_to_f32(static_cast<uint32_t>(vx) & 0xFFFF);
                fy[e] = q16_to_f32(static_cast<uint32_t>(vy) & 0xFFFF);
                ma[e] = sig[tile_elem(r, 18 + k)];
            }
#endif
            const bool yv = y0 >= 0 && y0 < HV;
#ifdef DCN_K_PAIRS
            {
                const uint32_t pd = rm + r * 2 * ROWB;
#ifdef DCN_K_NOREAD
                (void)pd;
                continue;
#endif
                if (yv && x0 >= -1 && x0 < WV) {
                    const uint32_t pp = (q / HW) * (HV * (WV + 1)) + static_cast<uint32_t>(y0 * (WV + 1) + x0 + 1);
                    noc_async_read(x_acc.get_noc_addr(pp), pd, 2 * ROWB);
                } else {
                    zero_l1(pd, 2 * ROWB);
                }
                continue;
            }
#endif
            const uint32_t nbase = (q / HW) * HW;
            const bool xv0 = x0 >= 0 && x0 < WV, xv1 = x0 + 1 >= 0 && x0 + 1 < WV;
            const uint32_t pos = nbase + static_cast<uint32_t>(y0 * WV + x0);
#ifndef DCN_K_NOREAD
            if (yv && xv0) {
#if defined(DCN_K_RLOCAL)
                noc_async_read(get_noc_addr(dst1), dst0, ROWB);      // (anatomy probe: own L1, no mesh hops)
#elif defined(DCN_K_RSMALL)
                noc_async_read(x_acc.get_noc_addr(pos), dst0, 32);   // (anatomy probe: same requests, 1/16 bytes)
#elif defined(DCN_K_RSELF)
                noc_async_read(x_acc.get_noc_addr(q), dst0, ROWB);   // (anatomy probe: the pixel's own page)
#else
                noc_async_read(x_acc.get_noc_addr(pos), dst0, ROWB);
#endif
            } else {
                zero_l1(dst0, ROWB);
            }
            if (yv && xv1) {
#if defined(DCN_K_RLOCAL)
                noc_async_read(get_noc_addr(dst0), dst1, ROWB);
#elif defined(DCN_K_RSMALL)
                noc_async_read(x_acc.get_noc_addr(pos + 1), dst1, 32);
#elif defined(DCN_K_RSELF)
                noc_async_read(x_acc.get_noc_addr(q), dst1, ROWB);
#else
                noc_async_read(x_acc.get_noc_addr(pos + 1), dst1, ROWB);
#endif
            } else {
                zero_l1(dst1, ROWB);
            }
#else
            (void)pos;
#endif
        }
        noc_async_read_barrier();
        if constexpr (WRITES_GEO) {
            cb_push_back(cb_geo, 3);
        }
        cb_push_back(cb_rm, 2 * CT);
#endif  // DCN_K_NEAREST
        if constexpr (WRITES_OUT) {
            if (w > w0) {
                const uint32_t bp = (w - 1) / 9;
                const uint32_t kp = (w - 1) - bp * 9;
                cb_wait_front(cb_out, CT);
                uint32_t rp = get_read_ptr(cb_out);
                const uint32_t t0 = bp * 9 * CT + kp * CT;
#ifndef DCN_K_NOWRITE
                for (uint32_t t = 0; t < CT; ++t) {
                    noc_async_write(rp + t * TB16, o_acc.get_noc_addr(t0 + t), TB16);
                }
#endif
                noc_async_writes_flushed();
                cb_pop_front(cb_out, CT);
            }
        }
    }
    if constexpr (WRITES_OUT) {
        const uint32_t bp = (w1 - 1) / 9;
        const uint32_t kp = (w1 - 1) - bp * 9;
        cb_wait_front(cb_out, CT);
        uint32_t rp = get_read_ptr(cb_out);
        const uint32_t t0 = bp * 9 * CT + kp * CT;
        for (uint32_t t = 0; t < CT; ++t) {
            noc_async_write(rp + t * TB16, o_acc.get_noc_addr(t0 + t), TB16);
        }
        noc_async_write_barrier();
        cb_pop_front(cb_out, CT);
    }
}