File size: 5,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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | // SPDX-License-Identifier: Apache-2.0
// K-IN rows (bevformer-p150 OPT_PLAN R2-2): one camera of the uploaded bf16 image rows -> the height-sharded L1
// input of the 7x7 / stride-2 stem conv (ttnn's L1 conv path), replacing the device reshape of the image rows to
// 6-byte pixel pages and the stem's 10 DRAM PaddedSlices.
//
// in : [1, 1, CAMS * H, W * C_IN] bf16 ROW_MAJOR, DRAM interleaved (one page per image row, ROW_BYTES = W * C_IN * 2)
// out: [1, 1, H * W, C_IN] bf16 ROW_MAJOR, L1 HEIGHT_SHARDED, shard [S, C_PAD] (C_PAD = 8: 16-byte rows), camera
// CAM's image rows ROW0 = CAM * H .. ROW0 + H - 1
//
// Core k (row-major over a GX-wide grid at the origin; ROW_MAJOR shard orientation) owns pixels
// [k * S, min((k + 1) * S, N)), split between SLOTS data-movement RISC-Vs (BRISC / NCRISC, each its own scratch CB).
// Each reads the image rows that hold its pixels into its scratch CB (whole pages), then copies each pixel's C_IN
// bf16 values and zeros in channels C_IN .. C_PAD - 1 straight into its own shard (local L1 stores; no NoC write).
// The copy is exact (bits moved, no conversion). Data-movement only, no CB push / pop, no semaphores, no multicast;
// every read waited on by noc_async_read_barrier before use. Runtime args: COMMON only [in_addr, out_addr] (P3).
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t N = get_compile_time_arg_val(0); // pixels of one camera (H * W)
constexpr uint32_t S = get_compile_time_arg_val(1); // shard rows (pixels per core)
constexpr uint32_t W = get_compile_time_arg_val(2); // pixels per image row
constexpr uint32_t C_IN = get_compile_time_arg_val(3); // channels (3)
constexpr uint32_t C_PAD = get_compile_time_arg_val(4); // shard width in elements (8: 16 B rows)
constexpr uint32_t NUM_CORES = get_compile_time_arg_val(5); // cores holding shards
constexpr uint32_t GX = get_compile_time_arg_val(6); // grid width (cores)
constexpr uint32_t CB_SCRATCH = get_compile_time_arg_val(7); // scratch CB: MAX_ROWS image rows
constexpr uint32_t MAX_ROWS = get_compile_time_arg_val(8); // rows a slot's pixels can touch
constexpr uint32_t SLOTS = get_compile_time_arg_val(9); // RISC-Vs sharing a core's shard (1 or 2)
constexpr uint32_t SLOT = get_compile_time_arg_val(10); // this kernel's slot: pixels [SLOT * S / SLOTS, ...)
constexpr uint32_t ROW0 = get_compile_time_arg_val(11); // first image row of the camera in the input
constexpr auto in_args = TensorAccessorArgs<12>();
constexpr uint32_t ROW_BYTES = W * C_IN * 2;
static_assert(C_IN == 3 && C_PAD == 8, "the copy loop handles 3 channels into 16-byte rows");
static_assert(ROW_BYTES % 64 == 0, "image rows must be 64-byte multiples (DRAM read alignment)");
void kernel_main() {
const uint32_t in_addr = get_common_arg_val<uint32_t>(0);
const uint32_t out_addr = get_common_arg_val<uint32_t>(1);
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NUM_CORES) {
return;
}
constexpr uint32_t PART = (S + SLOTS - 1) / SLOTS;
const uint32_t shard0 = core * S;
const uint32_t p0 = shard0 + SLOT * PART;
const uint32_t end = shard0 + ((SLOT + 1) * PART < S ? (SLOT + 1) * PART : S);
const uint32_t p1 = end < N ? end : N;
if (p0 >= p1) {
return;
}
const uint32_t row0 = p0 / W;
const uint32_t row1 = (p1 - 1) / W + 1; // exclusive
const auto ia = TensorAccessor(in_args, in_addr, ROW_BYTES);
const uint32_t scratch = (get_write_ptr(CB_SCRATCH) + 63u) & ~63u;
for (uint32_t r = row0; r < row1; ++r) {
noc_async_read(ia.get_noc_addr(ROW0 + r), scratch + (r - row0) * ROW_BYTES, ROW_BYTES);
}
noc_async_read_barrier();
asm volatile("" ::: "memory"); // no load of the rows may move above the barrier
// pixel p's 3 bf16 = 6 bytes at scratch + 6 (p - row0 W); an even pixel starts on a 4-byte boundary
const uint32_t src0 = scratch + (p0 - row0 * W) * C_IN * 2;
uint32_t* dst = reinterpret_cast<uint32_t*>(out_addr + (p0 - shard0) * C_PAD * 2);
uint32_t n = p1 - p0;
const uint16_t* sh = reinterpret_cast<const uint16_t*>(src0);
if (n && (reinterpret_cast<uint32_t>(sh) & 3u)) { // odd start pixel: one single
dst[0] = static_cast<uint32_t>(sh[0]) | (static_cast<uint32_t>(sh[1]) << 16);
dst[1] = sh[2];
dst[2] = 0;
dst[3] = 0;
sh += 3;
dst += 4;
--n;
}
const uint32_t* sw = reinterpret_cast<const uint32_t*>(sh);
for (; n >= 2; n -= 2) { // 2 pixels = 3 words: a0 b0 | c0 a1 | b1 c1
const uint32_t w0 = sw[0], w1 = sw[1], w2 = sw[2];
dst[0] = w0;
dst[1] = w1 & 0xFFFFu;
dst[2] = 0;
dst[3] = 0;
dst[4] = (w1 >> 16) | (w2 << 16);
dst[5] = w2 >> 16;
dst[6] = 0;
dst[7] = 0;
sw += 3;
dst += 8;
}
if (n) {
sh = reinterpret_cast<const uint16_t*>(sw);
dst[0] = static_cast<uint32_t>(sh[0]) | (static_cast<uint32_t>(sh[1]) << 16);
dst[1] = sh[2];
dst[2] = 0;
dst[3] = 0;
}
asm volatile("" ::: "memory");
}
|