File size: 4,748 Bytes
55474c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | // SPDX-License-Identifier: Apache-2.0
// pairmap2 (tt/pairmap_k.py pair_maps, BEVFORMER_PAIRMAP=v2; D20 round 2), data movement: the same output as
// pairmap_k.cpp (horizontal pixel-pair maps, 128-byte pages (m, y, i) = the sticks of x = i - 1 and x = i, zeros
// outside the row) from whole tiles instead of 32-byte face-row reads: items (batch, tile row, head) in contiguous
// ranges over the cores; per item ROLE 0 (NCRISC) reads the 2 KB tile (-> the compute kernel's pack_untilize: 32
// pixel sticks of 64 B, consecutive in cb_rm) and the previous pixel's stick (the last row of the tile above, two
// 32-byte face rows); ROLE 1 (BRISC) writes each page as ONE 128-byte write of two consecutive untilized sticks (two
// 64-byte writes at row starts / tile starts, and the extra end-of-row page). About 1.1 requests per page instead
// of 5. No semaphores, no multicast; common RT args only (P3).
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t ROLE = get_compile_time_arg_val(0);
constexpr uint32_t HV = get_compile_time_arg_val(1);
constexpr uint32_t WV = get_compile_time_arg_val(2);
constexpr uint32_t SUB = get_compile_time_arg_val(3);
constexpr uint32_t HEADS = get_compile_time_arg_val(4);
constexpr uint32_t NTR = get_compile_time_arg_val(5); // tile rows per batch
constexpr uint32_t NITEMS = get_compile_time_arg_val(6);
constexpr uint32_t NC = get_compile_time_arg_val(7);
constexpr uint32_t GX = get_compile_time_arg_val(8);
constexpr auto i_args = TensorAccessorArgs<9>();
constexpr auto o_args = TensorAccessorArgs<i_args.next_compile_time_args_offset()>();
constexpr uint32_t TB16 = 2048, PAGE = 128, STK = 64, W1 = WV + 1, MAP = HV * WV;
constexpr uint32_t cb_in = 0, cb_prev = 1, cb_rm = 16;
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 * NITEMS / NC;
const uint32_t w1 = (ci + 1) * NITEMS / NC;
const auto ia = TensorAccessor(i_args, get_common_arg_val<uint32_t>(0), TB16);
const auto oa = TensorAccessor(o_args, get_common_arg_val<uint32_t>(1), PAGE);
for (uint32_t w = w0; w < w1; ++w) {
const uint32_t j = w % HEADS;
const uint32_t bt_t = w / HEADS; // bt * NTR + t
const uint32_t t = bt_t % NTR;
const uint32_t bt = bt_t / NTR;
if constexpr (ROLE == 0) {
cb_reserve_back(cb_in, 1);
cb_reserve_back(cb_prev, 1);
noc_async_read(ia.get_noc_addr(bt_t * HEADS + j), get_write_ptr(cb_in), TB16);
if (t > 0) { // the previous pixel = row 31 of the tile above: faces 2 / 3, row 15
const uint64_t a = ia.get_noc_addr((bt_t - 1) * HEADS + j, 1024 + 15 * 32);
noc_async_read(a, get_write_ptr(cb_prev), 32);
noc_async_read(a + 512, get_write_ptr(cb_prev) + 32, 32);
}
noc_async_read_barrier();
cb_push_back(cb_in, 1);
cb_push_back(cb_prev, 1);
} else {
cb_wait_front(cb_rm, 1);
cb_wait_front(cb_prev, 1);
const uint32_t rm = get_read_ptr(cb_rm);
const uint32_t prev = get_read_ptr(cb_prev);
const uint32_t zeros = MEM_ZEROS_BASE;
for (uint32_t r = 0; r < 32; ++r) {
const uint32_t q = t * 32 + r; // pixel within the batch: (s, y, x)
if (q >= SUB * MAP) {
break;
}
const uint32_t s = q / MAP;
const uint32_t rem = q - s * MAP;
const uint32_t y = rem / WV;
const uint32_t x = rem - y * WV;
const uint32_t m = (bt * HEADS + j) * SUB + s;
const uint32_t pg = (m * HV + y) * W1 + x;
const uint32_t cur = rm + r * STK;
if (x > 0 && r > 0) {
noc_async_write(cur - STK, oa.get_noc_addr(pg), PAGE);
} else {
const uint64_t d = oa.get_noc_addr(pg);
noc_async_write(x == 0 ? zeros : prev, d, STK);
noc_async_write(cur, d + STK, STK);
}
if (x == WV - 1) { // the end-of-row page (y, WV) = [x = WV - 1 | zeros]
const uint64_t d = oa.get_noc_addr(pg + 1);
noc_async_write(cur, d, STK);
noc_async_write(zeros, d + STK, STK);
}
}
noc_async_writes_flushed();
cb_pop_front(cb_rm, 1);
cb_pop_front(cb_prev, 1);
}
}
if constexpr (ROLE == 1) {
noc_async_write_barrier();
}
}
|