hamer-p150 / code /hamer /kernels /cmm /help_writer.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
2.64 kB
// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// cmm GELU helper (round 8), BRISC: each GELU'd tile (CB 21) goes straight into its owner's activation shard (the cmm
// owner layout: core (K block, token tile row), tile kk at offset kk * page), then the owner's GELU-done semaphore is
// incremented. The write is ACKED (barrier) before the increment: atomics and writes are not ordered on the NoC.
// The owner multicasts its block once all its tiles are counted (kernels/cmm/mover1.cpp, CMM_GELU_HELP).
// Round 11 (HELP_WR_BATCH, default): all tiles already in CB 21 are written, then ONE barrier, then one increment per
// tile (instead of write + barrier + increment per tile).
// compile-time: semaphore id. runtime: act_addr, n, then n x (owner x, owner y, byte offset) (fixed length).
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"
void kernel_main() {
constexpr uint32_t sem_id = get_compile_time_arg_val(0);
const uint32_t act_addr = get_arg_val<uint32_t>(0);
const uint32_t n = get_arg_val<uint32_t>(1);
constexpr uint32_t cb = 21;
const uint32_t page = get_local_cb_interface(cb).fifo_page_size;
const uint32_t sem_addr = get_semaphore(sem_id);
#ifdef HELP_WR_BATCH
uint32_t i = 0;
while (i < n) {
cb_wait_front(cb, 1);
uint32_t k = 1;
while (i + k < n && k < HELP_WR_BATCH && cb_pages_available_at_front(cb, k + 1)) {
++k;
}
uint32_t src = get_read_ptr(cb);
for (uint32_t j = 0; j < k; ++j) {
const uint32_t a = 2 + 3 * (i + j);
noc_async_write(src, get_noc_addr(get_arg_val<uint32_t>(a), get_arg_val<uint32_t>(a + 1),
act_addr + get_arg_val<uint32_t>(a + 2)), page);
src += page;
}
noc_async_write_barrier();
for (uint32_t j = 0; j < k; ++j) {
const uint32_t a = 2 + 3 * (i + j);
noc_semaphore_inc(get_noc_addr(get_arg_val<uint32_t>(a), get_arg_val<uint32_t>(a + 1), sem_addr), 1);
}
cb_pop_front(cb, k);
i += k;
}
#else
for (uint32_t i = 0; i < n; ++i) {
const uint32_t ox = get_arg_val<uint32_t>(2 + 3 * i), oy = get_arg_val<uint32_t>(3 + 3 * i);
const uint32_t off = get_arg_val<uint32_t>(4 + 3 * i);
cb_wait_front(cb, 1);
noc_async_write(get_read_ptr(cb), get_noc_addr(ox, oy, act_addr + off), page);
noc_async_write_barrier();
noc_semaphore_inc(get_noc_addr(ox, oy, sem_addr), 1);
cb_pop_front(cb, 1);
}
#endif
noc_async_atomic_barrier();
}