code/hamer/kernels/hgemv/writer.cpp · changh95/hamer-p150 at main

hamer-p150 / code /hamer /kernels /hgemv /writer.cpp

Optimized build (2026-10-03): model call 3.49 ms, device 3.36 ms

58ba6df verified about 1 month ago

4.21 kB

// 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);
}

	// 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);
	}