Livt.ML

src/activation/SoftmaxApprox.lvt

namespace Livt.ML.Activation

using Livt.Math.Arithmetic

/**
 * Runtime division is scheduled; completion latency increases with work.
 * Serialize operations and keep inputs stable until the call or handshake completes.
 *
 * Computes an 8-element integer softmax approximation with scaled outputs.
 */
component SoftmaxApprox
{
	// Owned divider; calls complete before this component advances.
	divider: IntegerDivision<32>

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

	// Number of entries in the clamped exponential lookup table.
	public const LUT_SIZE: int = 16

	// Fixed-point probability scale used for stored outputs.
	public const SCALE_FACTOR: int = 256

	// Owned approximation table in the component numeric scale.
	expLut: int[16]

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

	/**
	 * Initializes the 16-entry exponential approximation table.
	 */
	new()
	{
		divider = new IntegerDivision<32>()
		expLut[0] = 1
		expLut[1] = 2
		expLut[2] = 7
		expLut[3] = 20
		expLut[4] = 54
		expLut[5] = 148
		expLut[6] = 403
		expLut[7] = 1096
		expLut[8] = 2980
		expLut[9] = 8103
		expLut[10] = 22026
		expLut[11] = 59874
		expLut[12] = 162754
		expLut[13] = 442413
		expLut[14] = 1202604
		expLut[15] = 3269446
	}

	/**
	 * Computes scaled softmax probabilities after clamping logits to [0, 15].
	 */
	public fn Compute(input: int[8])
	{
		this.ComputeCount(input, VECTOR_SIZE)
	}

	/**
	 * Computes softmax over the first count values and clears inactive outputs.
	 */
	public fn ComputeCount(input: int[8], count: int)
	{
		var activeCount = count

		if (activeCount < 0) {
			activeCount = 0
		}

		if (activeCount > VECTOR_SIZE) {
			activeCount = VECTOR_SIZE
		}

		for (var i = 0; i < VECTOR_SIZE; i++) {
			output[i] = 0
		}

		if (activeCount == 0) {
			return
		}

		var sumExp = 0

		for (var i = 0; i < activeCount; i++) {
			var clampedA = input[i]

			if (clampedA < 0) {
				clampedA = 0
			}

			if (clampedA > 15) {
				clampedA = 15
			}

			sumExp = sumExp + expLut[clampedA]
		}

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

		for (var i = 0; i < activeCount; i++) {
			var clampedB = input[i]

			if (clampedB < 0) {
				clampedB = 0
			}

			if (clampedB > 15) {
				clampedB = 15
			}

			var ev = expLut[clampedB]
			output[i] = divider.Divide(ev * SCALE_FACTOR, sumExp)
		}
	}

	/**
	 * Returns one scaled probability, 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 output[index]
	}
}