File size: 975 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
// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hattn writer (BRISC), round 6: row 0 of the 2 context tiles of head h -> tiles 2h, 2h+1 of the row-form context
// [1, 1, 32, H*64] (bf16; the head concat is this placement).  runtime: out_addr, h
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"

void kernel_main() {
    constexpr auto o_args = TensorAccessorArgs<0>();
    constexpr uint32_t c_out = 16, page = 2048, FB = 512;
    const uint32_t o_addr = get_arg_val<uint32_t>(0), h = get_arg_val<uint32_t>(1);
    const auto out = TensorAccessor(o_args, o_addr, page);
    cb_wait_front(c_out, 2);
    const uint32_t src = get_read_ptr(c_out);
    for (uint32_t d = 0; d < 2; ++d) {
        const uint64_t dst = out.get_noc_addr(2 * h + d);
        noc_async_write(src + d * page, dst, 32);
        noc_async_write(src + d * page + FB, dst + FB, 32);
    }
    noc_async_write_barrier();
    cb_pop_front(c_out, 2);
}