Livt.ML

src/linear/QuantizedMatrix.lvt

namespace Livt.ML.Linear

using Livt.Math.Arithmetic
using Livt.ML.Numeric
using Livt.ML.Storage
using Livt.IO

/**
 * Row-major packed signed INT4 projection with flattened scale groups and injected storage.
 */
component QuantizedMatrix<ROWS: int, COLUMNS: int, GROUP_SIZE: int,
	BITS: int, FRACTION: int, SCALE_FRACTION: int, WEIGHTS: IQuantizedWeights, INPUT: IRam<int>>
{
	// Shared worker for non-power-of-two scale groups.
	divider: UnsignedDivision<31>

	static assert ROWS > 0 && ROWS <= 65536, "invalid row count"

	static assert COLUMNS > 0 && COLUMNS <= 65536, "invalid column count"

	static assert ROWS <= 2147483647 / COLUMNS, "matrix index overflows int"

	static assert GROUP_SIZE > 0 && GROUP_SIZE <= 65536, "invalid group size"

	// Borrowed packed coefficients and activation memory; keep stable during Compute.
	weights: WEIGHTS
	input: INPUT

	// Owned arithmetic worker used by serialized operations.
	arithmetic: FixedPointMath<BITS, FRACTION, SCALE_FRACTION>

	// Owned bounded dot-product reducer.
	accumulator: FixedPointAccumulator<COLUMNS, BITS, FRACTION>

	// Validity of the latest output-row computation.
	failed: bool

	/**
	 * Borrows packed weights and owns fixed-point reconstruction and accumulation helpers.
	 * Keep weights and scales stable throughout each computation.
	 */
	new(weights: WEIGHTS, input: INPUT)
	{
		divider = new UnsignedDivision<31>()
		this.weights = weights
		this.input = input
		arithmetic = new FixedPointMath<BITS, FRACTION, SCALE_FRACTION>()
		accumulator = new FixedPointAccumulator<COLUMNS, BITS, FRACTION>()
	}

	/**
	 * Compute one output row; storage remains external and no weight ROM is elaborated.
	 */
	public fn Compute(row: int, inputBase: int) int
	{
		failed = row < 0 || row >= ROWS || inputBase < 0 || inputBase > 2147483647 - COLUMNS

		if (failed) {
			return 0
		}

		if (weights.WeightCount() < ROWS * COLUMNS ||
			weights.ScaleCount() < ((ROWS * COLUMNS - 1) / GROUP_SIZE) + 1 ||
			!input.IsValidAddress(inputBase) ||
			!input.IsValidAddress(inputBase + COLUMNS - 1)) {
			failed = true
			return 0
		}

		accumulator.Begin()
		var previousGroup = -1
		var scale: logic[32] = 0

		for (var column = 0; column < COLUMNS; column++) {
			var index = row * COLUMNS + column
			var group = this.ScaleGroup(index)

			if (group != previousGroup) {
				scale = weights.ReadScale(group)
				previousGroup = group
			}

			var packed = weights.ReadPacked(index / 2)
			var code = arithmetic.SignedNibble(packed, index % 2 != 0)
			var weight = arithmetic.Weight(code, scale)

			if (arithmetic.Failed()) {
				failed = true
				return 0
			}

			accumulator.Accumulate(input.Read(inputBase + column), weight)
		}

		var result = accumulator.Finish()
		failed = accumulator.Failed()
		return result
	}

	/**
	 * Reports whether the latest computation rejected its shape, storage, weights, or
	 * arithmetic inputs; a failed result must be discarded.
	 */
	public fn Failed() bool
	{
		return failed
	}

	/**
	 * Maps a nonnegative element index to its scale group. Power-of-two
	 * specialization retains cheap constant arithmetic; other sizes are scheduled.
	 */
	fn ScaleGroup(index: int) int
	{
		static if ((GROUP_SIZE & (GROUP_SIZE - 1)) == 0) {
			return index / GROUP_SIZE
		} else {
			return divider.Divide(index as uint, GROUP_SIZE as uint) as int
		}
	}
}