File size: 2,637 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
// 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();
}