File size: 2,246 Bytes
58ba6df
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
// 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
}