82 lines
3.4 KiB
C
82 lines
3.4 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/**
|
|
* @file ft_tt_op_rms_norm.h
|
|
* @brief Host driver for the RMS-norm triad — the fc_rms_norm equivalent (spec §5.1).
|
|
*
|
|
* y[r][c] = x[r][c] * gamma[c] * rsqrt( mean_c(x[r][c]^2) + eps )
|
|
*
|
|
* Two things the host owns that are not obvious from the kernel:
|
|
*
|
|
* 1. **The scaler and eps tiles.** The row sum comes from reduce_tile<SUM, REDUCE_ROW>,
|
|
* which multiplies by a scaler tile; the host builds one holding 1/W. eps is then one
|
|
* more accumulated reduce against a second constant tile whose every element is
|
|
* eps*W/32, because each row of it contributes 32 * (eps*W/32) * (1/W) = eps. That
|
|
* buys eps for one extra reduce_tile call instead of an add_tiles_bcast_scalar, an
|
|
* extra CB round trip and a pack (see the kernel's header note 1).
|
|
*
|
|
* ft_tt_op_rms_norm_build_constants() makes both, in the CB's own format, so nothing
|
|
* downstream has to know the trick.
|
|
*
|
|
* 2. **gamma is expanded to full tiles by the loader**, not broadcast in the kernel —
|
|
* ft_tt_fttw_upload_broadcast_rows(). The 32x inflation is 96 KiB for a 1536-wide
|
|
* model and it removes a broadcast reconfiguration from the inner loop.
|
|
*
|
|
* @status PCC and shape coverage in docs/tt_bringup_log.md §14.
|
|
* @issues None
|
|
* @todo None
|
|
*/
|
|
|
|
#ifndef FT_TT_OP_RMS_NORM_H
|
|
#define FT_TT_OP_RMS_NORM_H
|
|
|
|
#include <stddef.h>
|
|
#include <stdint.h>
|
|
|
|
#include "ft_tt_kernel_blob.h"
|
|
#include "ft_tt_launch.h"
|
|
#include "ft_tt_op_matmul.h" /* ft_tt_op_operand */
|
|
|
|
/** One RMS norm over `row_blocks * 32` rows of `tile_cols * 32` columns. */
|
|
typedef struct ft_tt_op_rms_norm_request {
|
|
uint32_t row_blocks; /**< rows / 32 */
|
|
uint32_t tile_cols; /**< W / 32 */
|
|
ft_tt_op_operand activation;
|
|
ft_tt_op_operand gamma;
|
|
/** Two adjacent tiles: the 1/W scaler, then the eps tile. */
|
|
ft_tt_op_operand constants;
|
|
ft_tt_op_operand output;
|
|
} ft_tt_op_rms_norm_request;
|
|
|
|
/**
|
|
* Build the two constant tiles into `out`, which must hold 2 * 2048 bytes (two Float16_b
|
|
* tiles). Upload them contiguously and point `constants` at the first.
|
|
*
|
|
* @param width W, the reduced dimension in ELEMENTS — not tiles. The eps tile's value
|
|
* depends on it, which is why it is not a fixed constant
|
|
* @param eps the epsilon the CPU reference uses
|
|
* @return FT_TT_OK, or FT_TT_EINVAL for a short buffer or a zero width
|
|
*
|
|
* Thread-safe: yes Reentrant: yes Signal-safe: yes
|
|
*/
|
|
int ft_tt_op_rms_norm_build_constants(uint32_t width, float eps, void *out, size_t out_bytes);
|
|
|
|
/**
|
|
* Run one RMS norm on a core whose firmware is already booted.
|
|
*
|
|
* L1 cost is `(3 * tile_cols + 5)` tiles of circular buffer: the activation, gamma and the
|
|
* x*x intermediate are all `tile_cols` pages deep, because the kernel needs a whole row
|
|
* block resident to reduce over it and then to scale it. At W=1536 that is ~300 KiB of
|
|
* L1's 1.5 MiB.
|
|
*
|
|
* @param timeout_ms bounds the wait; there is no untimed path (spec §7.3)
|
|
* @return FT_TT_OK; FT_TT_EIO on timeout with the core dumped; FT_TT_EINVAL
|
|
*
|
|
* Thread-safe: no Reentrant: no Signal-safe: no
|
|
*/
|
|
int ft_tt_op_rms_norm_run(const ft_tt_noc_access *noc, uint16_t x, uint16_t y,
|
|
const ft_tt_kernel_program *program,
|
|
const ft_tt_op_rms_norm_request *request, uint32_t timeout_ms);
|
|
|
|
#endif /* FT_TT_OP_RMS_NORM_H */
|