hamer-p150 / code /hamer /kernels /cmm /help_reader.cpp
changh95's picture
Optimized build (2026-10-03): model call 3.49 ms, device 3.36 ms
58ba6df verified
Raw History Blame Contribute Delete
3.5 kB
// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// cmm GELU helper (round 8), NCRISC: the pre-GELU fc1-output pages assigned to this helper core -> CB 20.
// Round 11 (HELP_RD_ALL, default): every read is issued up front, tagged with a transaction id (1 + i % 15), and page i
// is pushed after the trid-i read barrier (the library DRAM-sharded reader pattern; the plain-counter early signal of
// round 10 was the unsafe one), instead of one read + full barrier per page (~0.5 us of latency per page).
// compile-time: TensorAccessorArgs(pre). runtime: pre_addr, n, then n page indices (fixed length CMM_HELP_MAXT).
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"
void kernel_main() {
const uint32_t pre_addr = get_arg_val<uint32_t>(0);
const uint32_t n = get_arg_val<uint32_t>(1);
constexpr auto pre_args = TensorAccessorArgs<0>();
constexpr uint32_t cb = 20;
const uint32_t page = get_local_cb_interface(cb).fifo_page_size;
const auto pre = TensorAccessor(pre_args, pre_addr, page);
#ifdef HELP_RD_ALL
// CB 20 holds CMM_HELP_MAXT >= n pages: reserve all, issue all
cb_reserve_back(cb, n);
uint32_t dst = get_write_ptr(cb);
for (uint32_t i = 0; i < n; ++i) {
noc_async_read_set_trid(1 + i % 15);
noc_async_read(pre.get_noc_addr(get_arg_val<uint32_t>(2 + i)), dst, page);
dst += page;
}
for (uint32_t i = 0; i < n; ++i) {
noc_async_read_barrier_with_trid(1 + i % 15);
cb_push_back(cb, 1);
}
noc_async_read_set_trid(0);
noc_async_read_barrier();
#elif defined(HELP_RD_WIN)
// a window of HELP_RD_WIN reads in flight (trid 1 + i % W): page i is pushed after its trid barrier
constexpr uint32_t W = HELP_RD_WIN;
cb_reserve_back(cb, n);
const uint32_t base = get_write_ptr(cb);
uint32_t issued = 0;
for (uint32_t i = 0; i < n; ++i) {
while (issued < n && issued < i + W) {
noc_async_read_set_trid(1 + issued % W);
noc_async_read(pre.get_noc_addr(get_arg_val<uint32_t>(2 + issued)), base + issued * page, page);
++issued;
}
noc_async_read_barrier_with_trid(1 + i % W);
cb_push_back(cb, 1);
}
noc_async_read_set_trid(0);
noc_async_read_barrier();
#else
for (uint32_t i = 0; i < n; ++i) {
cb_reserve_back(cb, 1);
noc_async_read(pre.get_noc_addr(get_arg_val<uint32_t>(2 + i)), get_write_ptr(cb), page);
noc_async_read_barrier();
cb_push_back(cb, 1);
}
#endif
#ifdef HELP_PF
// round 11: after its GELU inputs, this core reads the first K blocks of a LATER xmm's packed weight (its own
// N-group: it is that xmm's weight sender) from DRAM into its persistent L1 prefetch buffer (tt_kernels.XmmPrefetch).
// runtime args 2 + HELP_MAXT ..: npf, w_addr, npf x (bank, bank offset, local dst, bytes)
constexpr uint32_t P = 2 + HELP_MAXT;
const uint32_t npf = get_arg_val<uint32_t>(P);
if (npf) {
const uint32_t w_addr = get_arg_val<uint32_t>(P + 1);
for (uint32_t i = 0; i < npf; ++i) {
const uint32_t bank = get_arg_val<uint32_t>(P + 2 + 4 * i), off = get_arg_val<uint32_t>(P + 3 + 4 * i);
const uint32_t ldst = get_arg_val<uint32_t>(P + 4 + 4 * i), bytes = get_arg_val<uint32_t>(P + 5 + 4 * i);
noc_async_read(get_noc_addr_from_bank_id<true>(bank, w_addr + off), ldst, bytes);
}
noc_async_read_barrier();
}
#endif
}