File size: 3,326 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 | // SPDX-License-Identifier: Apache-2.0
// dcnpair_k (tt/dcn_k.py dcn_pair_map, BEVFORMER_DCN_PAIRS; D20 round 2): the DCN input map -> horizontal pixel-pair
// map for dcn_k's corner reads (one 2C-wide read per row and corner pair: NW + NE, or SW + SE).
// in : bf16 ROW_MAJOR [1, 1, N HV WV, C], interleaved (one C-channel stick per page; L1)
// out: bf16 ROW_MAJOR [1, 1, N HV (WV + 1), 2 C], interleaved: page (n, y, i) = the sticks of pixels x = i - 1 and
// x = i of row y (zeros for x = -1 and x = WV)
// Work: pages in contiguous ranges over the cores, batches of BATCH pages alternating over the two data-movement
// RISC-Vs (SLOT 0 / 1, each its own scratch CB): 2 reads per page (or zeros), one barrier, BATCH 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 ROWB = get_compile_time_arg_val(4); // C * 2 bytes
constexpr uint32_t NPAGES = get_compile_time_arg_val(5);
constexpr uint32_t NC = get_compile_time_arg_val(6);
constexpr uint32_t GX = get_compile_time_arg_val(7);
constexpr uint32_t BATCH = 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 W1 = WV + 1;
static_assert(ROWB % MEM_ZEROS_SIZE == 0, "dcnpair_k: C * 2 a multiple of the zeros page");
static inline void zeros(uint32_t dst) {
const uint64_t z = get_noc_addr(MEM_ZEROS_BASE);
for (uint32_t o = 0; o < ROWB; o += MEM_ZEROS_SIZE) {
noc_async_read(z, dst + 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 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), ROWB);
const auto oa = TensorAccessor(o_args, get_common_arg_val<uint32_t>(1), 2 * ROWB);
cb_reserve_back(cb_s, 1);
const uint32_t buf = get_write_ptr(cb_s);
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 row = p / W1; // (n, y)
const uint32_t i = p - row * W1;
const uint32_t dst = buf + (p - b0) * 2 * ROWB;
const uint32_t q = row * WV + i; // input page of pixel x = i
if (i == 0) {
zeros(dst);
} else {
noc_async_read(ia.get_noc_addr(q - 1), dst, ROWB);
}
if (i == WV) {
zeros(dst + ROWB);
} else {
noc_async_read(ia.get_noc_addr(q), dst + ROWB, ROWB);
}
}
noc_async_read_barrier();
for (uint32_t p = b0; p < b1; ++p) {
noc_async_write(buf + (p - b0) * 2 * ROWB, oa.get_noc_addr(p), 2 * ROWB);
}
noc_async_write_barrier();
}
}
|