File size: 4,206 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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | // SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hgemv writer (BRISC, NOC_1), round 6. HG_LN: the normalized dense tile (row t = elements 32t..32t+31, bf16) is
// spread into the row-form matmul operand (tile t row 0), by local NoC copies. Then row 0 of the output tile goes
// to tile n of the row-form output tensor.
// HG_DOUT=D: row 0 is also written as row n of D replica tiles of a dense output vector [1, 1, 32*D, 32].
// compile-time: KT, OES (output element bytes), TensorAccessorArgs(out)[, TensorAccessorArgs(dense out)];
// runtime: out_addr, n[, dense_addr]
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"
void kernel_main() {
constexpr uint32_t KT = get_compile_time_arg_val(0);
constexpr uint32_t OES = get_compile_time_arg_val(1);
constexpr auto o_args = TensorAccessorArgs<2>();
constexpr uint32_t c_hd = 6, c_hrow = 5, c_out = 16;
const uint32_t o_addr = get_arg_val<uint32_t>(0), n = get_arg_val<uint32_t>(1);
#if (defined(HG_LN) || defined(HG_DIN)) && !defined(HG_TINY)
{
// dense tile (row t = elements 32t..32t+31) -> row-form operand (tile t row 0): the LN output (bf16), or the
// dense input itself without LN (element size HES)
#ifdef HG_LN
constexpr uint32_t HES = 2, c_src = c_hd;
#else
constexpr uint32_t HES = HG_IES, c_src = 0;
#endif
constexpr uint32_t HFB = 256 * HES, HPAGE = 1024 * HES;
cb_wait_front(c_src, 1);
cb_reserve_back(c_hrow, KT);
const uint32_t hd = get_read_ptr(c_src), hr = get_write_ptr(c_hrow);
for (uint32_t t = 0; t < KT; ++t) {
const uint32_t s = hd + ((t / 16) * 2) * HFB + (t % 16) * 16 * HES;
noc_async_read(get_noc_addr(s), hr + t * HPAGE, 16 * HES);
noc_async_read(get_noc_addr(s + HFB), hr + t * HPAGE + HFB, 16 * HES);
}
noc_async_read_barrier();
cb_push_back(c_hrow, KT);
cb_pop_front(c_src, 1);
}
#endif
constexpr uint32_t OFB = 256 * OES;
const auto out = TensorAccessor(o_args, o_addr, 1024 * OES);
cb_wait_front(c_out, 1);
const uint32_t src = get_read_ptr(c_out);
const uint64_t dst = out.get_noc_addr(n);
#ifdef HG_TINY
// round 10: c_out is a [1, 32] tiny tile: its two face rows are contiguous
constexpr uint32_t SRC1 = 16 * OES;
#else
constexpr uint32_t SRC1 = OFB;
#endif
#ifdef HG_REP
// replicate row 0 into all 32 rows (consumers that process whole tiles, e.g. SDPA on the query)
for (uint32_t r = 0; r < 32; ++r) {
const uint32_t off = ((r / 16) * 2) * OFB + (r % 16) * 16 * OES;
noc_async_write(src, dst + off, 16 * OES);
noc_async_write(src + SRC1, dst + off + OFB, 16 * OES);
}
#else
noc_async_write(src, dst, 16 * OES);
noc_async_write(src + SRC1, dst + OFB, 16 * OES);
#endif
#ifdef HG_DOUT
{
// dense replicas: row n of each of the HG_DOUT replica tiles of the dense output vector
constexpr auto d_args = TensorAccessorArgs<o_args.next_compile_time_args_offset()>();
const auto dout = TensorAccessor(d_args, get_arg_val<uint32_t>(2), 1024 * OES);
#ifdef HG_TINY
// round 10: row-major dense replicas: row n = 32 contiguous elements = this core's tiny output tile
for (uint32_t c = 0; c < HG_DOUT; ++c) {
noc_async_write(src, dout.get_noc_addr(c) + n * 32 * OES, 32 * OES);
}
#else
const uint32_t off = ((n / 16) * 2) * OFB + (n % 16) * 16 * OES;
for (uint32_t c = 0; c < HG_DOUT; ++c) {
const uint64_t d = dout.get_noc_addr(c) + off;
noc_async_write(src, d, 16 * OES);
noc_async_write(src + OFB, d + OFB, 16 * OES);
}
#endif
}
#endif
noc_async_write_barrier();
#ifdef HG_SIGNAL
// round 15: this program also runs the head attention (kernels/hattn2, HA2_WAIT) on other cores: the output column
// (a query column) is in L1 (acked above) -> increment the attention core's semaphore 0
noc_semaphore_inc(get_noc_addr(get_arg_val<uint32_t>(HG_SIGNAL), get_arg_val<uint32_t>(HG_SIGNAL + 1), get_semaphore(0)), 1);
noc_async_atomic_barrier();
#endif
cb_pop_front(c_out, 1);
}
|