File size: 1,063 Bytes
c699c4c | 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 | // SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
// SPDX-License-Identifier: Apache-2.0
//
// Score softmax core (SP_HEAD_SM=1), PROC 1: write each softmax tile row (3 tiles) to the s_sm tensor
// (TILE [N, 65], pages = tiles, row-major tile order). RT args: dst_addr, nrows, first tile row.
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"
void kernel_main() {
const uint32_t dst_addr = get_arg_val<uint32_t>(0);
const uint32_t nrows = get_arg_val<uint32_t>(1);
const uint32_t row0 = get_arg_val<uint32_t>(2);
constexpr uint32_t cb_out = get_compile_time_arg_val(0);
constexpr auto o_args = TensorAccessorArgs<1>();
const auto oacc = TensorAccessor(o_args, dst_addr, 2048);
for (uint32_t i = 0; i < nrows; ++i) {
cb_wait_front(cb_out, 3);
const uint32_t l1 = get_read_ptr(cb_out);
for (uint32_t j = 0; j < 3; ++j) {
noc_async_write(l1 + j * 2048, oacc.get_noc_addr((row0 + i) * 3 + j), 2048);
}
noc_async_write_barrier();
cb_pop_front(cb_out, 3);
}
}
|