Livt.ML

src/norm/RMSNorm.lvt

namespace Livt.ML.Norm

using Livt.Math.Arithmetic

using Livt.Math.Sqrt

/**
 * Runtime division is scheduled; completion latency increases with work.
 * Serialize operations and keep inputs stable until the call or handshake completes.
 *
 * Computes integer RMS normalization for a fixed 8-element vector.
 */
component RMSNorm
{
	// Owned divider; calls complete before this component advances.
	divider: IntegerDivision<32>

	// Number of integer elements processed by Compute.
	public const VECTOR_SIZE: int = 8

	// Owned arithmetic worker used by serialized operations.
	sqrtNr: SqrtNewtonRaphson

	// Latest computed result, retained for the caller.
	output: int[8]

	/**
	 * Creates the integer square-root helper used for RMS scaling.
	 */
	new()
	{
		divider = new IntegerDivision<32>()
		this.sqrtNr = new SqrtNewtonRaphson()
	}

	/**
	 * Normalizes input by the integer root mean square and stores the output.
	 */
	public fn Compute(input: int[8])
	{
		var sumSq = 0

		for (var i = 0; i < VECTOR_SIZE; i++) {
			sumSq = sumSq + input[i] * input[i]
		}

		var meanSq = sumSq / VECTOR_SIZE
		var rms = this.sqrtNr.Compute(meanSq)

		if (rms == 0) {
			rms = 1
		}

		for (var i = 0; i < VECTOR_SIZE; i++) {
			this.output[i] = divider.Divide(input[i], rms)
		}
	}

	/**
	 * Returns one normalized output element, or zero when index is out of range.
	 */
	public fn GetOutput(index: int) int
	{
		if (index < 0) {
			return 0
		}

		if (index >= VECTOR_SIZE) {
			return 0
		}

		return this.output[index]
	}
}