Livt.ML

src/pooling/MaxPool8x28To14Ram.lvt

namespace Livt.ML.Pooling

using Livt.ML.Storage

/**
 * RAM-backed non-streaming 2x2 stride-two max pool for eight 28x28 planes.
 */
component MaxPool8x28To14Ram
{
	// Compile-time dimensions, bounds, and numeric constants.
	public const INPUT_SIZE: int = 6272
	public const OUTPUT_SIZE: int = 1568

	// Input tensor storage; populate the required span before computation.
	inputValues: ByteTensorRam8192

	// Owned output tensor RAM; Compute writes the result before readback.
	outputValues: ByteTensorRam2048

	/**
	 * Allocates owned input and output tensor RAMs. Write the complete input before Compute;
	 * construction does not initialize RAM contents.
	 */
	new()
	{
		this.inputValues = new ByteTensorRam8192(INPUT_SIZE)
		this.outputValues = new ByteTensorRam2048(OUTPUT_SIZE)
	}

	/**
	 * Stores one flattened NCHW input value and ignores an invalid index.
	 */
	public fn SetInput(index: int, value: int)
	{
		if (index < 0) {
			return
		}

		if (index >= INPUT_SIZE) {
			return
		}

		this.inputValues.Write(index, value as byte)
	}

	/**
	 * Computes eight row-major 14x14 output planes.
	 */
	public fn Compute()
	{
		for (var channel = 0; channel < 8; channel++) {
			for (var row = 0; row < 14; row++) {
				for (var column = 0; column < 14; column++) {
					var maxValue: byte = 0x00

					for (var poolRow = 0; poolRow < 2; poolRow++) {
						for (var poolColumn = 0; poolColumn < 2; poolColumn++) {
							var inputRow = row * 2 + poolRow
							var inputColumn = column * 2 + poolColumn
							var inputIndex = channel * 784 + inputRow * 28 + inputColumn
							var value = this.inputValues.Read(inputIndex)

							if (value > maxValue) {
								maxValue = value
							}
						}
					}

					this.outputValues.Write(channel * 196 + row * 14 + column, maxValue)
				}
			}
		}
	}

	/**
	 * Returns one flattened NCHW output value, or zero for an invalid index.
	 */
	public fn GetOutput(index: int) int
	{
		if (index < 0) {
			return 0
		}

		if (index >= OUTPUT_SIZE) {
			return 0
		}

		return this.outputValues.Read(index) as int
	}
}