From d974ed93aa3b9c9c320f468518ef1589091ed092 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sat, 16 May 2026 22:53:48 +0200 Subject: [PATCH] Add Bf16TensorData + Bf16DenseTensorData (#609) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 1 of the three-phase BF16 dispatch chain. Foundation only — no loader or dispatch changes here; those follow in Phase 2 / Phase 3. `Bf16TensorData : TensorData` is the recognition surface for "this weight is packed BF16 bytes" so the upcoming dispatch in `DefaultCpuOpsJvm.chooseQuantizedMatmul` (Phase 3) can route via the `Bf16MatmulKernel` SPI without dequant-to-FP32 at load. Surface area: - `Bf16TensorData` interface (commonMain) — exposes `packedData: ByteArray` (2 bytes per element, little-endian) for zero-copy hand-off to SIMD matmul kernels, plus `Bf16TensorData.Companion .floatToBf16Bits` / `bf16BitsToFloat` static helpers. - `Bf16DenseTensorData` concrete impl — backed by a packed ByteArray. `get(*indices): Float` decodes BF16 → FP32 on read; `set` truncates FP32 → BF16 (lossy by construction). Bulk `copyToFloatArray()` for consumers that just want all values. - `Bf16DenseTensorData.fromFloatArray(shape, FloatArray)` factory for tests and offline round-tripping. - `Bf16TensorData.toFloatArray()` extension for dequant fallback. 11 unit tests in commonTest cover: round-trip within BF16 precision (1e-2 abs), raw byte-order check (FP32 1.0 = BF16 0x3F80 = bytes [0x80, 0x3F]), get/set primitives, signed zero preservation, 2D / 3D shape strides, bulk copyToFloatArray parity with element-by- element decode, undersized-buffer / out-of-bounds / wrong-rank rejections, and bit-identity for FP32 values that have zero in the low 16 bits. Passes jvmTest + linuxX64Test on the lang-core module. Refs #609. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ainet/lang/tensor/data/Bf16TensorData.kt | 165 ++++++++++++++++++ .../lang/tensor/data/Bf16TensorDataTest.kt | 147 ++++++++++++++++ 2 files changed, 312 insertions(+) create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Bf16TensorData.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/Bf16TensorDataTest.kt diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Bf16TensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Bf16TensorData.kt new file mode 100644 index 000000000..f97c56f05 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/Bf16TensorData.kt @@ -0,0 +1,165 @@ +package sk.ainet.lang.tensor.data + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.DType + +/** + * Tensor data interface for **dense BF16** (bfloat16) values. + * + * Each element is stored as 2 packed little-endian bytes — the high 16 + * bits of an IEEE FP32 value. Conversion to FP32 is the bit-shift + * identity: + * + * `float_bits = (bf16 & 0xFFFF) shl 16` + * + * Unlike block-quantized formats (Q4_K / Q6_K / Q8_0), BF16 carries no + * per-block scale — there is no block structure at all. The encoding + * is `Dense(bytesPerElement = 2)`. + * + * ## Why this type exists + * + * The `Bf16MatmulKernel` SPI (`skainet-backend-api`) needs a way to + * recognize "this weight tensor's data is packed BF16 bytes" so the + * matmul dispatch can route to the SIMD-vectorized BF16 kernel + * (Panama Vector / native FFM, priorities 50 / 100) instead of + * falling back to a dequant-then-FP32-matmul path. This interface is + * the recognition surface: `is Bf16TensorData` in + * `DefaultCpuOpsJvm.chooseQuantizedMatmul` will land in a follow-up. + * + * ## `get` / `set` semantics + * + * `get(*indices): Float` **decodes** BF16 → FP32 on read, so the + * Tensor surface looks like a regular FP32 tensor to consumers that + * don't care about the underlying storage. `set(*indices, value)` + * **truncates** FP32 → BF16 (high 16 bits, zero rounding) — lossy by + * construction, documented at every call site. + * + * For zero-copy access to the packed bytes (e.g. from a SIMD matmul + * kernel that reads 2 bytes per element directly), use [packedData]. + */ +public interface Bf16TensorData : TensorData { + + /** + * Raw packed BF16 bytes — 2 per logical element, little-endian. + * Length is `shape.volume * 2`. Safe for direct hand-off to the + * native / Panama matmul kernels without an intermediate copy. + */ + public val packedData: ByteArray + + public companion object { + /** Bytes per BF16 element. */ + public const val BYTES_PER_ELEMENT: Int = 2 + + /** Convert FP32 → BF16 bits (high 16 bits, zero rounding). */ + public fun floatToBf16Bits(value: Float): Int = + (value.toRawBits() ushr 16) and 0xFFFF + + /** Convert BF16 bits (low 16 bits used) → FP32. */ + public fun bf16BitsToFloat(bf16Bits: Int): Float = + Float.fromBits((bf16Bits and 0xFFFF) shl 16) + } +} + +/** + * Dense BF16 tensor data backed by a packed byte array. + * + * Memory layout: row-major; element at flat index `i` occupies bytes + * `[i*2 .. i*2 + 1]`, low byte first (little-endian). + * + * @param initialShape the logical shape of the tensor (in elements, not bytes). + * @param data the raw packed BF16 byte array, length ≥ `shape.volume * 2`. + */ +public class Bf16DenseTensorData( + initialShape: Shape, + private val data: ByteArray, +) : Bf16TensorData { + + override val shape: Shape = Shape(initialShape.dimensions.copyOf()) + private val strides: IntArray = shape.computeStrides() + override val packedData: ByteArray get() = data + + init { + val requiredBytes = shape.volume * Bf16TensorData.BYTES_PER_ELEMENT + require(data.size >= requiredBytes) { + "Data size ${data.size} is less than required $requiredBytes bytes for ${shape.volume} BF16 elements" + } + } + + override fun get(vararg indices: Int): Float { + val flatIndex = calcFlatIndex(indices) + val byteIdx = flatIndex * Bf16TensorData.BYTES_PER_ELEMENT + val lo = data[byteIdx].toInt() and 0xFF + val hi = data[byteIdx + 1].toInt() and 0xFF + val bf16Bits = (hi shl 8) or lo + return Float.fromBits(bf16Bits shl 16) + } + + override fun set(vararg indices: Int, value: Float) { + val flatIndex = calcFlatIndex(indices) + val byteIdx = flatIndex * Bf16TensorData.BYTES_PER_ELEMENT + val bf16Bits = Bf16TensorData.floatToBf16Bits(value) + data[byteIdx] = (bf16Bits and 0xFF).toByte() + data[byteIdx + 1] = ((bf16Bits ushr 8) and 0xFF).toByte() + } + + override fun copyToFloatArray(): FloatArray { + val volume = shape.volume + val out = FloatArray(volume) + for (i in 0 until volume) { + val byteIdx = i * Bf16TensorData.BYTES_PER_ELEMENT + val lo = data[byteIdx].toInt() and 0xFF + val hi = data[byteIdx + 1].toInt() and 0xFF + out[i] = Float.fromBits(((hi shl 8) or lo) shl 16) + } + return out + } + + private fun calcFlatIndex(indices: IntArray): Int { + require(indices.size == shape.dimensions.size) { + "Number of indices (${indices.size}) must match tensor dimensions (${shape.dimensions.size})" + } + var flatIndex = 0 + for (i in indices.indices) { + val idx = indices[i] + require(idx >= 0 && idx < shape.dimensions[i]) { + "Index $idx out of bounds for dimension $i with size ${shape.dimensions[i]}" + } + flatIndex += idx * strides[i] + } + return flatIndex + } + + public companion object { + /** + * Construct a [Bf16DenseTensorData] from raw BF16 bytes. The byte + * array length must be at least `shape.volume * 2`. + */ + public fun fromRawBytes(shape: Shape, bytes: ByteArray): Bf16DenseTensorData = + Bf16DenseTensorData(shape, bytes) + + /** + * Build a [Bf16DenseTensorData] from a `FloatArray` by truncating + * each value into BF16. Lossy — useful for tests and for + * round-tripping through the dense layout. + */ + public fun fromFloatArray(shape: Shape, values: FloatArray): Bf16DenseTensorData { + require(values.size >= shape.volume) { + "FloatArray length ${values.size} is less than ${shape.volume} BF16 elements required" + } + val bytes = ByteArray(shape.volume * Bf16TensorData.BYTES_PER_ELEMENT) + for (i in 0 until shape.volume) { + val bf16Bits = Bf16TensorData.floatToBf16Bits(values[i]) + bytes[i * 2] = (bf16Bits and 0xFF).toByte() + bytes[i * 2 + 1] = ((bf16Bits ushr 8) and 0xFF).toByte() + } + return Bf16DenseTensorData(shape, bytes) + } + } +} + +/** + * Dequantize a [Bf16TensorData] to a fresh FloatArray. Convenience over + * `tensor.data.copyToFloatArray()` when only the raw FP32 values are + * needed (e.g. parity checks against a reference). + */ +public fun Bf16TensorData.toFloatArray(): FloatArray = copyToFloatArray() diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/Bf16TensorDataTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/Bf16TensorDataTest.kt new file mode 100644 index 000000000..ecc7db47e --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/Bf16TensorDataTest.kt @@ -0,0 +1,147 @@ +package sk.ainet.lang.tensor.data + +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue +import sk.ainet.lang.tensor.Shape + +/** + * Unit tests for [Bf16DenseTensorData]. Covers the FP32 ↔ BF16 round + * trip (within BF16 precision), raw-byte access through [packedData], + * shape / stride correctness on 2D and 3D layouts, and edge cases at + * the value extremes (zero, sign, infinity). + */ +class Bf16TensorDataTest { + + /** BF16 has 7 mantissa bits — relative precision ≈ 1/128 ≈ 0.78%. */ + private val bf16AbsTol = 1e-2f + + @Test + fun fromFloatArray_then_toFloatArray_roundTrips_within_bf16_precision() { + val values = floatArrayOf(0.0f, 1.0f, -1.0f, 0.5f, -0.5f, 3.14159f, -2.71828f, 100.0f) + val tensor = Bf16DenseTensorData.fromFloatArray(Shape(values.size), values) + val out = tensor.toFloatArray() + assertEquals(values.size, out.size) + for (i in values.indices) { + val diff = abs(values[i] - out[i]) + val rel = if (values[i] == 0f) 0f else diff / abs(values[i]) + assertTrue( + diff <= bf16AbsTol || rel <= bf16AbsTol, + "BF16 round-trip exceeds tolerance at $i: in=${values[i]} out=${out[i]} diff=$diff", + ) + } + } + + @Test + fun packedData_exposes_two_bytes_per_element_little_endian() { + // FP32 1.0 = 0x3F800000 → BF16 0x3F80 → bytes [0x80, 0x3F]. + val tensor = Bf16DenseTensorData.fromFloatArray(Shape(1), floatArrayOf(1.0f)) + assertEquals(2, tensor.packedData.size) + assertEquals(0x80.toByte(), tensor.packedData[0], "low byte of BF16(1.0)") + assertEquals(0x3F.toByte(), tensor.packedData[1], "high byte of BF16(1.0)") + } + + @Test + fun get_decodes_packed_bytes_correctly() { + // BF16 1.0 packed → expected via get(). + val bytes = byteArrayOf(0x80.toByte(), 0x3F.toByte()) + val tensor = Bf16DenseTensorData(Shape(1), bytes) + assertEquals(1.0f, tensor.get(0)) + } + + @Test + fun set_truncates_fp32_to_bf16_high_bits() { + val tensor = Bf16DenseTensorData(Shape(1), ByteArray(2)) + tensor.set(0, value = 1.0f) + assertEquals(0x80.toByte(), tensor.packedData[0]) + assertEquals(0x3F.toByte(), tensor.packedData[1]) + assertEquals(1.0f, tensor.get(0)) + } + + @Test + fun zero_round_trips_exactly() { + val tensor = Bf16DenseTensorData.fromFloatArray(Shape(2), floatArrayOf(0.0f, -0.0f)) + assertEquals(0.0f, tensor.get(0)) + // -0.0f in BF16 is bit pattern 0x8000 — distinct from +0.0 (0x0000). + // Float.fromBits(0x8000 << 16) = -0.0f, which compares == 0.0f but has different bits. + assertEquals((-0.0f).toRawBits(), tensor.get(1).toRawBits()) + } + + @Test + fun two_d_shape_strides_correctly() { + // 3×2 matrix of consecutive FP32 values. + val rows = 3 + val cols = 2 + val values = FloatArray(rows * cols) { it.toFloat() } + val tensor = Bf16DenseTensorData.fromFloatArray(Shape(rows, cols), values) + + for (r in 0 until rows) { + for (c in 0 until cols) { + val expected = (r * cols + c).toFloat() + assertEquals(expected, tensor.get(r, c), "mismatch at ($r, $c)") + } + } + } + + @Test + fun three_d_shape_strides_correctly() { + // 2×3×4 tensor — typical "batch, seq, dim" layout. + val a = 2; val b = 3; val c = 4 + val values = FloatArray(a * b * c) { it.toFloat() } + val tensor = Bf16DenseTensorData.fromFloatArray(Shape(a, b, c), values) + + for (i in 0 until a) { + for (j in 0 until b) { + for (k in 0 until c) { + val expected = (i * b * c + j * c + k).toFloat() + assertEquals(expected, tensor.get(i, j, k), "mismatch at ($i, $j, $k)") + } + } + } + } + + @Test + fun copyToFloatArray_matches_element_by_element_decode() { + val values = FloatArray(16) { it * 0.1f } + val tensor = Bf16DenseTensorData.fromFloatArray(Shape(16), values) + val bulk = tensor.copyToFloatArray() + val elementByElement = FloatArray(16) { tensor.get(it) } + assertContentEquals(elementByElement, bulk) + } + + @Test + fun rejects_undersized_byte_buffer() { + // Shape demands 4 elements × 2 bytes = 8 bytes; pass 6. + assertFailsWith { + Bf16DenseTensorData(Shape(4), ByteArray(6)) + } + } + + @Test + fun rejects_index_out_of_bounds() { + val tensor = Bf16DenseTensorData(Shape(3), ByteArray(6)) + assertFailsWith { tensor.get(3) } + assertFailsWith { tensor.get(-1) } + } + + @Test + fun rejects_wrong_number_of_indices() { + val tensor = Bf16DenseTensorData(Shape(2, 3), ByteArray(12)) + assertFailsWith { tensor.get(0) } + assertFailsWith { tensor.get(0, 0, 0) } + } + + @Test + fun floatToBf16Bits_and_back_is_bit_identity_on_clean_values() { + // Values whose FP32 mantissa is zero in low 16 bits → BF16 round-trip is exact. + val cleanValues = floatArrayOf(0.0f, 1.0f, -1.0f, 2.0f, 0.5f, 256.0f) + for (v in cleanValues) { + val bf16 = Bf16TensorData.floatToBf16Bits(v) + val recovered = Bf16TensorData.bf16BitsToFloat(bf16) + assertEquals(v, recovered, "round-trip mismatch on $v: bf16=0x${bf16.toString(16)} recovered=$recovered") + } + } +}