File size: 3,503 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
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
// 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
}