File size: 4,288 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 | // SPDX-License-Identifier: Apache-2.0
// pairmap_k (tt/pairmap_k.py, BEVFORMER_MSDA_PAIRS): the SCA value projection -> horizontal pixel-pair maps for
// msda_k's 128-byte corner reads (one read per point and row: the NW + NE or SW + SE sticks).
//
// in : bf16 TILE [B, 1, SUB * HV * WV, HEADS * 32], L1 interleaved (32-byte reads at 32-byte offsets); rows =
// camera-major pixels, head j in columns j * 32 ..: the stick of (map m = (b HEADS + j) SUB + s, y, x) = row
// s * HV * WV + y * WV + x of batch b, tile column j (TPB tiles per batch): two 32-byte face rows (faces
// 2 (r / 16) and + 1, row r % 16, r = row % 32)
// out: bf16 ROW_MAJOR [B * HEADS * SUB, HV, WV + 1, 64], interleaved, 128-byte pages: page (m, y, i) = the sticks of
// pixels x = i - 1 and x = i of row y (zeros for x = -1 and x = WV)
// Work: pages split in contiguous ranges over the cores, then alternately over the two data-movement RISC-Vs (SLOT 0 /
// 1, each its own scratch CB of BATCH pages): per batch, 4 reads per page (zeros page outside the row) into the
// scratch, one barrier, BATCH page writes, one write barrier. No semaphores, no multicast; common RT args only (P3).
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t SLOT = get_compile_time_arg_val(0);
constexpr uint32_t cb_s = get_compile_time_arg_val(1);
constexpr uint32_t HV = get_compile_time_arg_val(2);
constexpr uint32_t WV = get_compile_time_arg_val(3);
constexpr uint32_t SUB = get_compile_time_arg_val(4);
constexpr uint32_t HEADS = get_compile_time_arg_val(5);
constexpr uint32_t NPAGES = 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 uint32_t BATCH = get_compile_time_arg_val(9);
constexpr uint32_t TPB = get_compile_time_arg_val(10); // input tiles per batch
constexpr auto i_args = TensorAccessorArgs<11>();
constexpr auto o_args = TensorAccessorArgs<i_args.next_compile_time_args_offset()>();
constexpr uint32_t TB16 = 2048, PAGE = 128, W1 = WV + 1, MAP = HV * WV;
template <typename A>
static inline void stick(const A& ia, uint32_t m, uint32_t y, int32_t x, uint64_t zeros, uint32_t dst) {
if (x < 0 || x >= static_cast<int32_t>(WV)) {
noc_async_read(zeros, dst, 32);
noc_async_read(zeros, dst + 32, 32);
return;
}
const uint32_t bt = m / (HEADS * SUB), jm = m - bt * HEADS * SUB;
const uint32_t j = jm / SUB, s = jm - j * SUB;
const uint32_t row = s * MAP + y * WV + static_cast<uint32_t>(x);
const uint32_t r = row & 31;
const uint32_t off = ((r >> 4) << 10) + ((r & 15) << 5);
const uint64_t a = ia.get_noc_addr(bt * TPB + (row >> 5) * HEADS + j, off);
noc_async_read(a, dst, 32); // channels 0..15 (face 2 (r / 16))
noc_async_read(a + 512, dst + 32, 32); // channels 16..31 (the next face)
}
void kernel_main() {
const uint32_t ci = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (ci >= NC) {
return;
}
const uint32_t p0 = ci * NPAGES / NC;
const uint32_t p1 = (ci + 1) * NPAGES / 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);
const uint64_t zeros = get_noc_addr(MEM_ZEROS_BASE);
cb_reserve_back(cb_s, 1);
const uint32_t buf = get_write_ptr(cb_s);
// batches b0 = p0 + (2 SLOT) BATCH, step 2 BATCH: the two RISC-Vs alternate whole batches
for (uint32_t b0 = p0 + SLOT * BATCH; b0 < p1; b0 += 2 * BATCH) {
const uint32_t b1 = b0 + BATCH < p1 ? b0 + BATCH : p1;
for (uint32_t p = b0; p < b1; ++p) {
const uint32_t m = p / (HV * W1);
const uint32_t rem = p - m * HV * W1;
const uint32_t y = rem / W1;
const int32_t i = static_cast<int32_t>(rem - y * W1);
const uint32_t dst = buf + (p - b0) * PAGE;
stick(ia, m, y, i - 1, zeros, dst);
stick(ia, m, y, i, zeros, dst + 64);
}
noc_async_read_barrier();
for (uint32_t p = b0; p < b1; ++p) {
noc_async_write(buf + (p - b0) * PAGE, oa.get_noc_addr(p), PAGE);
}
noc_async_write_barrier();
}
}
|