code/hamer/kernels/eltwise/reader.cpp · changh95/hamer-p150 at main

hamer-p150 / code /hamer /kernels /eltwise /reader.cpp

Optimized build (2026-10-03): model call 3.49 ms, device 3.36 ms

58ba6df verified about 1 month ago

2.25 kB

// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// Unary eltwise reader: pages [start, start+n) of the source (interleaved or sharded) into cb 0 in one
// batch (the CB holds the core's whole range).  compile-time: TensorAccessorArgs(src).
// runtime: src_addr, start, n [, npf, w_addr, npf x (bank, bank offset, local dst, bytes)]
//   npf > 0 (cross-op weight prefetch, tt_kernels.gelu(prefetch=...)): after its own tiles this core reads
//   npf contiguous DRAM blocks (the next cmm_linear's first weight blocks) into a persistent L1 buffer
//   (plain reads + one plain barrier before the kernel ends, so the data is in L1 when the next op starts)
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"

void kernel_main() {
    const uint32_t src_addr = get_arg_val<uint32_t>(0);
    const uint32_t start = get_arg_val<uint32_t>(1);
    const uint32_t n = get_arg_val<uint32_t>(2);
    constexpr auto src_args = TensorAccessorArgs<0>();
    constexpr uint32_t cb = 0;
    const uint32_t page = get_local_cb_interface(cb).fifo_page_size;
    const auto src = TensorAccessor(src_args, src_addr, page);
    // in up to 4 chunks so the compute starts on the first tiles while the rest are in flight
    cb_reserve_back(cb, n);
    uint32_t dst = get_write_ptr(cb);
    const uint32_t chunk = (n + 7) / 8;
    for (uint32_t j0 = 0; j0 < n; j0 += chunk) {
        const uint32_t m = (n - j0) < chunk ? (n - j0) : chunk;
        for (uint32_t j = 0; j < m; ++j) {
            noc_async_read(src.get_noc_addr(start + j0 + j), dst, page);
            dst += page;
        }
        noc_async_read_barrier();
        cb_push_back(cb, m);
    }
#ifdef ELT_PREFETCH
    const uint32_t npf = get_arg_val<uint32_t>(3);
    if (npf) {
        const uint32_t w_addr = get_arg_val<uint32_t>(4);
        for (uint32_t i = 0; i < npf; ++i) {
            const uint32_t bank = get_arg_val<uint32_t>(5 + 4 * i), off = get_arg_val<uint32_t>(6 + 4 * i);
            const uint32_t ldst = get_arg_val<uint32_t>(7 + 4 * i), bytes = get_arg_val<uint32_t>(8 + 4 * i);
            noc_async_read(get_noc_addr_from_bank_id<true>(bank, w_addr + off), ldst, bytes);
        }
        noc_async_read_barrier();
    }
#endif
}

	// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
	// SPDX-License-Identifier: Apache-2.0
	// Unary eltwise reader: pages [start, start+n) of the source (interleaved or sharded) into cb 0 in one
	// batch (the CB holds the core's whole range). compile-time: TensorAccessorArgs(src).
	// runtime: src_addr, start, n [, npf, w_addr, npf x (bank, bank offset, local dst, bytes)]
	// npf > 0 (cross-op weight prefetch, tt_kernels.gelu(prefetch=...)): after its own tiles this core reads
	// npf contiguous DRAM blocks (the next cmm_linear's first weight blocks) into a persistent L1 buffer
	// (plain reads + one plain barrier before the kernel ends, so the data is in L1 when the next op starts)
	#include <stdint.h>
	#include "api/dataflow/dataflow_api.h"

	void kernel_main() {
	const uint32_t src_addr = get_arg_val<uint32_t>(0);
	const uint32_t start = get_arg_val<uint32_t>(1);
	const uint32_t n = get_arg_val<uint32_t>(2);
	constexpr auto src_args = TensorAccessorArgs<0>();
	constexpr uint32_t cb = 0;
	const uint32_t page = get_local_cb_interface(cb).fifo_page_size;
	const auto src = TensorAccessor(src_args, src_addr, page);
	// in up to 4 chunks so the compute starts on the first tiles while the rest are in flight
	cb_reserve_back(cb, n);
	uint32_t dst = get_write_ptr(cb);
	const uint32_t chunk = (n + 7) / 8;
	for (uint32_t j0 = 0; j0 < n; j0 += chunk) {
	const uint32_t m = (n - j0) < chunk ? (n - j0) : chunk;
	for (uint32_t j = 0; j < m; ++j) {
	noc_async_read(src.get_noc_addr(start + j0 + j), dst, page);
	dst += page;
	}
	noc_async_read_barrier();
	cb_push_back(cb, m);
	}
	#ifdef ELT_PREFETCH
	const uint32_t npf = get_arg_val<uint32_t>(3);
	if (npf) {
	const uint32_t w_addr = get_arg_val<uint32_t>(4);
	for (uint32_t i = 0; i < npf; ++i) {
	const uint32_t bank = get_arg_val<uint32_t>(5 + 4 * i), off = get_arg_val<uint32_t>(6 + 4 * i);
	const uint32_t ldst = get_arg_val<uint32_t>(7 + 4 * i), bytes = get_arg_val<uint32_t>(8 + 4 * i);
	noc_async_read(get_noc_addr_from_bank_id<true>(bank, w_addr + off), ldst, bytes);
	}
	noc_async_read_barrier();
	}
	#endif
	}