Livt.ML

src/transformer/FixedTransformerKernels.lvt

namespace Livt.ML.Transformer

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

/**
 * Serial correctness implementation over external memory; one instance serves all layers.
 * Unscaled attention, affine RMSNorm, separately rounded GELU/gate and residuals.
 * Does not allocate model-sized on-chip arrays. Two scratch rows use borrowed scheduled RAM, never function-local arrays.
 */
component FixedTransformerKernels<MAX_TOKENS: int, MAX_WIDTH: int, BITS: int,
	FRACTION: int, PROB_FRACTION: int, EPSILON: int,
	WEIGHTS: IWeightTensor, TABLES: IFixedTransformerTables, MEMORY: IRam<int>,
	SCRATCH: IRam<int>> : IIncrementalTransformerKernels
{
	static assert MAX_TOKENS > 0 && MAX_TOKENS <= 1024, "invalid token capacity"

	static assert MAX_WIDTH > 0 && MAX_WIDTH <= 65536, "invalid vector capacity"

	static assert BITS >= 2 && BITS <= 30, "invalid activation width"

	static assert FRACTION >= 0 && FRACTION < BITS, "invalid activation fraction"

	static assert PROB_FRACTION >= 0 && PROB_FRACTION < BITS - 1, "invalid probability width"

	static assert EPSILON > 0, "epsilon must be positive"

	static assert 2 * BITS + Bits.RequiredFor(MAX_WIDTH) <= 64, "unsafe accumulation"

	static assert 2 * BITS + Bits.RequiredFor(MAX_TOKENS) <= 64, "unsafe context accumulation"

	// Symmetric raw activation magnitude allowed by the configured bit width.
	public const LIMIT: int = (1 << (BITS - 1)) - 1

	// Borrowed model providers and activation memory; callers serialize access.
	weights: WEIGHTS
	tables: TABLES
	memory: MEMORY

	// Owned wide arithmetic, shape division, and streaming argmax workers.
	math: WideArithmetic<BITS>
	divider: UnsignedDivision<31>
	argmax: StreamingArgMax

	// Exclusive scratch: scores at 0, probabilities at MAX_TOKENS.
	scratch: SCRATCH

	/**
	 * Borrows providers, activation memory and exclusive scratch for the instance lifetime.
	 * Scratch must expose contiguous addresses 0..2*MAX_TOKENS-1 and must not alias activation
	 * memory. Serialize every kernel operation and all scratch access. Reads/writes complete
	 * through the scheduled IRam contract; storage latency changes cycles, not arithmetic.
	 * Reset caller and RAM transaction state together. Cells need no reset: each attention
	 * query writes every consumed slot before reading it, including after failure/reset.
	 */
	new(weights: WEIGHTS, tables: TABLES, memory: MEMORY, scratch: SCRATCH)
	{
		this.weights = weights
		this.tables = tables
		this.memory = memory
		this.scratch = scratch
		math = new WideArithmetic<BITS>()
		divider = new UnsignedDivision<31>()
		argmax = new StreamingArgMax()
	}

	/**
	 * Checks the symmetric raw activation range -LIMIT..LIMIT.
	 */
	static fn ValueValid[](x: int) bool
	{
		return x >= -LIMIT && x <= LIMIT
	}

	/**
	 * Checks a nonempty word span for address overflow and valid endpoints. Assumes the memory
	 * exposes a contiguous valid range.
	 */
	fn Span(base: int, count: int) bool
	{
		if (base < 0 || count <= 0 || base > 2147483647 - count) {
			return false
		}

		return memory.IsValidAddress(base) && memory.IsValidAddress(base + count - 1)
	}

	/**
	 * Tests intersection of two half-open word spans whose endpoints were validated by Span.
	 */
	static fn Overlap[](a: int, na: int, b: int, nb: int) bool
	{
		return a < b + nb && b < a + na
	}

	/**
	 * Checks positive token and vector dimensions against MAX_TOKENS and MAX_WIDTH.
	 */
	static fn Shape[](length: int, width: int) bool
	{
		return length > 0 && length <= MAX_TOKENS && width > 0 && width <= MAX_WIDTH
	}

	/**
	 * Divides a signed wide intermediate by 2^fraction with ties to even, then clamps to the
	 * symmetric activation range. Requires a supported nonnegative shift.
	 */
	fn Round(value: logic[64], fraction: int) int
	{
		return math.Clamp(math.DivideEven(value, (1 << fraction) as logic[64]))
	}

	/**
	 * Copies length embedding rows selected by token IDs into a length-by-width target.
	 * Addresses count int32 words; target must not overlap the token span. Returns true after
	 * all writes; false may leave partial output.
	 */
	override fn Embed(tensor: int, tokens: int, target: int, length: int, width: int) bool
	{
		if (!Shape(length, width) || !Span(tokens, length) || !Span(target, length * width)) {
			return false
		}

		if (Overlap(tokens, length, target, length * width) || !weights.Select(tensor)) {
			return false
		}

		if (weights.Columns() != width) {
			return false
		}

		for (var t = 0; t < length; t++) {
			var id = memory.Read(tokens + t)

			if (id < 0 || id >= weights.Rows()) {
				return false
			}
		}

		for (var t = 0; t < length; t++) {
			var id = memory.Read(tokens + t)

			for (var i = 0; i < width; i++) {
				var x = weights.Read(id * width + i)

				if (weights.Failed() || !ValueValid(x)) {
					return false
				}

				memory.Write(target + t * width + i, x)
			}
		}

		return true
	}

	/**
	 * Projects length token-major vectors through a rows-by-columns weight tensor. Source and
	 * target word spans must be disjoint. Returns true after all rows are written; false may
	 * leave partial output.
	 */
	override fn Linear(tensor: int, source: int, target: int, length: int, rows: int, columns: int) bool
	{
		var rowsValid = Shape(length, rows)
		var columnsValid = Shape(length, columns)

		if (!rowsValid || !columnsValid) {
			return false
		}

		if (!Span(source, length * columns) || !Span(target, length * rows)) {
			return false
		}

		if (Overlap(source, length * columns, target, length * rows) || !weights.Select(tensor)) {
			return false
		}

		if (weights.Rows() != rows || weights.Columns() != columns) {
			return false
		}

		for (var t = 0; t < length; t++) {
			for (var r = 0; r < rows; r++) {
				var sum: logic[64] = 0

				for (var c = 0; c < columns; c++) {
					var x = memory.Read(source + t * columns + c)
					var w = weights.Read(r * columns + c)

					var xValid = ValueValid(x)
					var wValid = ValueValid(w)

					if (weights.Failed() || !xValid || !wValid) {
						return false
					}

					sum = math.Add(sum, math.Multiply(x as logic[64], w as logic[64]))
				}

				memory.Write(target + t * rows + r, Round(sum, FRACTION))
			}
		}

		return true
	}

	/**
	 * Applies affine RMS normalization to length vectors of width elements, using a 1-by-width
	 * weight tensor. Source and target word spans must be disjoint. True means completed
	 * writes; false may leave partial output.
	 */
	override fn Norm(tensor: int, source: int, target: int, length: int, width: int) bool
	{
		if (!Shape(length, width) || !Span(source, length * width) || !Span(target, length * width)) {
			return false
		}

		if (Overlap(source, length * width, target, length * width) || !weights.Select(tensor)) {
			return false
		}

		if (weights.Rows() != 1 || weights.Columns() != width) {
			return false
		}

		for (var t = 0; t < length; t++) {
			var sum: logic[64] = 0

			for (var i = 0; i < width; i++) {
				var x = memory.Read(source + t * width + i)

				if (!ValueValid(x)) {
					return false
				}

				sum = math.Add(sum, math.Multiply(x as logic[64], x as logic[64]))
			}

			var mean = math.DivideFloor(sum, width as logic[64])
			var root = math.Sqrt(math.Add(mean, EPSILON as logic[64]))

			for (var i = 0; i < width; i++) {
				var affineInput = memory.Read(source + t * width + i)
				var gamma = weights.Read(i)

				if (weights.Failed() || !ValueValid(gamma)) {
					return false
				}

				var numerator = math.Multiply(affineInput as logic[64], (1 << FRACTION) as logic[64])
				var normalized = math.Clamp(math.DivideEven(numerator, root))
				var result = Round(math.Multiply(normalized as logic[64], gamma as logic[64]), FRACTION)
				memory.Write(target + t * width + i, result)
			}
		}

		return true
	}

	/**
	 * Writes count saturated elementwise sums. Target may equal either input base, but partial
	 * overlap is rejected. Addresses count int32 words. True means completed writes; false may
	 * leave partial output.
	 */
	override fn Add(a: int, b: int, target: int, count: int) bool
	{
		if (count <= 0 || count > MAX_TOKENS * MAX_WIDTH) {
			return false
		}

		if (!Span(a, count) || !Span(b, count) || !Span(target, count)) {
			return false
		}

		var aOverlaps = Overlap(a, count, target, count)
		var bOverlaps = Overlap(b, count, target, count)

		if ((a != target && aOverlaps) || (b != target && bOverlaps)) {
			return false
		}

		for (var i = 0; i < count; i++) {
			var x = memory.Read(a + i)
			var y = memory.Read(b + i)

			var xValid = ValueValid(x)
			var yValid = ValueValid(y)

			if (!xValid || !yValid) {
				return false
			}

			memory.Write(target + i, math.Clamp(math.Add(x as logic[64], y as logic[64])))
		}

		return true
	}

	/**
	 * Writes GELU(a) multiplied by b for count fixed-point elements. All addresses count int32
	 * words; target must be disjoint from both inputs. True means completed writes; false may
	 * leave partial output.
	 */
	override fn Gate(a: int, b: int, target: int, count: int) bool
	{
		if (count <= 0 || count > MAX_TOKENS * MAX_WIDTH) {
			return false
		}

		if (!Span(a, count) || !Span(b, count) || !Span(target, count)) {
			return false
		}

		var aOverlaps = Overlap(a, count, target, count)
		var bOverlaps = Overlap(b, count, target, count)

		if (aOverlaps || bOverlaps) {
			return false
		}

		for (var i = 0; i < count; i++) {
			var x = memory.Read(a + i)
			var y = memory.Read(b + i)

			var xValid = ValueValid(x)
			var yValid = ValueValid(y)

			if (!xValid || !yValid) {
				return false
			}

			var activated = tables.Gelu(x)

			if (!ValueValid(activated)) {
				return false
			}

			memory.Write(target + i, Round(math.Multiply(activated as logic[64], y as logic[64]), FRACTION))
		}

		return true
	}

	/**
	 * Computes unscaled attention over token-major Q/K/V, with heads groups of width features
	 * and nonzero key masks. Uses learned relative-position bias. Target must be disjoint from
	 * Q/K/V and masks; at least one key must be enabled. True means completed writes; false
	 * may leave partial output.
	 */
	override fn Attention(tensor: int, q: int, k: int, v: int, target: int, masks: int,
		length: int, heads: int, width: int) bool
	{
		if (!scratch.IsValidAddress(0) || !scratch.IsValidAddress(2 * MAX_TOKENS - 1)) {
			return false
		}

		if (!Shape(length, width)) {
			return false
		}

		if (heads <= 0 || heads > (divider.Divide(MAX_WIDTH as uint, width as uint) as int)) {
			return false
		}

		var features = heads * width
		var count = length * features

		if (!Span(q, count) || !Span(k, count) || !Span(v, count) || !Span(target, count) || !Span(masks, length)) {
			return false
		}

		var qOverlaps = Overlap(q, count, target, count)
		var kOverlaps = Overlap(k, count, target, count)
		var vOverlaps = Overlap(v, count, target, count)
		var masksOverlaps = Overlap(masks, length, target, count)

		if (qOverlaps || kOverlaps || vOverlaps || masksOverlaps) {
			return false
		}

		if (!weights.Select(tensor) || weights.Columns() != heads) {
			return false
		}

		var any = false

		for (var i = 0; i < length; i++) {
			if (memory.Read(masks + i) != 0) {
				any = true
			}
		}

		if (!any) {
			return false
		}

		for (var h = 0; h < heads; h++) {
			for (var t = 0; t < length; t++) {
				var maximum = -LIMIT

				for (var j = 0; j < length; j++) {
					var sum: logic[64] = 0

					for (var d = 0; d < width; d++) {
						var x = memory.Read(q + t * features + h * width + d)
						var y = memory.Read(k + j * features + h * width + d)

						var xValid = ValueValid(x)
						var yValid = ValueValid(y)

						if (!xValid || !yValid) {
							return false
						}

						sum = math.Add(sum, math.Multiply(x as logic[64], y as logic[64]))
					}

					var bucket = tables.Bucket(j - t)

					if (bucket < 0 || bucket >= weights.Rows()) {
						return false
					}

					var bias = weights.Read(bucket * heads + h)

					if (weights.Failed() || !ValueValid(bias)) {
						return false
					}

					var score = Round(sum, FRACTION)
					score = math.Clamp(math.Add(score as logic[64], bias as logic[64]))
					scratch.Write(j, score)

					if (memory.Read(masks + j) != 0 && score > maximum) {
						maximum = score
					}
				}

				var denominator: logic[64] = 0

				for (var j = 0; j < length; j++) {
					var exponent = 0

					if (memory.Read(masks + j) != 0) {
						exponent = tables.Exponent(maximum - scratch.Read(j))

						if (exponent < 0) {
							return false
						}
					}

					scratch.Write(MAX_TOKENS + j, exponent)
					denominator = math.Add(denominator, exponent as logic[64])
				}

				if (denominator == 0) {
					return false
				}

				for (var j = 0; j < length; j++) {
					var numerator = math.Multiply(scratch.Read(MAX_TOKENS + j) as logic[64], (1 << PROB_FRACTION) as logic[64])
					var probability = math.DivideEven(numerator, denominator) as int
					scratch.Write(MAX_TOKENS + j, probability)

					if (probability < 0 || probability > (1 << PROB_FRACTION)) {
						return false
					}
				}

				for (var d = 0; d < width; d++) {
					var contextSum: logic[64] = 0

					for (var j = 0; j < length; j++) {
						var value = memory.Read(v + j * features + h * width + d)

						if (!ValueValid(value)) {
							return false
						}

						contextSum = math.Add(contextSum, math.Multiply(scratch.Read(MAX_TOKENS + j) as logic[64], value as logic[64]))
					}

					memory.Write(target + t * features + h * width + d, Round(contextSum, PROB_FRACTION))
				}
			}
		}

		return true
	}

	/**
	 * Computes one query against length token-major cached K/V vectors. Width is features per
	 * head; position is the query token position. Causal mode excludes keys beyond position;
	 * zero masks exclude keys. Tensor -1 disables bias, otherwise buckets holds 2*MAX_TOKENS-1
	 * entries centered on relative position zero. Target must not overlap any consumed input
	 * span. True means completed writes; false may leave partial output.
	 */
	override fn CachedAttention(tensor: int, q: int, k: int, v: int, target: int, masks: int,
		length: int, heads: int, width: int, position: int, causal: bool, buckets: int) bool
	{
		if (!scratch.IsValidAddress(0) || !scratch.IsValidAddress(2 * MAX_TOKENS - 1)) {
			return false
		}

		if (!Shape(length, width) || position < 0 || position >= MAX_TOKENS || tensor < -1) {
			return false
		}

		if (heads <= 0 || heads > (divider.Divide(MAX_WIDTH as uint, width as uint) as int)) {
			return false
		}

		var features = heads * width
		var count = length * features

		if (!Span(q, features) || !Span(k, count) || !Span(v, count) || !Span(target, features) || !Span(masks, length)) {
			return false
		}

		var qOverlaps = Overlap(q, features, target, features)
		var kOverlaps = Overlap(k, count, target, features)
		var vOverlaps = Overlap(v, count, target, features)
		var masksOverlaps = Overlap(masks, length, target, features)

		if (qOverlaps || kOverlaps || vOverlaps || masksOverlaps) {
			return false
		}

		if (tensor >= 0) {
			if (!Span(buckets, 2 * MAX_TOKENS - 1)) {
				return false
			}

			if (Overlap(buckets, 2 * MAX_TOKENS - 1, target, features)) {
				return false
			}

			if (!weights.Select(tensor) || weights.Columns() != heads) {
				return false
			}
		}

		var any = false

		for (var i = 0; i < length; i++) {
			if ((!causal || i <= position) && memory.Read(masks + i) != 0) {
				any = true
			}
		}

		if (!any) {
			return false
		}

		for (var h = 0; h < heads; h++) {
			var maximum = -LIMIT

			for (var j = 0; j < length; j++) {
				scratch.Write(j, 0)

				if ((!causal || j <= position) && memory.Read(masks + j) != 0) {
					var sum: logic[64] = 0

					for (var d = 0; d < width; d++) {
						var x = memory.Read(q + h * width + d)
						var y = memory.Read(k + j * features + h * width + d)

						var xValid = ValueValid(x)
						var yValid = ValueValid(y)

						if (!xValid || !yValid) {
							return false
						}

						sum = math.Add(sum, math.Multiply(x as logic[64], y as logic[64]))
					}

					var score = Round(sum, FRACTION)

					if (tensor >= 0) {
						var bucket = memory.Read(buckets + j - position + MAX_TOKENS - 1)

						if (bucket < 0 || bucket >= weights.Rows()) {
							return false
						}

						var bias = weights.Read(bucket * heads + h)

						if (weights.Failed() || !ValueValid(bias)) {
							return false
						}

						score = math.Clamp(math.Add(score as logic[64], bias as logic[64]))
					}

					scratch.Write(j, score)

					if (score > maximum) {
						maximum = score
					}
				}
			}

			var denominator: logic[64] = 0

			for (var j = 0; j < length; j++) {
				var exponent = 0

				if ((!causal || j <= position) && memory.Read(masks + j) != 0) {
					exponent = tables.Exponent(maximum - scratch.Read(j))

					if (exponent < 0) {
						return false
					}
				}

				scratch.Write(MAX_TOKENS + j, exponent)
				denominator = math.Add(denominator, exponent as logic[64])
			}

			if (denominator == 0) {
				return false
			}

			for (var j = 0; j < length; j++) {
				var numerator = math.Multiply(scratch.Read(MAX_TOKENS + j) as logic[64], (1 << PROB_FRACTION) as logic[64])
				var probability = math.DivideEven(numerator, denominator) as int
				scratch.Write(MAX_TOKENS + j, probability)

				if (probability < 0 || probability > (1 << PROB_FRACTION)) {
					return false
				}
			}

			for (var d = 0; d < width; d++) {
				var contextSum: logic[64] = 0

				for (var j = 0; j < length; j++) {
					if ((!causal || j <= position) && memory.Read(masks + j) != 0) {
						var value = memory.Read(v + j * features + h * width + d)

						if (!ValueValid(value)) {
							return false
						}

						contextSum = math.Add(contextSum, math.Multiply(scratch.Read(MAX_TOKENS + j) as logic[64], value as logic[64]))
					}
				}

				memory.Write(target + h * width + d, Round(contextSum, PROB_FRACTION))
			}
		}

		return true
	}

	/**
	 * Projects one columns-element vector through rows weight rows without storing logits.
	 * Returns the lowest row index among equal rounded/clipped maxima, or -1 on invalid shape,
	 * storage, or weights. Does not write activation memory.
	 */
	override fn ProjectArgMax(tensor: int, source: int, rows: int, columns: int) int
	{
		if (rows <= 0 || rows > 65536 || columns <= 0 || columns > MAX_WIDTH) {
			return -1
		}

		if (!Span(source, columns) || !weights.Select(tensor)) {
			return -1
		}

		if (weights.Rows() != rows || weights.Columns() != columns) {
			return -1
		}

		// The earlier row bound is 65536. Smaller configured widths prove safety
		// at compile time; larger configurations retain an exact scheduled guard.
		static if (MAX_WIDTH > 32767) {
			if (rows > (divider.Divide(2147483647 as uint, columns as uint) as int)) {
				return -1
			}
		}

		argmax.Reset()

		for (var row = 0; row < rows; row++) {
			var sum: logic[64] = 0

			for (var c = 0; c < columns; c++) {
				var x = memory.Read(source + c)
				var w = weights.Read(row * columns + c)

				var xValid = ValueValid(x)
				var wValid = ValueValid(w)

				if (weights.Failed() || !xValid || !wValid) {
					return -1
				}

				sum = math.Add(sum, math.Multiply(x as logic[64], w as logic[64]))
			}

			argmax.Accept(row, Round(sum, FRACTION))
		}

		return argmax.GetIndex()
	}
}