From dff4a6b625f45409a8ae414a2f6cb912127969a9 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 26 Jul 2026 18:02:03 +0200 Subject: [PATCH 1/4] feat(lang): shared narrow-float codec + dense storage for BF16 and FP16 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First stage (B1) of #884. Introduces one 16-bit float layer parameterized by codec, so BF16 and FP16 share an implementation instead of growing two half-built stacks. BF16 had storage, kernels, dispatch and a SafeTensors KEEP_NATIVE path but no trace capture, no elementwise ops and no `convert` target; FP16 had the dtype tag, `convert` and trace capture but no storage, no kernel and no loader. The gaps are complementary, so finishing either alone leaves the other broken. Added: - `NarrowFloatCodec` with `Bf16Codec` and `Fp16Codec` (skainet-lang-core, commonMain — pure integer bit math, no JDK dependency, so every KMP target gets it). Becomes the canonical home for float<->half conversion; the six duplicated private `halfToFloat`/`decodeHalf` copies can now collapse onto it. - `NarrowFloatTensorData` / `NarrowFloatDenseTensorData` — the shared 2-byte dense storage and the recognition surface for narrow-float kernel dispatch. - `Fp16DenseTensorData` — the class the SafeTensors and GGUF loaders already name in their "no FP16 backing yet" errors, which is why `Require(FP16)` had to be rejected outright. Unblocks the loader work (B3). `Bf16Codec.encode` TRUNCATES rather than rounding to nearest. This is contractual, not an oversight: it matches the existing `floatToBf16Bits`, the safetensors writer, and the on-device A/B that qualified bf16 as a bit-exact drop-in for the f16 vmfb. A regression test pins the two implementations together so a future "fix" fails loudly. `Fp16Codec` implements full IEEE 754 binary16 with round-to-nearest-ties-to-even: normals, gradual underflow into subnormals, overflow to +/-Inf past 65504, and NaN payload preservation (never collapsing a NaN into an infinity). Fixed two latent bugs found while wiring this up: - `TraceSession.refOf` had no `BF16::class` arm, so a BF16 tensor traced as FP32. The captured dtype is what the StableHLO converter reads to pick an MLIR element type, so this is precisely why bf16 weights could only be produced by regex-rewriting the emitted MLIR text after the fact (SKaiNET-transformers#248). - `TensorStorageFactory.toTensorData` decoded FLOAT16/BFLOAT16 through the FLOAT32 branch's hard 4-byte reader, yielding half the element count with each "float" assembled from two adjacent narrow elements — silent garbage. Readers already tagged these `Dense(bytesPerElement = 2)`; only the decode side ignored it. Widening to f32 here preserves the method's existing contract; preserving 2-byte storage end-to-end is the loader KEEP_NATIVE path, not this one. `Bf16DenseTensorData` is re-parented onto `NarrowFloatDenseTensorData` and `Bf16TensorData` onto `NarrowFloatTensorData`. Public API is preserved: `floatToBf16Bits`/`bf16BitsToFloat` now delegate to `Bf16Codec`, and the moved members remain resolvable via normal JVM hierarchy lookup. NOTE: `Bf16TensorData` gains `getCodec()` with a default body, so an out-of-tree class implementing that interface and compiled against the old API would hit AbstractMethodError. `Bf16DenseTensorData` is the only implementor in this repo. Tests: 41 new (21 codec, 16 dense storage, 4 storage-decode regression). The pre-existing BF16 suites pass unmodified — Bf16TensorDataTest (12), PanamaVectorBf16MatmulKernelParityTest (10), Bf16MatmulDispatchTest (3); the last of these confirms `is Bf16TensorData` still matches after re-parenting. apiDump regenerated; apiCheck green. Full engine build clean (2133 tasks). Refs #884 Co-Authored-By: Claude Opus 5 (1M context) --- .../api/jvm/skainet-lang-core.api | 88 ++++++- .../ainet/lang/tensor/data/Bf16TensorData.kt | 87 ++----- .../lang/tensor/data/NarrowFloatTensorData.kt | 165 ++++++++++++ .../tensor/storage/TensorStorageFactory.kt | 32 ++- .../sk/ainet/lang/trace/TraceSession.kt | 5 + .../sk/ainet/lang/types/NarrowFloatCodec.kt | 152 +++++++++++ .../data/NarrowFloatDenseTensorDataTest.kt | 185 ++++++++++++++ .../storage/NarrowFloatStorageDecodeTest.kt | 93 +++++++ .../ainet/lang/types/NarrowFloatCodecTest.kt | 241 ++++++++++++++++++ 9 files changed, 968 insertions(+), 80 deletions(-) create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/NarrowFloatTensorData.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/NarrowFloatCodec.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatDenseTensorDataTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/NarrowFloatStorageDecodeTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/NarrowFloatCodecTest.kt diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 338bceb0..2ce8cc12 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -2862,16 +2862,10 @@ public final class sk/ainet/lang/tensor/benchmark/SlicingBenchmarksKt { public static synthetic fun slicingBenchmarkSuite$default (Ljava/lang/String;Lsk/ainet/lang/tensor/benchmark/SlicingBenchmarkConfig;Lkotlin/jvm/functions/Function0;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lsk/ainet/benchmark/BenchmarkSuite; } -public final class sk/ainet/lang/tensor/data/Bf16DenseTensorData : sk/ainet/lang/tensor/data/Bf16TensorData { +public final class sk/ainet/lang/tensor/data/Bf16DenseTensorData : sk/ainet/lang/tensor/data/NarrowFloatDenseTensorData, sk/ainet/lang/tensor/data/Bf16TensorData { public static final field Companion Lsk/ainet/lang/tensor/data/Bf16DenseTensorData$Companion; public fun (Lsk/ainet/lang/tensor/Shape;[B)V - public fun copyToFloatArray ()[F - public fun get ([I)Ljava/lang/Float; - public synthetic fun get ([I)Ljava/lang/Object; - public fun getPackedData ()[B - public fun getShape ()Lsk/ainet/lang/tensor/Shape; - public fun set ([IF)V - public synthetic fun set ([ILjava/lang/Object;)V + public fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; } public final class sk/ainet/lang/tensor/data/Bf16DenseTensorData$Companion { @@ -2879,10 +2873,10 @@ public final class sk/ainet/lang/tensor/data/Bf16DenseTensorData$Companion { public final fun fromRawBytes (Lsk/ainet/lang/tensor/Shape;[B)Lsk/ainet/lang/tensor/data/Bf16DenseTensorData; } -public abstract interface class sk/ainet/lang/tensor/data/Bf16TensorData : sk/ainet/lang/tensor/data/TensorData { +public abstract interface class sk/ainet/lang/tensor/data/Bf16TensorData : sk/ainet/lang/tensor/data/NarrowFloatTensorData { public static final field BYTES_PER_ELEMENT I public static final field Companion Lsk/ainet/lang/tensor/data/Bf16TensorData$Companion; - public abstract fun getPackedData ()[B + public fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; } public final class sk/ainet/lang/tensor/data/Bf16TensorData$Companion { @@ -2893,6 +2887,7 @@ public final class sk/ainet/lang/tensor/data/Bf16TensorData$Companion { public final class sk/ainet/lang/tensor/data/Bf16TensorData$DefaultImpls { public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/Bf16TensorData;)[F + public static fun getCodec (Lsk/ainet/lang/tensor/data/Bf16TensorData;)Lsk/ainet/lang/types/NarrowFloatCodec; } public final class sk/ainet/lang/tensor/data/Bf16TensorDataKt { @@ -2967,6 +2962,16 @@ public final class sk/ainet/lang/tensor/data/FloatBufferTensorData$DefaultImpls public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/FloatBufferTensorData;)[F } +public final class sk/ainet/lang/tensor/data/Fp16DenseTensorData : sk/ainet/lang/tensor/data/NarrowFloatDenseTensorData { + public static final field Companion Lsk/ainet/lang/tensor/data/Fp16DenseTensorData$Companion; + public fun (Lsk/ainet/lang/tensor/Shape;[B)V +} + +public final class sk/ainet/lang/tensor/data/Fp16DenseTensorData$Companion { + public final fun fromFloatArray (Lsk/ainet/lang/tensor/Shape;[F)Lsk/ainet/lang/tensor/data/Fp16DenseTensorData; + public final fun fromRawBytes (Lsk/ainet/lang/tensor/Shape;[B)Lsk/ainet/lang/tensor/data/Fp16DenseTensorData; +} + public abstract interface class sk/ainet/lang/tensor/data/IntArrayTensorData : sk/ainet/lang/tensor/data/TensorData { public abstract fun getBuffer ()[I } @@ -3073,6 +3078,42 @@ public final class sk/ainet/lang/tensor/data/MmapTensorSource$Companion { public static synthetic fun fromChannel$default (Lsk/ainet/lang/tensor/data/MmapTensorSource$Companion;Ljava/nio/channels/FileChannel;JJILjava/lang/Object;)Lsk/ainet/lang/tensor/data/MmapTensorSource; } +public class sk/ainet/lang/tensor/data/NarrowFloatDenseTensorData : sk/ainet/lang/tensor/data/NarrowFloatTensorData { + public static final field Companion Lsk/ainet/lang/tensor/data/NarrowFloatDenseTensorData$Companion; + public fun (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/types/NarrowFloatCodec;)V + public fun copyToFloatArray ()[F + public fun get ([I)Ljava/lang/Float; + public synthetic fun get ([I)Ljava/lang/Object; + public fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; + public fun getPackedData ()[B + public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun set ([IF)V + public synthetic fun set ([ILjava/lang/Object;)V +} + +public final class sk/ainet/lang/tensor/data/NarrowFloatDenseTensorData$Companion { + public final fun fromFloatArray (Lsk/ainet/lang/tensor/Shape;[FLsk/ainet/lang/types/NarrowFloatCodec;)Lsk/ainet/lang/tensor/data/NarrowFloatDenseTensorData; +} + +public abstract interface class sk/ainet/lang/tensor/data/NarrowFloatTensorData : sk/ainet/lang/tensor/data/TensorData { + public static final field BYTES_PER_ELEMENT I + public static final field Companion Lsk/ainet/lang/tensor/data/NarrowFloatTensorData$Companion; + public abstract fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; + public abstract fun getPackedData ()[B +} + +public final class sk/ainet/lang/tensor/data/NarrowFloatTensorData$Companion { + public static final field BYTES_PER_ELEMENT I +} + +public final class sk/ainet/lang/tensor/data/NarrowFloatTensorData$DefaultImpls { + public static fun copyToFloatArray (Lsk/ainet/lang/tensor/data/NarrowFloatTensorData;)[F +} + +public final class sk/ainet/lang/tensor/data/NarrowFloatTensorDataKt { + public static final fun toFloatArray (Lsk/ainet/lang/tensor/data/NarrowFloatTensorData;)[F +} + public final class sk/ainet/lang/tensor/data/PprintKt { public static final fun pprint (Lsk/ainet/lang/tensor/data/TensorData;)Ljava/lang/String; } @@ -5714,6 +5755,14 @@ public final class sk/ainet/lang/types/BF16 : sk/ainet/lang/types/DType { public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } +public final class sk/ainet/lang/types/Bf16Codec : sk/ainet/lang/types/NarrowFloatCodec { + public static final field INSTANCE Lsk/ainet/lang/types/Bf16Codec; + public fun decode (I)F + public fun encode (F)I + public fun getBytesPerElement ()I + public fun getDtype ()Lsk/ainet/lang/types/DType; +} + public abstract interface class sk/ainet/lang/types/DType { public static final field Companion Lsk/ainet/lang/types/DType$Companion; public static fun bf16 ()Lsk/ainet/lang/types/DType; @@ -5845,6 +5894,14 @@ public final class sk/ainet/lang/types/FP64 : sk/ainet/lang/types/DType { public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } +public final class sk/ainet/lang/types/Fp16Codec : sk/ainet/lang/types/NarrowFloatCodec { + public static final field INSTANCE Lsk/ainet/lang/types/Fp16Codec; + public fun decode (I)F + public fun encode (F)I + public fun getBytesPerElement ()I + public fun getDtype ()Lsk/ainet/lang/types/DType; +} + public final class sk/ainet/lang/types/Int16 : sk/ainet/lang/types/DType { public static final field INSTANCE Lsk/ainet/lang/types/Int16; public fun getName ()Ljava/lang/String; @@ -5885,6 +5942,17 @@ public final class sk/ainet/lang/types/Int8 : sk/ainet/lang/types/DType { public fun promoteTo (Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; } +public abstract interface class sk/ainet/lang/types/NarrowFloatCodec { + public abstract fun decode (I)F + public abstract fun encode (F)I + public fun getBytesPerElement ()I + public abstract fun getDtype ()Lsk/ainet/lang/types/DType; +} + +public final class sk/ainet/lang/types/NarrowFloatCodec$DefaultImpls { + public static fun getBytesPerElement (Lsk/ainet/lang/types/NarrowFloatCodec;)I +} + public final class sk/ainet/lang/types/PrecisionValidationKt { public static final fun isValidPrecisionChain (Ljava/util/List;)Z public static final fun validateCompatibility (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;)V 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 index f97c56f0..48b08989 100644 --- 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 @@ -1,7 +1,8 @@ package sk.ainet.lang.tensor.data import sk.ainet.lang.tensor.Shape -import sk.ainet.lang.types.DType +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.NarrowFloatCodec /** * Tensor data interface for **dense BF16** (bfloat16) values. @@ -37,26 +38,26 @@ import sk.ainet.lang.types.DType * 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 { +public interface Bf16TensorData : NarrowFloatTensorData { - /** - * 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 + /** Always [Bf16Codec] — BF16 is the high 16 bits of an FP32 value. */ + override val codec: NarrowFloatCodec get() = Bf16Codec 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 FP32 → BF16 bits (high 16 bits, zero rounding). + * + * Delegates to [Bf16Codec] so there is a single implementation. The truncating (rather + * than round-to-nearest) behaviour is contractual — emitted weight archives and the + * on-device bf16-vs-f16 A/B depend on it. + */ + public fun floatToBf16Bits(value: Float): Int = Bf16Codec.encode(value) /** Convert BF16 bits (low 16 bits used) → FP32. */ - public fun bf16BitsToFloat(bf16Bits: Int): Float = - Float.fromBits((bf16Bits and 0xFFFF) shl 16) + public fun bf16BitsToFloat(bf16Bits: Int): Float = Bf16Codec.decode(bf16Bits) } } @@ -71,63 +72,11 @@ public interface Bf16TensorData : TensorData { */ 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 + data: ByteArray, +) : NarrowFloatDenseTensorData(initialShape, data, Bf16Codec), Bf16TensorData { - 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 - } + /** Fixed to [Bf16Codec]; resolves the base class / [Bf16TensorData] inheritance. */ + override val codec: NarrowFloatCodec get() = Bf16Codec public companion object { /** diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/NarrowFloatTensorData.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/NarrowFloatTensorData.kt new file mode 100644 index 00000000..5f95a53b --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/NarrowFloatTensorData.kt @@ -0,0 +1,165 @@ +package sk.ainet.lang.tensor.data + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.DType +import sk.ainet.lang.types.Fp16Codec +import sk.ainet.lang.types.NarrowFloatCodec + +/** + * Tensor data whose elements are stored as **packed 16-bit floats** — two bytes per element, + * little-endian, no block structure and no per-block scale. The encoding is `Dense(2)`. + * + * This is the recognition surface for narrow-float dispatch: a matmul kernel can test + * `is NarrowFloatTensorData` and read [packedData] directly at 2 bytes per element instead of + * forcing a dequant-to-FP32 copy first. [codec] says which 16-bit format the bytes are in, so one + * kernel implementation serves both BF16 and FP16. + * + * `get`/`copyToFloatArray` decode to FP32, so consumers that do not care about the storage width + * see an ordinary float tensor. `set` re-encodes and is lossy by construction. + * + * @see sk.ainet.lang.types.Bf16Codec + * @see sk.ainet.lang.types.Fp16Codec + */ +public interface NarrowFloatTensorData : TensorData { + + /** + * Raw packed bytes — 2 per logical element, little-endian. Length is `shape.volume * 2`. + * Safe for direct hand-off to native / Panama kernels without an intermediate copy. + */ + public val packedData: ByteArray + + /** The 16-bit format [packedData] is encoded in. */ + public val codec: NarrowFloatCodec + + public companion object { + /** Bytes per element for every narrow float format. */ + public const val BYTES_PER_ELEMENT: Int = 2 + } +} + +/** + * Dense narrow-float 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. The concrete 16-bit format is supplied as a [NarrowFloatCodec], which is what lets BF16 + * and FP16 share one implementation rather than maintaining two parallel stacks. + * + * @param initialShape logical shape, in elements (not bytes). + * @param data packed bytes, length ≥ `shape.volume * 2`. + * @param codec the 16-bit format of [data]. + */ +public open class NarrowFloatDenseTensorData( + initialShape: Shape, + private val data: ByteArray, + override val codec: NarrowFloatCodec, +) : NarrowFloatTensorData { + + 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 * NarrowFloatTensorData.BYTES_PER_ELEMENT + require(data.size >= requiredBytes) { + "Data size ${data.size} is less than required $requiredBytes bytes " + + "for ${shape.volume} ${codec.dtype.name} elements" + } + } + + override fun get(vararg indices: Int): Float { + val byteIdx = calcFlatIndex(indices) * NarrowFloatTensorData.BYTES_PER_ELEMENT + val lo = data[byteIdx].toInt() and 0xFF + val hi = data[byteIdx + 1].toInt() and 0xFF + return codec.decode((hi shl 8) or lo) + } + + override fun set(vararg indices: Int, value: Float) { + val byteIdx = calcFlatIndex(indices) * NarrowFloatTensorData.BYTES_PER_ELEMENT + val bits = codec.encode(value) + data[byteIdx] = (bits and 0xFF).toByte() + data[byteIdx + 1] = ((bits 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 * NarrowFloatTensorData.BYTES_PER_ELEMENT + val lo = data[byteIdx].toInt() and 0xFF + val hi = data[byteIdx + 1].toInt() and 0xFF + out[i] = codec.decode((hi shl 8) or lo) + } + 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 { + /** Pack a `FloatArray` into [codec]'s 16-bit format. Lossy by construction. */ + public fun fromFloatArray( + shape: Shape, + values: FloatArray, + codec: NarrowFloatCodec, + ): NarrowFloatDenseTensorData { + require(values.size >= shape.volume) { + "FloatArray length ${values.size} is less than ${shape.volume} elements required" + } + val bytes = ByteArray(shape.volume * NarrowFloatTensorData.BYTES_PER_ELEMENT) + for (i in 0 until shape.volume) { + val bits = codec.encode(values[i]) + bytes[i * 2] = (bits and 0xFF).toByte() + bytes[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return NarrowFloatDenseTensorData(shape, bytes, codec) + } + } +} + +/** + * Dense **IEEE binary16** tensor data. + * + * The counterpart to [Bf16DenseTensorData], and the class the SafeTensors and GGUF loaders name in + * their "no FP16 backing yet" errors — its absence was the reason `Require(FP16)` had to be + * rejected outright. + */ +public class Fp16DenseTensorData( + initialShape: Shape, + data: ByteArray, +) : NarrowFloatDenseTensorData(initialShape, data, Fp16Codec) { + + public companion object { + /** Wrap raw packed FP16 bytes; length must be ≥ `shape.volume * 2`. */ + public fun fromRawBytes(shape: Shape, bytes: ByteArray): Fp16DenseTensorData = + Fp16DenseTensorData(shape, bytes) + + /** Build from FP32 values, rounding each to nearest binary16 (ties to even). */ + public fun fromFloatArray(shape: Shape, values: FloatArray): Fp16DenseTensorData { + require(values.size >= shape.volume) { + "FloatArray length ${values.size} is less than ${shape.volume} FP16 elements required" + } + val bytes = ByteArray(shape.volume * NarrowFloatTensorData.BYTES_PER_ELEMENT) + for (i in 0 until shape.volume) { + val bits = Fp16Codec.encode(values[i]) + bytes[i * 2] = (bits and 0xFF).toByte() + bytes[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return Fp16DenseTensorData(shape, bytes) + } + } +} + +/** Decode a [NarrowFloatTensorData] to a fresh FloatArray. */ +public fun NarrowFloatTensorData.toFloatArray(): FloatArray = copyToFloatArray() diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorageFactory.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorageFactory.kt index a93b463e..10ef17e3 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorageFactory.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/TensorStorageFactory.kt @@ -10,9 +10,12 @@ import sk.ainet.lang.tensor.data.Q4_KTensorData import sk.ainet.lang.tensor.data.Q8_0BlockTensorData import sk.ainet.lang.tensor.data.Q8_0TensorData import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.Bf16Codec import sk.ainet.lang.types.DType import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec import sk.ainet.lang.types.Int32 +import sk.ainet.lang.types.NarrowFloatCodec /** * Factory methods for constructing [TensorStorage] from existing SKaiNET types @@ -175,10 +178,22 @@ public object TensorStorageFactory { return when (storage.encoding) { is TensorEncoding.Dense -> when (storage.logicalType) { - LogicalDType.FLOAT32, LogicalDType.FLOAT16, LogicalDType.BFLOAT16 -> { + LogicalDType.FLOAT32 -> { val floats = bytesToFloatArray(bytes) DenseFloatArrayTensorData(storage.shape, floats) as TensorData } + // 2 bytes per element, NOT 4. These previously shared the FLOAT32 branch, so + // `bytesToFloatArray` produced half the element count with each "float" assembled + // from two adjacent narrow elements — silent garbage. Readers already tag these + // correctly as `TensorEncoding.Dense(bytesPerElement = 2)`; the decode side simply + // ignored it. Widening to f32 here keeps this method's existing contract (it has + // always returned float-backed data); preserving 2-byte storage end-to-end is the + // loader `KEEP_NATIVE` path, not this one. + LogicalDType.FLOAT16, LogicalDType.BFLOAT16 -> { + val codec = if (storage.logicalType == LogicalDType.FLOAT16) Fp16Codec else Bf16Codec + val floats = narrowBytesToFloatArray(bytes, codec) + DenseFloatArrayTensorData(storage.shape, floats) as TensorData + } LogicalDType.INT32 -> { val ints = bytesToIntArray(bytes) DenseIntArrayTensorData(storage.shape, ints) as TensorData @@ -214,6 +229,21 @@ public object TensorStorageFactory { ) } + /** + * Decode packed little-endian 16-bit floats to FP32 using [codec]. One element per 2 bytes — + * the element count is `bytes.size / 2`, which is the whole point of this helper existing + * separately from [bytesToFloatArray]. + */ + private fun narrowBytesToFloatArray(bytes: ByteArray, codec: NarrowFloatCodec): FloatArray { + val count = bytes.size / codec.bytesPerElement + return FloatArray(count) { i -> + val off = i * 2 + val lo = bytes[off].toInt() and 0xFF + val hi = bytes[off + 1].toInt() and 0xFF + codec.decode((hi shl 8) or lo) + } + } + private fun bytesToFloatArray(bytes: ByteArray): FloatArray { val count = bytes.size / 4 return FloatArray(count) { i -> diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TraceSession.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TraceSession.kt index 98a64f84..fb610fd8 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TraceSession.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TraceSession.kt @@ -18,10 +18,15 @@ public open class TraceSession { public open fun refOf(tensor: Tensor<*, *>): TensorRef { val key = unwrap(tensor) return tensorToRef.getOrPut(key) { + // The captured dtype is what the StableHLO converter reads to pick an MLIR element + // type, so a missing arm here silently downgrades the emitted graph. BF16 was absent + // and fell through to FP32, which is why bf16 weights could only be produced by + // rewriting the emitted MLIR text after the fact. val dtypeInstance: DType = when (tensor.dtype) { Int32::class -> Int32 FP32::class -> FP32 FP16::class -> FP16 + BF16::class -> BF16 Int8::class -> Int8 Int4::class -> Int4 Ternary::class -> Ternary diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/NarrowFloatCodec.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/NarrowFloatCodec.kt new file mode 100644 index 00000000..512152ec --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/NarrowFloatCodec.kt @@ -0,0 +1,152 @@ +package sk.ainet.lang.types + +/** + * Bit-level codec for a **dense 16-bit float** format. + * + * Both supported formats store one element per 2 bytes and are decoded to `Float` for compute — + * narrow floats are a *storage* width here, never an accumulate width. Kernels widen to f32 lanes, + * accumulate in f32, and narrow again on store, which is what PyTorch/JAX/tensor cores do. + * + * The two formats are not interchangeable: + * + * | | exponent | mantissa | relative step | max finite | + * |---|---|---|---|---| + * | [Bf16Codec] | 8 bits | 7 (+1 implicit) | 2⁻⁸ ≈ 0.39% | ~3.39e38 (f32 range) | + * | [Fp16Codec] | 5 bits | 10 (+1 implicit) | 2⁻¹¹ ≈ 0.049% | 65504 | + * + * BF16 keeps FP32's full exponent range, so overflow is structurally impossible on conversion; + * FP16 buys three more mantissa bits at the cost of a 65504 ceiling. For weights dequantized from + * block-quantized sources (Q4_K/Q5_K/…) the quantization error dominates both. + * + * Implementations are pure integer bit math with no JVM/JDK dependency, so they are usable from + * every Kotlin Multiplatform target. (The JVM's `Float.floatToFloat16` is JDK 20+ and unavailable + * in `commonMain`.) + */ +public interface NarrowFloatCodec { + + /** The [DType] this codec encodes — [BF16] or [FP16]. */ + public val dtype: DType + + /** Always 2. Present so callers can size buffers without branching on [dtype]. */ + public val bytesPerElement: Int get() = 2 + + /** + * Encode an FP32 value into this format's 16-bit pattern (returned in the low 16 bits). + * Lossy by construction; see each implementation for its rounding rule. + */ + public fun encode(value: Float): Int + + /** Decode a 16-bit pattern (low 16 bits used) back to FP32. Always exact — f32 is a superset. */ + public fun decode(bits: Int): Float +} + +/** + * **bfloat16** — the high 16 bits of an IEEE FP32 value. + * + * Conversion is the bit-shift identity in both directions, which is why it costs essentially + * nothing: `float_bits = bf16 shl 16`. + * + * [encode] **truncates** (round-toward-zero) rather than rounding to nearest. This is deliberate + * and must not be "fixed": it matches the existing `Bf16TensorData.floatToBf16Bits` contract, the + * safetensors writer, and the on-device A/B that verified bf16 as a bit-exact drop-in for the f16 + * vmfb. Changing it would silently alter every emitted weight archive. + */ +public object Bf16Codec : NarrowFloatCodec { + + override val dtype: DType get() = BF16 + + /** FP32 → BF16 by truncation (high 16 bits, zero rounding). Bit-exact with the legacy path. */ + override fun encode(value: Float): Int = (value.toRawBits() ushr 16) and 0xFFFF + + /** BF16 → FP32 — exact; the low 16 mantissa bits are zero-filled. */ + override fun decode(bits: Int): Float = Float.fromBits((bits and 0xFFFF) shl 16) +} + +/** + * **IEEE 754 binary16** (half precision): 1 sign / 5 exponent / 10 mantissa bits. + * + * [encode] implements **round-to-nearest, ties-to-even** across the full input domain: + * normals, gradual underflow into binary16 subnormals, overflow to ±Inf beyond 65504, and + * NaN payload preservation (never collapsing a NaN into an infinity). + */ +public object Fp16Codec : NarrowFloatCodec { + + override val dtype: DType get() = FP16 + + /** + * FP32 → FP16, round-to-nearest-ties-to-even. + * + * Values above binary16's max finite (65504) plus half an ulp go to ±Inf; values at or below + * half the smallest subnormal (2⁻²⁵) go to ±0. NaN stays NaN with as much payload as fits, + * and is forced quiet. + */ + override fun encode(value: Float): Int { + val bits = value.toRawBits() + val sign = (bits ushr 16) and 0x8000 + val absBits = bits and 0x7FFF_FFFF + + // Inf / NaN — exponent field all ones. + if (absBits >= 0x7F80_0000) { + if (absBits == 0x7F80_0000) return sign or 0x7C00 // ±Inf + // NaN: keep the top payload bits, force quiet, and never let the payload become 0 + // (that would turn a NaN into an Inf). + val payload = (absBits ushr 13) and 0x03FF + return sign or 0x7C00 or 0x0200 or payload + } + + val exp = (absBits ushr 23) - 127 // unbiased FP32 exponent + val mant = absBits and 0x007F_FFFF + + // Beyond binary16's exponent range — round-to-nearest still lands on Inf. + if (exp >= 16) return sign or 0x7C00 + + if (exp >= -14) { + // Normal binary16: drop 13 mantissa bits, round to nearest even. + val half = ((exp + 15) shl 10) or (mant ushr 13) + val rem = mant and 0x1FFF + val roundUp = rem > 0x1000 || (rem == 0x1000 && (half and 1) == 1) + // A carry out of the mantissa flows into the exponent field, which is the correct + // IEEE result (and produces Inf at the top of the range). + return sign or (half + if (roundUp) 1 else 0) + } + + // Gradual underflow: binary16 subnormals represent m * 2^-24 for m in [1, 1023]. + // Anything below 2^-25 (half the smallest subnormal) rounds to zero. + if (exp < -25) return sign + val full = mant or 0x0080_0000 // restore the implicit leading 1 + val shift = -1 - exp // exp <= -15 => shift >= 14 + val q = full ushr shift + val rem = full and ((1 shl shift) - 1) + val halfBit = 1 shl (shift - 1) + val roundUp = rem > halfBit || (rem == halfBit && (q and 1) == 1) + // q may carry from 1023 to 1024, yielding the smallest normal — also correct IEEE. + return sign or (q + if (roundUp) 1 else 0) + } + + /** FP16 → FP32 — always exact, including subnormals (which renormalize into FP32 normals). */ + override fun decode(bits: Int): Float { + val h = bits and 0xFFFF + val sign = (h and 0x8000) shl 16 + val exp = (h ushr 10) and 0x1F + val mant = h and 0x03FF + + return when (exp) { + 0 -> { + if (mant == 0) { + Float.fromBits(sign) // ±0 + } else { + // Subnormal: renormalize until the implicit bit position is set. + var m = mant + var e = -1 + do { + m = m shl 1 + e++ + } while (m and 0x0400 == 0) + Float.fromBits(sign or ((127 - 15 - e) shl 23) or ((m and 0x03FF) shl 13)) + } + } + 0x1F -> Float.fromBits(sign or 0x7F80_0000 or (mant shl 13)) // ±Inf / NaN + else -> Float.fromBits(sign or ((exp - 15 + 127) shl 23) or (mant shl 13)) + } + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatDenseTensorDataTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatDenseTensorDataTest.kt new file mode 100644 index 00000000..e7734124 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatDenseTensorDataTest.kt @@ -0,0 +1,185 @@ +package sk.ainet.lang.tensor.data + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.Fp16Codec +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertNotEquals +import kotlin.test.assertTrue + +/** + * Contract for the shared narrow-float storage: [NarrowFloatDenseTensorData] and its + * [Fp16DenseTensorData] specialization. Mirrors `Bf16TensorDataTest`, which continues to pin the + * BF16 side independently. + */ +class NarrowFloatDenseTensorDataTest { + + @Test + fun fp16_round_trips_exactly_representable_values() { + val values = floatArrayOf(1.0f, -2.0f, 0.5f, 3.0f, -0.25f, 1024.0f) + val data = Fp16DenseTensorData.fromFloatArray(Shape(values.size), values) + val out = data.copyToFloatArray() + for (i in values.indices) { + assertEquals(values[i], out[i], "element $i") + } + } + + @Test + fun fp16_stores_two_bytes_per_element_little_endian() { + val data = Fp16DenseTensorData.fromFloatArray(Shape(1), floatArrayOf(1.0f)) + assertEquals(2, data.packedData.size) + // 1.0f is 0x3C00 in binary16 -> low byte first + assertEquals(0x00.toByte(), data.packedData[0]) + assertEquals(0x3C.toByte(), data.packedData[1]) + } + + @Test + fun fp16_reports_its_codec_and_dtype() { + val data = Fp16DenseTensorData.fromFloatArray(Shape(2), floatArrayOf(1f, 2f)) + assertEquals(Fp16Codec, data.codec) + assertEquals(FP16, data.codec.dtype) + assertEquals(2, data.codec.bytesPerElement) + } + + @Test + fun bf16_reports_its_codec_and_dtype() { + val data = Bf16DenseTensorData.fromFloatArray(Shape(2), floatArrayOf(1f, 2f)) + assertEquals(Bf16Codec, data.codec) + assertEquals(BF16, data.codec.dtype) + } + + @Test + fun bf16_dense_is_still_recognized_as_both_supertypes() { + // Dispatch sites test `is Bf16TensorData`; narrow-float kernels test `is NarrowFloatTensorData`. + // Re-parenting Bf16DenseTensorData must keep both working. Erased to Any so these are real + // runtime checks rather than statically-known-true ones. + val data: Any = Bf16DenseTensorData.fromFloatArray(Shape(2), floatArrayOf(1f, 2f)) + assertTrue(data is Bf16TensorData, "must still satisfy the legacy BF16 dispatch check") + assertTrue(data is NarrowFloatTensorData, "must satisfy the shared narrow-float check") + } + + @Test + fun fp16_is_narrow_but_not_bf16() { + val data: Any = Fp16DenseTensorData.fromFloatArray(Shape(2), floatArrayOf(1f, 2f)) + assertTrue(data is NarrowFloatTensorData) + assertTrue(data !is Bf16TensorData, "FP16 must not be routed to a BF16 kernel") + } + + @Test + fun the_two_formats_disagree_on_the_same_bytes() { + // Same value, different encodings — proves the codec is doing real work per format. + val v = floatArrayOf(1.0f) + val fp16 = Fp16DenseTensorData.fromFloatArray(Shape(1), v) + val bf16 = Bf16DenseTensorData.fromFloatArray(Shape(1), v) + assertNotEquals( + fp16.packedData.toList(), bf16.packedData.toList(), + "1.0 encodes as 0x3C00 in fp16 but 0x3F80 in bf16", + ) + // ...yet both decode back to the same value. + assertEquals(1.0f, fp16.copyToFloatArray()[0]) + assertEquals(1.0f, bf16.copyToFloatArray()[0]) + } + + @Test + fun generic_factory_honours_the_supplied_codec() { + val values = floatArrayOf(1.5f, -0.75f, 256.0f) + val asFp16 = NarrowFloatDenseTensorData.fromFloatArray(Shape(3), values, Fp16Codec) + val asBf16 = NarrowFloatDenseTensorData.fromFloatArray(Shape(3), values, Bf16Codec) + + // All three are exactly representable in both formats. + for (i in values.indices) { + assertEquals(values[i], asFp16.copyToFloatArray()[i]) + assertEquals(values[i], asBf16.copyToFloatArray()[i]) + } + assertEquals(Fp16Codec, asFp16.codec) + assertEquals(Bf16Codec, asBf16.codec) + } + + @Test + fun get_and_set_decode_and_re_encode() { + val data = Fp16DenseTensorData.fromFloatArray(Shape(2, 2), floatArrayOf(1f, 2f, 3f, 4f)) + assertEquals(1.0f, data.get(0, 0)) + assertEquals(2.0f, data.get(0, 1)) + assertEquals(3.0f, data.get(1, 0)) + assertEquals(4.0f, data.get(1, 1)) + + data.set(1, 0, value = 9.0f) + assertEquals(9.0f, data.get(1, 0)) + assertEquals(4.0f, data.get(1, 1), "neighbouring element must be untouched") + } + + @Test + fun multi_dimensional_striding_is_row_major() { + val values = FloatArray(24) { it.toFloat() } + val data = Fp16DenseTensorData.fromFloatArray(Shape(2, 3, 4), values) + assertEquals(0.0f, data.get(0, 0, 0)) + assertEquals(7.0f, data.get(0, 1, 3)) + assertEquals(23.0f, data.get(1, 2, 3)) + assertEquals(values.toList(), data.copyToFloatArray().toList()) + } + + @Test + fun bulk_and_elementwise_reads_agree() { + val values = FloatArray(16) { (it - 8) * 0.5f } + val data = Fp16DenseTensorData.fromFloatArray(Shape(16), values) + val bulk = data.copyToFloatArray() + for (i in values.indices) { + assertEquals(bulk[i], data.get(i), "bulk vs elementwise at $i") + } + } + + @Test + fun signed_zero_is_preserved_bit_exactly() { + val data = Fp16DenseTensorData.fromFloatArray(Shape(2), floatArrayOf(0.0f, -0.0f)) + assertEquals(0x0000, (data.packedData[1].toInt() and 0xFF shl 8) or (data.packedData[0].toInt() and 0xFF)) + assertEquals(0x8000, (data.packedData[3].toInt() and 0xFF shl 8) or (data.packedData[2].toInt() and 0xFF)) + } + + @Test + fun fp16_saturates_to_infinity_beyond_its_range() { + // The practical difference from bf16: 70000 does not fit in binary16. + val data = Fp16DenseTensorData.fromFloatArray(Shape(1), floatArrayOf(70000f)) + assertTrue(data.copyToFloatArray()[0].isInfinite(), "fp16 overflows where bf16 would not") + + val bf16 = Bf16DenseTensorData.fromFloatArray(Shape(1), floatArrayOf(70000f)) + assertTrue(bf16.copyToFloatArray()[0].isFinite(), "bf16 keeps f32 exponent range") + } + + @Test + fun undersized_buffers_are_rejected() { + assertFailsWith { + Fp16DenseTensorData(Shape(4), ByteArray(6)) // needs 8 bytes + } + assertFailsWith { + NarrowFloatDenseTensorData(Shape(4), ByteArray(6), Bf16Codec) + } + } + + @Test + fun out_of_bounds_and_wrong_arity_indices_are_rejected() { + val data = Fp16DenseTensorData.fromFloatArray(Shape(2, 2), floatArrayOf(1f, 2f, 3f, 4f)) + assertFailsWith { data.get(2, 0) } + assertFailsWith { data.get(0) } + assertFailsWith { data.get(0, 0, 0) } + } + + @Test + fun fp16_precision_beats_bf16_on_the_same_tensor() { + val values = floatArrayOf(1.1f, 2.2f, 3.3f, 4.4f, 5.5f) + val fp16 = Fp16DenseTensorData.fromFloatArray(Shape(5), values).copyToFloatArray() + val bf16 = Bf16DenseTensorData.fromFloatArray(Shape(5), values).copyToFloatArray() + + var fp16Err = 0.0 + var bf16Err = 0.0 + for (i in values.indices) { + fp16Err += abs(fp16[i] - values[i]).toDouble() + bf16Err += abs(bf16[i] - values[i]).toDouble() + } + assertTrue(fp16Err < bf16Err, "fp16 err=$fp16Err should beat bf16 err=$bf16Err") + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/NarrowFloatStorageDecodeTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/NarrowFloatStorageDecodeTest.kt new file mode 100644 index 00000000..5122acee --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/storage/NarrowFloatStorageDecodeTest.kt @@ -0,0 +1,93 @@ +package sk.ainet.lang.tensor.storage + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.Fp16Codec +import sk.ainet.lang.types.NarrowFloatCodec +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * Regression guard for the narrow-float width bug in [TensorStorageFactory.toTensorData]. + * + * `FLOAT16`/`BFLOAT16` used to share the `FLOAT32` branch and were decoded by a hard 4-byte + * reader. For an N-element narrow tensor that produced N/2 elements, each one assembled from two + * adjacent narrow values — silently wrong data rather than an error. Readers already tagged these + * as `TensorEncoding.Dense(bytesPerElement = 2)`; only the decode side ignored it. + */ +class NarrowFloatStorageDecodeTest { + + private fun packed(values: FloatArray, codec: NarrowFloatCodec): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = codec.encode(values[i]) + out[i * 2] = (bits and 0xFF).toByte() + out[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return out + } + + private fun storageOf(values: FloatArray, codec: NarrowFloatCodec): TensorStorage { + val bytes = packed(values, codec) + return TensorStorage( + shape = Shape(values.size), + logicalType = if (codec === Fp16Codec) LogicalDType.FLOAT16 else LogicalDType.BFLOAT16, + encoding = TensorEncoding.Dense(2), + buffer = BufferHandle.Owned(bytes, 0, bytes.size.toLong()), + ) + } + + @Test + fun bf16_dense_storage_decodes_at_two_bytes_per_element() { + val values = floatArrayOf(1.0f, -2.0f, 0.5f, 3.0f, -0.25f, 8.0f) + val data: TensorData = TensorStorageFactory.toTensorData(storageOf(values, Bf16Codec)) + + val out = data.copyToFloatArray() + assertEquals(values.size, out.size, "element count must not halve") + for (i in values.indices) { + assertEquals(values[i], out[i], "element $i") + } + } + + @Test + fun fp16_dense_storage_decodes_at_two_bytes_per_element() { + val values = floatArrayOf(1.0f, -2.0f, 0.5f, 3.0f, -0.25f, 8.0f) + val data: TensorData = TensorStorageFactory.toTensorData(storageOf(values, Fp16Codec)) + + val out = data.copyToFloatArray() + assertEquals(values.size, out.size, "element count must not halve") + for (i in values.indices) { + assertEquals(values[i], out[i], "element $i") + } + } + + @Test + fun odd_element_counts_survive_the_narrow_decode() { + // An odd count is where a 4-byte reader loses the trailing element entirely. + val values = floatArrayOf(1.0f, 2.0f, 4.0f, 8.0f, 16.0f) + val out = TensorStorageFactory.toTensorData(storageOf(values, Bf16Codec)) + .copyToFloatArray() + assertEquals(5, out.size) + assertEquals(16.0f, out[4], "trailing element must not be dropped") + } + + @Test + fun bf16_and_fp16_decode_differently_for_the_same_bytes() { + // Proves the codec is actually selected by logicalType rather than one being aliased + // to the other: the bit pattern 0x3C00 is 1.0 in fp16 but ~0.0078 in bf16. + val bytes = byteArrayOf(0x00, 0x3C) + val asFp16 = TensorStorageFactory.toTensorData( + TensorStorage(Shape(1), LogicalDType.FLOAT16, TensorEncoding.Dense(2), BufferHandle.Owned(bytes, 0, 2)), + ).copyToFloatArray()[0] + val asBf16 = TensorStorageFactory.toTensorData( + TensorStorage(Shape(1), LogicalDType.BFLOAT16, TensorEncoding.Dense(2), BufferHandle.Owned(bytes, 0, 2)), + ).copyToFloatArray()[0] + + assertEquals(1.0f, asFp16) + assertTrue(asBf16 < 0.01f && asBf16 > 0.0f, "bf16 reading of 0x3C00 should be ~0.0078, got $asBf16") + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/NarrowFloatCodecTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/NarrowFloatCodecTest.kt new file mode 100644 index 00000000..865fddb2 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/NarrowFloatCodecTest.kt @@ -0,0 +1,241 @@ +package sk.ainet.lang.types + +import sk.ainet.lang.tensor.data.Bf16TensorData +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * Bit-level contract for [Bf16Codec] and [Fp16Codec]. + * + * The FP16 cases pin IEEE 754 binary16 semantics — normals, gradual underflow, the 65504 ceiling, + * ties-to-even, and NaN preservation. The BF16 cases pin *truncation* and bit-exact agreement with + * the pre-existing `Bf16TensorData.floatToBf16Bits`, which the emitted weight archives depend on. + */ +class NarrowFloatCodecTest { + + // ---------------------------------------------------------------- BF16 + + @Test + fun bf16_is_the_high_half_of_the_float() { + assertEquals(0x3F80, Bf16Codec.encode(1.0f)) + assertEquals(0xC000, Bf16Codec.encode(-2.0f)) + assertEquals(0x0000, Bf16Codec.encode(0.0f)) + assertEquals(0x8000, Bf16Codec.encode(-0.0f)) + } + + @Test + fun bf16_truncates_rather_than_rounds() { + // 1.0f + 1 ulp(f32) has mantissa bits below the bf16 cut; truncation must drop them + // (round-to-nearest would also give 0x3F80 here, so use a value that discriminates: + // a mantissa whose dropped bits exceed half — truncation still yields the LOWER value). + val justUnderTwo = Float.fromBits(0x3FFFFFFF) // 1.9999999… + assertEquals(0x3FFF, Bf16Codec.encode(justUnderTwo), "truncation must not round up to 0x4000") + } + + @Test + fun bf16_matches_the_legacy_floatToBf16Bits_bit_for_bit() { + // Regression guard: the on-device A/B that qualified bf16 as a drop-in for the f16 vmfb + // depends on this exact encoding. Any divergence silently changes every weight archive. + val samples = floatArrayOf( + 0.0f, -0.0f, 1.0f, -1.0f, 2.0f, 0.5f, 3.14159265f, -2.71828f, + 1e-8f, -1e-8f, 1e8f, -1e8f, 65504f, 1.0e-38f, 6.7e-5f, + Float.MAX_VALUE, Float.MIN_VALUE, Float.POSITIVE_INFINITY, Float.NEGATIVE_INFINITY, + ) + for (v in samples) { + assertEquals( + Bf16TensorData.floatToBf16Bits(v), Bf16Codec.encode(v), + "encode mismatch for $v", + ) + assertEquals( + Bf16TensorData.bf16BitsToFloat(Bf16Codec.encode(v)), Bf16Codec.decode(Bf16Codec.encode(v)), + "decode mismatch for $v", + ) + } + } + + @Test + fun bf16_keeps_f32_exponent_range() { + // The whole point of bf16 over fp16: no overflow on conversion. + assertTrue(Bf16Codec.decode(Bf16Codec.encode(1e38f)).isFinite()) + assertTrue(Bf16Codec.decode(Bf16Codec.encode(-1e38f)).isFinite()) + } + + @Test + fun bf16_round_trips_within_its_precision() { + for (v in floatArrayOf(1.0f, -3.5f, 1234.0f, 0.001f, -7.25e10f)) { + val back = Bf16Codec.decode(Bf16Codec.encode(v)) + // bf16 has 8 significand bits -> relative step 2^-8; truncation doubles the bound. + assertTrue(abs(back - v) <= abs(v) * 0.008f, "bf16 round-trip $v -> $back") + } + } + + // ---------------------------------------------------------------- FP16 normals + + @Test + fun fp16_encodes_known_normal_patterns() { + assertEquals(0x3C00, Fp16Codec.encode(1.0f)) + assertEquals(0xC000, Fp16Codec.encode(-2.0f)) + assertEquals(0x3800, Fp16Codec.encode(0.5f)) + assertEquals(0x0000, Fp16Codec.encode(0.0f)) + assertEquals(0x8000, Fp16Codec.encode(-0.0f)) + } + + @Test + fun fp16_decodes_known_normal_patterns() { + assertEquals(1.0f, Fp16Codec.decode(0x3C00)) + assertEquals(-2.0f, Fp16Codec.decode(0xC000)) + assertEquals(0.5f, Fp16Codec.decode(0x3800)) + assertEquals(0.0f, Fp16Codec.decode(0x0000)) + assertEquals(-0.0f, Fp16Codec.decode(0x8000)) + } + + @Test + fun fp16_max_finite_is_65504() { + assertEquals(0x7BFF, Fp16Codec.encode(65504f)) + assertEquals(65504f, Fp16Codec.decode(0x7BFF)) + } + + // ---------------------------------------------------------------- FP16 overflow + + @Test + fun fp16_overflows_to_infinity_past_the_ceiling() { + // Half an ulp above max finite (ulp at 65504 is 32) is the round-to-Inf threshold. + assertEquals(0x7C00, Fp16Codec.encode(65520f), "65520 is the tie -> even -> Inf") + assertEquals(0x7C00, Fp16Codec.encode(70000f)) + assertEquals(0xFC00, Fp16Codec.encode(-70000f)) + assertEquals(0x7C00, Fp16Codec.encode(1e30f), "far beyond range") + assertTrue(Fp16Codec.decode(Fp16Codec.encode(70000f)).isInfinite()) + } + + @Test + fun fp16_just_below_the_threshold_stays_finite() { + // 65519 rounds down to max finite, not up to Inf. + assertEquals(0x7BFF, Fp16Codec.encode(65519f)) + } + + // ---------------------------------------------------------------- FP16 subnormals + + @Test + fun fp16_smallest_normal_and_largest_subnormal() { + val smallestNormal = 1.0f / 16384.0f // 2^-14 + assertEquals(0x0400, Fp16Codec.encode(smallestNormal)) + assertEquals(smallestNormal, Fp16Codec.decode(0x0400)) + + val largestSubnormal = 1023.0f / 16777216.0f // 1023 * 2^-24 + assertEquals(0x03FF, Fp16Codec.encode(largestSubnormal)) + assertEquals(largestSubnormal, Fp16Codec.decode(0x03FF)) + } + + @Test + fun fp16_smallest_subnormal_is_2_pow_minus_24() { + val smallestSubnormal = 1.0f / 16777216.0f // 2^-24 + assertEquals(0x0001, Fp16Codec.encode(smallestSubnormal)) + assertEquals(smallestSubnormal, Fp16Codec.decode(0x0001)) + } + + @Test + fun fp16_underflows_to_zero_with_ties_to_even() { + val half = 1.0f / 33554432.0f // 2^-25 — exactly half the smallest subnormal + assertEquals(0x0000, Fp16Codec.encode(half), "exact tie rounds to even, i.e. zero") + assertEquals(0x8000, Fp16Codec.encode(-half), "sign is preserved on underflow") + + // Just above the tie must round up to the smallest subnormal. + val justAbove = Float.fromBits(half.toRawBits() + 1) + assertEquals(0x0001, Fp16Codec.encode(justAbove)) + + // Float subnormals are far below binary16's range and must not corrupt the exponent math. + assertEquals(0x0000, Fp16Codec.encode(Float.MIN_VALUE)) + assertEquals(0x8000, Fp16Codec.encode(-Float.MIN_VALUE)) + } + + @Test + fun fp16_subnormals_decode_to_exact_values() { + // Every binary16 subnormal is m * 2^-24 exactly, and f32 represents all of them exactly. + for (m in 1..1023) { + val expected = m.toFloat() / 16777216.0f + assertEquals(expected, Fp16Codec.decode(m), "subnormal m=$m") + assertEquals(m, Fp16Codec.encode(expected), "subnormal round-trip m=$m") + } + } + + // ---------------------------------------------------------------- FP16 rounding + + @Test + fun fp16_rounds_ties_to_even() { + // ulp at 1.0 is 2^-10, so 1 + 2^-11 sits exactly between 0x3C00 and 0x3C01. + val tieDown = 1.0f + 1.0f / 2048.0f + assertEquals(0x3C00, Fp16Codec.encode(tieDown), "tie with even LSB rounds down") + + // 1 + 3*2^-11 sits exactly between 0x3C01 (odd) and 0x3C02 -> rounds up to even. + val tieUp = 1.0f + 3.0f / 2048.0f + assertEquals(0x3C02, Fp16Codec.encode(tieUp), "tie with odd LSB rounds up") + } + + @Test + fun fp16_rounds_to_nearest_off_ties() { + assertEquals(0x3C01, Fp16Codec.encode(1.0f + 1.0f / 1024.0f), "exactly representable") + assertEquals(0x3C01, Fp16Codec.encode(1.0f + 0.9f / 1024.0f), "nearer the upper neighbour") + assertEquals(0x3C00, Fp16Codec.encode(1.0f + 0.1f / 1024.0f), "nearer the lower neighbour") + } + + // ---------------------------------------------------------------- FP16 inf / NaN + + @Test + fun fp16_preserves_infinities() { + assertEquals(0x7C00, Fp16Codec.encode(Float.POSITIVE_INFINITY)) + assertEquals(0xFC00, Fp16Codec.encode(Float.NEGATIVE_INFINITY)) + assertTrue(Fp16Codec.decode(0x7C00) == Float.POSITIVE_INFINITY) + assertTrue(Fp16Codec.decode(0xFC00) == Float.NEGATIVE_INFINITY) + } + + @Test + fun fp16_nan_never_collapses_into_infinity() { + val encoded = Fp16Codec.encode(Float.NaN) + assertTrue(Fp16Codec.decode(encoded).isNaN(), "NaN must survive the round trip") + assertTrue((encoded and 0x03FF) != 0, "payload must be non-zero or it would decode as Inf") + + // A NaN whose payload lives entirely in the low mantissa bits still must not become Inf. + val lowPayloadNaN = Float.fromBits(0x7F80_0001) + val encodedLow = Fp16Codec.encode(lowPayloadNaN) + assertTrue(Fp16Codec.decode(encodedLow).isNaN(), "low-payload NaN must survive") + } + + // ---------------------------------------------------------------- shared contract + + @Test + fun codecs_report_their_dtype_and_width() { + assertEquals(BF16, Bf16Codec.dtype) + assertEquals(FP16, Fp16Codec.dtype) + assertEquals(2, Bf16Codec.bytesPerElement) + assertEquals(2, Fp16Codec.bytesPerElement) + } + + @Test + fun fp16_is_more_precise_than_bf16_in_range() { + // fp16 has three more mantissa bits, so it can never be worse — but for individual values + // the two can coincide exactly (at pi both land on 3.140625), so assert over a sample: + // never worse pointwise, and strictly better in aggregate. + val samples = floatArrayOf( + 1.1f, 3.14159265f, 2.71828f, 0.1f, 123.456f, -7.77f, 1000.5f, 0.007f, + ) + var fp16Total = 0.0 + var bf16Total = 0.0 + for (v in samples) { + val fp16Err = abs(Fp16Codec.decode(Fp16Codec.encode(v)) - v) + val bf16Err = abs(Bf16Codec.decode(Bf16Codec.encode(v)) - v) + assertTrue(fp16Err <= bf16Err, "fp16 must never be worse: $v -> fp16=$fp16Err bf16=$bf16Err") + fp16Total += fp16Err.toDouble() + bf16Total += bf16Err.toDouble() + } + assertTrue(fp16Total < bf16Total, "fp16 total err=$fp16Total should beat bf16 total=$bf16Total") + } + + @Test + fun encode_ignores_bits_above_the_low_sixteen_on_decode() { + // decode must mask, so callers can pass a sign-extended Short without corruption. + assertEquals(Fp16Codec.decode(0x3C00), Fp16Codec.decode(0x3C00 or 0xFFFF_0000.toInt())) + assertEquals(Bf16Codec.decode(0x3F80), Bf16Codec.decode(0x3F80 or 0xFFFF_0000.toInt())) + } +} From 87a28fce7a4d8c542e0e1e44064ded46874dc75f Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 26 Jul 2026 18:33:13 +0200 Subject: [PATCH 2/4] feat(io): narrow-float KEEP_NATIVE for SafeTensors F16 and GGUF F16/BF16 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second stage (B3) of #884. Both loaders can now hand 16-bit weights through in their on-disk 2-bytes-per-element layout instead of widening every one to FP32 at load. Builds on the Fp16DenseTensorData added in B1. SafeTensors: - F16 gains the KEEP_NATIVE arm that BF16 already had. It was absent purely because no Fp16DenseTensorData existed to back it. - `mapPolicyToNarrow(policy, native)` replaces the two hand-written mappings, so one function resolves the policy for either format. - `Bf16LoadPolicy` becomes a typealias to the new `NarrowFloatLoadPolicy`, with `toDTypePolicy()` moved to an extension function. Existing call sites and tests compile unchanged. GGUF: - F16 and BF16 stop unconditionally widening; `Require(BF16)` / `Require(FP16)` are accepted instead of throwing. This is the KEEP_NATIVE GGUF path the CHANGELOG had parked, and it is what makes Require(BF16) real on GGUF. - `keepsNative(policy, native)` is the single decision point. A policy naming one narrow format WIDENS the other rather than keeping it packed. This matters more than it looks: the two encodings are different bit layouts, so mis-tagging F16 bytes as BF16 would decode to silent garbage rather than failing. Enforced by `policy.target == native` in both loaders and covered by a dedicated test on each side. Three pre-existing tests asserted the old limitations (`require_fp16_fails_with_explicit_message`, and GGUF's two `require_*_fails_fast_with_clear_message`). Rewritten to assert the new capability — the behaviour they guarded is gone by design, not by accident. Tests: 4 new SafeTensors FP16 integration tests (packed bytes preserved verbatim, decode bit-identical to the widening path, the two policies independent on a mixed F16+BF16 file) plus 5 policy-mapping tests across both loaders. Existing BF16 suites pass unmodified. apiDump regenerated; full engine build clean (2133 tasks). Refs #884 Co-Authored-By: Claude Opus 5 (1M context) --- .../io/gguf/StreamingGgufParametersLoader.kt | 77 +++++--- ...StreamingGgufParametersLoaderPolicyTest.kt | 34 ++-- .../sk/ainet/io/safetensors/Bf16LoadPolicy.kt | 60 ++---- .../io/safetensors/NarrowFloatLoadPolicy.kt | 40 ++++ .../SafeTensorsParametersLoader.kt | 71 ++++--- .../SafeTensorsParametersLoaderPolicyTest.kt | 44 ++++- ...feTensorsParametersLoaderFp16PolicyTest.kt | 184 ++++++++++++++++++ 7 files changed, 393 insertions(+), 117 deletions(-) create mode 100644 skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/NarrowFloatLoadPolicy.kt create mode 100644 skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderFp16PolicyTest.kt diff --git a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt index 04d32f42..d1ae2937 100644 --- a/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt +++ b/skainet-io/skainet-io-gguf/src/commonMain/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoader.kt @@ -5,6 +5,8 @@ import sk.ainet.io.ParametersLoader import sk.ainet.io.RandomAccessSource import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Bf16DenseTensorData +import sk.ainet.lang.tensor.data.Fp16DenseTensorData import sk.ainet.lang.tensor.data.Q4_KBlockTensorData import sk.ainet.lang.tensor.data.Q5_KBlockTensorData import sk.ainet.lang.tensor.data.Q6_KBlockTensorData @@ -33,7 +35,16 @@ import kotlin.reflect.KClass */ public class StreamingGgufParametersLoader( private val sourceProvider: () -> RandomAccessSource, - private val onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> } + private val onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> }, + /** + * Keep `F16` source tensors in their on-disk 2-bytes-per-element layout instead of widening + * them to FP32 at load. Off by default — flip via `withPolicy(Require(FP16))`. + */ + private val keepF16Native: Boolean = false, + /** + * Keep `BF16` source tensors packed. Off by default — flip via `withPolicy(Require(BF16))`. + */ + private val keepBf16Native: Boolean = false, ) : ParametersLoader { @Suppress("UNCHECKED_CAST") @@ -68,20 +79,28 @@ public class StreamingGgufParametersLoader( } } - GGMLQuantizationType.F16 -> { - val floats = dequantF16(rawBytes) - when (dtype) { - FP32::class -> ctx.fromFloatArray(shape, dtype, floats) as Tensor - else -> null + GGMLQuantizationType.F16 -> when (dtype) { + FP32::class -> if (keepF16Native) { + // Zero-widening path: hand the on-disk bytes straight through as + // packed binary16. Consumers still see Float on read. + @Suppress("UNCHECKED_CAST") + val packed = Fp16DenseTensorData(shape, rawBytes) + ctx.fromData(packed as sk.ainet.lang.tensor.data.TensorData, dtype) + } else { + ctx.fromFloatArray(shape, dtype, dequantF16(rawBytes)) as Tensor } + else -> null } - GGMLQuantizationType.BF16 -> { - val floats = dequantBF16(rawBytes) - when (dtype) { - FP32::class -> ctx.fromFloatArray(shape, dtype, floats) as Tensor - else -> null + GGMLQuantizationType.BF16 -> when (dtype) { + FP32::class -> if (keepBf16Native) { + @Suppress("UNCHECKED_CAST") + val packed = Bf16DenseTensorData(shape, rawBytes) + ctx.fromData(packed as sk.ainet.lang.tensor.data.TensorData, dtype) + } else { + ctx.fromFloatArray(shape, dtype, dequantBF16(rawBytes)) as Tensor } + else -> null } GGMLQuantizationType.Q4_K -> { @@ -202,7 +221,25 @@ public class StreamingGgufParametersLoader( onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> }, ): StreamingGgufParametersLoader { validatePolicy(policy) - return StreamingGgufParametersLoader(sourceProvider, onProgress) + return StreamingGgufParametersLoader( + sourceProvider = sourceProvider, + onProgress = onProgress, + keepF16Native = keepsNative(policy, FP16), + keepBf16Native = keepsNative(policy, BF16), + ) + } + + /** + * Whether [policy] asks for [native] tensors to stay in their on-disk 16-bit layout. + * + * Only the format the policy actually names is kept — neither narrow format can be turned + * into the other without a lossy re-encode, so `Require(BF16)` must still widen F16 sources. + */ + internal fun keepsNative(policy: DTypePolicy, native: DType): Boolean = when (policy) { + DTypePolicy.Any -> false + is DTypePolicy.Require -> policy.target == native + is DTypePolicy.Prefer -> policy.target == native + is DTypePolicy.OneOf -> native in policy.allowed } internal fun validatePolicy(policy: DTypePolicy) { @@ -211,18 +248,10 @@ public class StreamingGgufParametersLoader( is DTypePolicy.Prefer -> Unit is DTypePolicy.OneOf -> Unit is DTypePolicy.Require -> when (policy.target) { - FP32 -> Unit - BF16 -> throw IllegalArgumentException( - "StreamingGgufParametersLoader: Require(BF16) is not supported — " + - "GGUF BF16 sources are dequanted to FP32 by this loader today (no KEEP_NATIVE " + - "GGUF path yet). Use Any or Prefer(BF16) to accept the dequant fallback, or " + - "wait for the policy-aware GGUF reader to land.", - ) - FP16 -> throw IllegalArgumentException( - "StreamingGgufParametersLoader: Require(FP16) is not supported — " + - "GGUF F16 sources are dequanted to FP32 by this loader today (no Fp16DenseTensorData " + - "backing yet). Use Any or Prefer(FP16) to accept the dequant fallback.", - ) + // FP16 / BF16 are satisfiable for sources already in that format: the loader + // hands the packed bytes through via Fp16/Bf16DenseTensorData. Sources in any + // other format still widen to FP32 rather than being re-encoded. + FP32, FP16, BF16 -> Unit else -> throw IllegalArgumentException( "StreamingGgufParametersLoader: Require(${policy.target.name}) is not satisfiable — " + "this loader produces FP32 / Int32 / Q4_K / Q8_0 tensors only, and does not cast " + diff --git a/skainet-io/skainet-io-gguf/src/commonTest/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoaderPolicyTest.kt b/skainet-io/skainet-io-gguf/src/commonTest/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoaderPolicyTest.kt index bf32f5a0..f496b252 100644 --- a/skainet-io/skainet-io-gguf/src/commonTest/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoaderPolicyTest.kt +++ b/skainet-io/skainet-io-gguf/src/commonTest/kotlin/sk/ainet/io/gguf/StreamingGgufParametersLoaderPolicyTest.kt @@ -30,21 +30,31 @@ class StreamingGgufParametersLoaderPolicyTest { } @Test - fun require_bf16_fails_fast_with_clear_message() { - val ex = assertFailsWith { - StreamingGgufParametersLoader.validatePolicy(DTypePolicy.Require(BF16)) - } - val msg = ex.message ?: "" - assertTrue(msg.contains("Require(BF16)"), msg) - assertTrue(msg.contains("KEEP_NATIVE"), msg) + fun require_bf16_is_now_accepted_and_keeps_bf16_sources_packed() { + // Previously rejected: the GGUF loader had no KEEP_NATIVE path and always widened. + StreamingGgufParametersLoader.validatePolicy(DTypePolicy.Require(BF16)) + assertTrue(StreamingGgufParametersLoader.keepsNative(DTypePolicy.Require(BF16), BF16)) } @Test - fun require_fp16_fails_fast_with_clear_message() { - val ex = assertFailsWith { - StreamingGgufParametersLoader.validatePolicy(DTypePolicy.Require(FP16)) - } - assertTrue(ex.message?.contains("Require(FP16)") == true, ex.message ?: "") + fun require_fp16_is_now_accepted_and_keeps_f16_sources_packed() { + // Previously rejected for want of an Fp16DenseTensorData backing, which now exists. + StreamingGgufParametersLoader.validatePolicy(DTypePolicy.Require(FP16)) + assertTrue(StreamingGgufParametersLoader.keepsNative(DTypePolicy.Require(FP16), FP16)) + } + + @Test + fun neither_narrow_format_is_kept_native_for_the_other() { + // Converting between the two is a lossy re-encode, so a policy naming one must widen + // sources in the other rather than mis-tagging their bytes. + assertTrue(!StreamingGgufParametersLoader.keepsNative(DTypePolicy.Require(BF16), FP16)) + assertTrue(!StreamingGgufParametersLoader.keepsNative(DTypePolicy.Require(FP16), BF16)) + } + + @Test + fun any_policy_still_widens_both_narrow_formats() { + assertTrue(!StreamingGgufParametersLoader.keepsNative(DTypePolicy.Any, BF16)) + assertTrue(!StreamingGgufParametersLoader.keepsNative(DTypePolicy.Any, FP16)) } @Test diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/Bf16LoadPolicy.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/Bf16LoadPolicy.kt index 016a9669..df54a26e 100644 --- a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/Bf16LoadPolicy.kt +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/Bf16LoadPolicy.kt @@ -27,53 +27,17 @@ import sk.ainet.lang.types.FP32 * `DefaultCpuOpsJvm.chooseQuantizedMatmul` (Phase 3 follow-up, * separately tracked). */ -public enum class Bf16LoadPolicy { - /** - * Default. Dequantise every BFLOAT16 tensor to FP32 at load time - * via the existing `dequantBF16` helper, then wrap as a - * `FloatArrayTensorData` (same as `FLOAT32` source tensors). - * - * Memory cost: 2× the on-disk size for each BF16 tensor. - * Runtime: zero extra dispatch — every matmul gets FP32 operands. - */ - DEQUANT_TO_FP32, +public typealias Bf16LoadPolicy = NarrowFloatLoadPolicy - /** - * Keep BFLOAT16 tensors in their on-disk packed-2-bytes-per-element - * layout. The loader emits a `Bf16DenseTensorData` (in - * `skainet-lang-core`) instead of dequanting; the tensor still - * advertises FP32 dtype to consumers (the underlying `get` decodes - * on read), but its `tensor.data` is recognisable as - * `Bf16TensorData` so a matmul dispatch can route to the SIMD - * `Bf16MatmulKernel` SPI. - * - * Memory cost: identical to the on-disk size — no doubling. - * Runtime: matmul dispatch picks up the BF16 SPI kernel when one - * is registered; falls back to per-multiply dequant otherwise. - * - * **Caveat**: any non-matmul op that touches the BF16 tensor pays - * a per-element decode cost via `get`. Don't flip this unless the - * model's hot path is dominated by matmuls (the typical - * transformer case). - */ - KEEP_NATIVE, - ; - - /** - * Maps this BF16-specific enum onto the generalised - * [DTypePolicy] sealed type. [DEQUANT_TO_FP32] becomes - * `Require(FP32)` (the loader must hand consumers an FP32 - * tensor); [KEEP_NATIVE] becomes `Require(BF16)` (consumers - * dispatch on the native BF16 dtype). - * - * Bridge for the RFC's policy-driven loader work - * (`rfc.md`, issue #615): existing call sites keep using this - * enum verbatim while new code paths can flow through - * [DTypePolicy] uniformly. The two are equivalent for BF16 — - * this method is the explicit equivalence proof. - */ - public fun toDTypePolicy(): DTypePolicy = when (this) { - DEQUANT_TO_FP32 -> DTypePolicy.Require(FP32) - KEEP_NATIVE -> DTypePolicy.Require(BF16) - } +/** + * Maps this narrow-float policy onto the generalised [DTypePolicy] sealed type. + * [NarrowFloatLoadPolicy.DEQUANT_TO_FP32] becomes `Require(FP32)` (the loader must hand consumers + * an FP32 tensor); [NarrowFloatLoadPolicy.KEEP_NATIVE] becomes `Require(BF16)`. + * + * Kept BF16-specific for source compatibility with existing call sites. For FP16, build the + * policy directly (`DTypePolicy.Require(FP16)`) — see `SafeTensorsParametersLoader.mapPolicyToFp16`. + */ +public fun Bf16LoadPolicy.toDTypePolicy(): DTypePolicy = when (this) { + NarrowFloatLoadPolicy.DEQUANT_TO_FP32 -> DTypePolicy.Require(FP32) + NarrowFloatLoadPolicy.KEEP_NATIVE -> DTypePolicy.Require(BF16) } diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/NarrowFloatLoadPolicy.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/NarrowFloatLoadPolicy.kt new file mode 100644 index 00000000..55d9febc --- /dev/null +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/NarrowFloatLoadPolicy.kt @@ -0,0 +1,40 @@ +package sk.ainet.io.safetensors + +/** + * Controls how a loader handles **16-bit float** source tensors (`FLOAT16` / `BFLOAT16`). + * + * The choice is whether to **widen to FP32 at load time** — doubling the resident footprint of + * every narrow tensor — or to **keep the on-disk bytes** and let the matmul dispatch route to a + * narrow-float kernel that widens per-multiply instead. + * + * Both paths are numerically equivalent when they reach the same kernel: the narrow → FP32 + * conversion is exact in either direction (FP32 is a strict superset of both formats), applied + * either once at load ([DEQUANT_TO_FP32]) or per-multiply ([KEEP_NATIVE]). + * + * This supersedes the BF16-only `Bf16LoadPolicy`, which remains as a typealias. + */ +public enum class NarrowFloatLoadPolicy { + /** + * Default. Widen every narrow-float tensor to FP32 at load time and wrap it as + * `FloatArrayTensorData`, exactly as a `FLOAT32` source tensor would be. + * + * Memory: 2× the on-disk size for each narrow tensor. Runtime: no extra dispatch — every + * matmul receives FP32 operands. + */ + DEQUANT_TO_FP32, + + /** + * Keep narrow-float tensors in their on-disk packed 2-bytes-per-element layout. The loader + * emits a `NarrowFloatDenseTensorData` (`Bf16DenseTensorData` / `Fp16DenseTensorData`) rather + * than widening. The tensor still presents `Float` on read — `get` decodes — but its + * `tensor.data` is recognisable as `NarrowFloatTensorData`, which is what lets a matmul + * dispatch pick the narrow-float SPI kernel. + * + * Memory: identical to the on-disk size — no doubling. Runtime: the narrow kernel is used when + * one is registered; otherwise consumers fall back to per-element decode. + * + * **Caveat**: any non-matmul op touching the tensor pays a per-element decode via `get`. Worth + * it when the hot path is matmul-dominated (the typical transformer case), not otherwise. + */ + KEEP_NATIVE, +} diff --git a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt index 2d54ed93..ef89d4e9 100644 --- a/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt +++ b/skainet-io/skainet-io-safetensors/src/commonMain/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoader.kt @@ -7,6 +7,7 @@ import sk.ainet.io.model.DataType import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.data.Bf16DenseTensorData +import sk.ainet.lang.tensor.data.Fp16DenseTensorData import sk.ainet.lang.tensor.data.TensorData import sk.ainet.lang.types.BF16 import sk.ainet.lang.types.DType @@ -46,7 +47,8 @@ import kotlin.reflect.KClass class SafeTensorsParametersLoader( private val sourceProvider: () -> RandomAccessSource, private val onProgress: (current: Long, total: Long, message: String?) -> Unit = { _, _, _ -> }, - private val bf16Policy: Bf16LoadPolicy = Bf16LoadPolicy.DEQUANT_TO_FP32, + private val bf16Policy: Bf16LoadPolicy = NarrowFloatLoadPolicy.DEQUANT_TO_FP32, + private val fp16Policy: NarrowFloatLoadPolicy = NarrowFloatLoadPolicy.DEQUANT_TO_FP32, ) : ParametersLoader { override suspend fun load( @@ -86,10 +88,22 @@ class SafeTensorsParametersLoader( DataType.FLOAT16 -> { require(dtype == FP32::class) { - "SafeTensors F16 tensor '${tensorInfo.name}' requires FP32 dtype (dequant), got ${dtype.simpleName}" + "SafeTensors F16 tensor '${tensorInfo.name}' requires FP32 dtype, got ${dtype.simpleName}" + } + when (fp16Policy) { + NarrowFloatLoadPolicy.DEQUANT_TO_FP32 -> { + val floats = dequantF16(bytes) + ctx.wrapFloatArray(shape, dtype, floats) as Tensor + } + NarrowFloatLoadPolicy.KEEP_NATIVE -> { + // Mirrors the BF16 arm below: wrap the on-disk F16 bytes directly. + // dtype stays FP32 from the consumer's POV (the tensor data decodes + // on read); the storage type is what a narrow-float matmul dispatch + // pattern-matches on. + val fp16Data = Fp16DenseTensorData(shape, bytes) + ctx.fromData(fp16Data as TensorData, dtype) + } } - val floats = dequantF16(bytes) - ctx.wrapFloatArray(shape, dtype, floats) as Tensor } DataType.BFLOAT16 -> { @@ -325,28 +339,37 @@ class SafeTensorsParametersLoader( sourceProvider = sourceProvider, onProgress = onProgress, bf16Policy = mapPolicyToBf16(policy), + fp16Policy = mapPolicyToFp16(policy), ) - internal fun mapPolicyToBf16(policy: DTypePolicy): Bf16LoadPolicy = when (policy) { - DTypePolicy.Any -> Bf16LoadPolicy.DEQUANT_TO_FP32 - is DTypePolicy.Require -> when (policy.target) { - BF16 -> Bf16LoadPolicy.KEEP_NATIVE - FP32 -> Bf16LoadPolicy.DEQUANT_TO_FP32 - FP16 -> throw IllegalArgumentException( - "SafeTensorsParametersLoader: Require(FP16) is not supported — " + - "F16 KEEP_NATIVE has no Fp16DenseTensorData backing yet. " + - "Use Require(FP32) to dequant F16 sources, or Any to inherit the adaptive default.", - ) - else -> throw IllegalArgumentException( - "SafeTensorsParametersLoader: Require(${policy.target.name}) is not satisfiable — " + - "the loader produces FP32 / BF16 / Int32 / Int8 tensors depending on source dtype; " + - "it cannot fabricate ${policy.target.name} from arbitrary sources.", - ) + internal fun mapPolicyToBf16(policy: DTypePolicy): Bf16LoadPolicy = + mapPolicyToNarrow(policy, BF16) + + internal fun mapPolicyToFp16(policy: DTypePolicy): NarrowFloatLoadPolicy = + mapPolicyToNarrow(policy, FP16) + + /** + * Resolve [policy] for one narrow-float source format. A tensor is kept native only when + * the policy names *that* format — `Require(BF16)` must not keep F16 tensors packed, and + * vice versa, since neither can be converted to the other without a lossy re-encode. + */ + private fun mapPolicyToNarrow(policy: DTypePolicy, native: DType): NarrowFloatLoadPolicy = + when (policy) { + DTypePolicy.Any -> NarrowFloatLoadPolicy.DEQUANT_TO_FP32 + is DTypePolicy.Require -> when (policy.target) { + native -> NarrowFloatLoadPolicy.KEEP_NATIVE + // The other narrow format, or FP32: this format still widens. + BF16, FP16, FP32 -> NarrowFloatLoadPolicy.DEQUANT_TO_FP32 + else -> throw IllegalArgumentException( + "SafeTensorsParametersLoader: Require(${policy.target.name}) is not satisfiable — " + + "the loader produces FP32 / BF16 / FP16 / Int32 / Int8 tensors depending on " + + "source dtype; it cannot fabricate ${policy.target.name} from arbitrary sources.", + ) + } + is DTypePolicy.Prefer -> if (policy.target == native) NarrowFloatLoadPolicy.KEEP_NATIVE + else NarrowFloatLoadPolicy.DEQUANT_TO_FP32 + is DTypePolicy.OneOf -> if (native in policy.allowed) NarrowFloatLoadPolicy.KEEP_NATIVE + else NarrowFloatLoadPolicy.DEQUANT_TO_FP32 } - is DTypePolicy.Prefer -> if (policy.target == BF16) Bf16LoadPolicy.KEEP_NATIVE - else Bf16LoadPolicy.DEQUANT_TO_FP32 - is DTypePolicy.OneOf -> if (BF16 in policy.allowed) Bf16LoadPolicy.KEEP_NATIVE - else Bf16LoadPolicy.DEQUANT_TO_FP32 - } } } diff --git a/skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderPolicyTest.kt b/skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderPolicyTest.kt index 7296508d..6b7143c1 100644 --- a/skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderPolicyTest.kt +++ b/skainet-io/skainet-io-safetensors/src/commonTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderPolicyTest.kt @@ -43,15 +43,41 @@ class SafeTensorsParametersLoaderPolicyTest { } @Test - fun require_fp16_fails_with_explicit_message() { - val ex = assertFailsWith { - SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Require(FP16)) - } - // The error message must point the operator at the alternative — - // RFC says "fail-fast with clear diagnostics," not just throw. - val msg = ex.message ?: "" - assertEquals(true, msg.contains("Require(FP16)"), "msg: $msg") - assertEquals(true, msg.contains("Fp16DenseTensorData"), "msg: $msg") + fun require_fp16_now_keeps_f16_sources_native() { + // Previously this threw: there was no Fp16DenseTensorData to back a KEEP_NATIVE F16 path. + // That backing now exists, so Require(FP16) is satisfiable for F16 sources. + assertEquals( + NarrowFloatLoadPolicy.KEEP_NATIVE, + SafeTensorsParametersLoader.mapPolicyToFp16(DTypePolicy.Require(FP16)), + ) + } + + @Test + fun the_two_narrow_formats_do_not_keep_each_other_native() { + // Neither format can be produced from the other without a lossy re-encode, so a policy + // naming one must widen the other rather than silently mis-tagging it. + assertEquals( + NarrowFloatLoadPolicy.DEQUANT_TO_FP32, + SafeTensorsParametersLoader.mapPolicyToFp16(DTypePolicy.Require(BF16)), + "Require(BF16) must not keep F16 sources packed", + ) + assertEquals( + NarrowFloatLoadPolicy.DEQUANT_TO_FP32, + SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Require(FP16)), + "Require(FP16) must not keep BF16 sources packed", + ) + } + + @Test + fun require_fp16_still_keeps_bf16_behaviour_intact() { + assertEquals( + NarrowFloatLoadPolicy.KEEP_NATIVE, + SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Require(BF16)), + ) + assertEquals( + NarrowFloatLoadPolicy.DEQUANT_TO_FP32, + SafeTensorsParametersLoader.mapPolicyToBf16(DTypePolicy.Require(FP32)), + ) } @Test diff --git a/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderFp16PolicyTest.kt b/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderFp16PolicyTest.kt new file mode 100644 index 00000000..0934efd5 --- /dev/null +++ b/skainet-io/skainet-io-safetensors/src/jvmTest/kotlin/sk/ainet/io/safetensors/SafeTensorsParametersLoaderFp16PolicyTest.kt @@ -0,0 +1,184 @@ +package sk.ainet.io.safetensors + +import kotlinx.coroutines.runBlocking +import org.junit.Test +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Bf16TensorData +import sk.ainet.lang.tensor.data.FloatArrayTensorData +import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.NarrowFloatTensorData +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec +import java.io.File +import java.nio.ByteBuffer +import java.nio.ByteOrder +import java.nio.file.Files +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * FP16 counterpart to [SafeTensorsParametersLoaderBf16PolicyTest]. + * + * Before `Fp16DenseTensorData` existed the loader had no choice but to widen every F16 tensor to + * FP32, and `Require(FP16)` was rejected outright. These tests pin the KEEP_NATIVE path: on-disk + * bytes preserved verbatim, values decoded on read, and BF16 handling untouched. + */ +class SafeTensorsParametersLoaderFp16PolicyTest { + + /** Encode FP32 to IEEE binary16 bytes, little-endian. */ + private fun fp32ToFp16Bytes(values: FloatArray): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = Fp16Codec.encode(values[i]) + out[i * 2] = (bits and 0xFF).toByte() + out[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return out + } + + private fun createF16SafeTensorsFile(name: String, values: FloatArray): File { + val bytes = fp32ToFp16Bytes(values) + val headerJson = + "{\"$name\": {\"dtype\": \"F16\", \"shape\": [${values.size}], \"data_offsets\": [0, ${bytes.size}]}}" + val headerBytes = headerJson.toByteArray(Charsets.UTF_8) + val tempFile = Files.createTempFile("test_f16_safetensors", ".safetensors").toFile() + tempFile.deleteOnExit() + tempFile.outputStream().use { out -> + out.write( + ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN) + .putLong(headerBytes.size.toLong()).array(), + ) + out.write(headerBytes) + out.write(bytes) + } + return tempFile + } + + /** Build a file holding one F16 tensor and one BF16 tensor, to prove the policies are independent. */ + private fun createF16AndBf16File(f16Values: FloatArray, bf16Values: FloatArray): File { + val f16Bytes = fp32ToFp16Bytes(f16Values) + val bf16Bytes = ByteArray(bf16Values.size * 2) + for (i in bf16Values.indices) { + val b = (bf16Values[i].toRawBits() ushr 16) and 0xFFFF + bf16Bytes[i * 2] = (b and 0xFF).toByte() + bf16Bytes[i * 2 + 1] = ((b ushr 8) and 0xFF).toByte() + } + val f16End = f16Bytes.size.toLong() + val bf16End = f16End + bf16Bytes.size + val headerJson = + "{\"w_f16\": {\"dtype\": \"F16\", \"shape\": [${f16Values.size}], \"data_offsets\": [0, $f16End]}," + + "\"w_bf16\": {\"dtype\": \"BF16\", \"shape\": [${bf16Values.size}], " + + "\"data_offsets\": [$f16End, $bf16End]}}" + val headerBytes = headerJson.toByteArray(Charsets.UTF_8) + val tempFile = Files.createTempFile("test_f16_bf16_safetensors", ".safetensors").toFile() + tempFile.deleteOnExit() + tempFile.outputStream().use { out -> + out.write( + ByteBuffer.allocate(8).order(ByteOrder.LITTLE_ENDIAN) + .putLong(headerBytes.size.toLong()).array(), + ) + out.write(headerBytes) + out.write(f16Bytes) + out.write(bf16Bytes) + } + return tempFile + } + + private fun loadAll( + file: File, + fp16: NarrowFloatLoadPolicy = NarrowFloatLoadPolicy.DEQUANT_TO_FP32, + bf16: NarrowFloatLoadPolicy = NarrowFloatLoadPolicy.DEQUANT_TO_FP32, + ): Map> = runBlocking { + val ctx = DirectCpuExecutionContext.create() + val loader = SafeTensorsParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(file) }, + bf16Policy = bf16, + fp16Policy = fp16, + ) + val out = mutableMapOf>() + loader.load(ctx, FP32::class) { name, tensor -> out[name] = tensor } + out + } + + @Test + fun fp16_default_policy_widens_to_fp32_floatArray() { + val values = floatArrayOf(0.0f, 1.0f, -1.0f, 0.5f, 3.0f, -2.5f) + val file = createF16SafeTensorsFile("weight", values) + + val weight = loadAll(file)["weight"] ?: error("missing 'weight'") + assertTrue( + weight.data is FloatArrayTensorData<*>, + "default policy must widen, got ${weight.data::class.simpleName}", + ) + // All of these are exactly representable in binary16. + assertContentEquals(values, weight.data.copyToFloatArray()) + } + + @Test + fun fp16_keep_native_emits_fp16DenseTensorData_with_on_disk_bytes() { + val values = floatArrayOf(0.0f, 1.0f, -1.0f, 0.5f, 3.0f, -2.5f, 100.0f, -64.0f) + val file = createF16SafeTensorsFile("weight", values) + + val weight = loadAll(file, fp16 = NarrowFloatLoadPolicy.KEEP_NATIVE)["weight"] + ?: error("missing 'weight'") + + assertTrue( + weight.data is Fp16DenseTensorData, + "KEEP_NATIVE must produce Fp16DenseTensorData, got ${weight.data::class.simpleName}", + ) + assertTrue(weight.data is NarrowFloatTensorData, "must be recognizable to narrow dispatch") + assertTrue( + weight.data !is Bf16TensorData, + "an F16 tensor must never be mistaken for BF16 — the bit layouts differ", + ) + + // Byte-for-byte identity proves no widening pass ran. + assertContentEquals( + fp32ToFp16Bytes(values), + (weight.data as Fp16DenseTensorData).packedData, + "KEEP_NATIVE must preserve on-disk F16 bytes verbatim", + ) + assertEquals( + values.size * 2, (weight.data as Fp16DenseTensorData).packedData.size, + "storage must stay 2 bytes per element", + ) + } + + @Test + fun fp16_keep_native_decodes_bit_identically_to_the_widening_path() { + // Both policies apply the same binary16 -> f32 decode; only the timing differs. + val values = FloatArray(64) { (it - 32) * 0.25f } + val file = createF16SafeTensorsFile("w", values) + + val widened = loadAll(file)["w"]!!.data.copyToFloatArray() + val native = loadAll(file, fp16 = NarrowFloatLoadPolicy.KEEP_NATIVE)["w"]!!.data.copyToFloatArray() + + assertEquals(widened.size, native.size) + for (i in widened.indices) { + assertEquals( + widened[i].toRawBits(), native[i].toRawBits(), + "bit-identity expected at $i: widened=${widened[i]} native=${native[i]}", + ) + } + } + + @Test + fun the_two_narrow_policies_are_independent() { + val f16Values = floatArrayOf(1.0f, 2.0f, 4.0f, 8.0f) + val bf16Values = floatArrayOf(0.5f, 0.25f, 16.0f, -3.0f) + val file = createF16AndBf16File(f16Values, bf16Values) + + // Keep F16 packed, widen BF16. + val a = loadAll(file, fp16 = NarrowFloatLoadPolicy.KEEP_NATIVE) + assertTrue(a["w_f16"]!!.data is Fp16DenseTensorData, "F16 should be packed") + assertTrue(a["w_bf16"]!!.data is FloatArrayTensorData<*>, "BF16 should be widened") + + // ...and the mirror image. + val b = loadAll(file, bf16 = NarrowFloatLoadPolicy.KEEP_NATIVE) + assertTrue(b["w_f16"]!!.data is FloatArrayTensorData<*>, "F16 should be widened") + assertTrue(b["w_bf16"]!!.data is Bf16TensorData, "BF16 should be packed") + } +} From df03f625d51c0b79643518876ac272f4a75d29e3 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 26 Jul 2026 19:02:52 +0200 Subject: [PATCH 3/4] feat(backend-cpu): FP16 matmul kernels and codec-driven narrow dispatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Third stage (B2) of #884. FP16 tensors now reach a real kernel instead of paying a per-element decode through `get`, which is what made the B3 KEEP_NATIVE path correct-but-slow for FP16. SPI: - `NarrowFloatMatmulKernel` carries a `codec` and the existing matmul signature; `Bf16MatmulKernel` and the new `Fp16MatmulKernel` extend it. The two formats differ by a decode step and nothing else. - `KernelProvider.matmulFp16()`, defaulting to null so providers without an FP16 kernel cascade as before. Kernels: - `ScalarFp16MatmulKernel` (priority 0, every KMP target) — the correctness reference. - `PanamaVectorFp16MatmulKernel` (priority 50) — vectorized FP32 FMA over n. The DECODE stays scalar into a lane-width scratch buffer: BF16 widens with an integer shift, but binary16 needs exponent rebiasing and mainstream JDKs expose no FP16 vector species. Expect this to trail the BF16 kernel; that is the format on this platform, not a defect here. Dispatch: - `DefaultCpuOpsJvm.chooseQuantizedMatmul` matches `NarrowFloatTensorData` rather than `Bf16TensorData` and picks the kernel by `codec`. Selecting by codec rather than width is load-bearing. Both formats are 2 bytes per element, so a mix-up cannot fail loudly — decoding F16 bytes through the BF16 shift yields plausible-looking wrong numbers with no exception. `fp16_weights_are_not_decoded_by_the_bf16_kernel` pins this from both sides: the dispatch must match the FP16 decode, AND the BF16 decode of the same bytes must differ by >10x, so the test cannot silently pass if the two converge. Accumulation is FP32 in every implementation. The narrow format is a storage width, never an accumulate width. Tests: 11 new (7 parity incl. lane-aligned/tail/k==0/empty shapes and an exact-FP32-reference check, 4 dispatch). Existing BF16 suites pass unmodified (10 parity + 3 dispatch). apiDump regenerated; full engine build clean. NOT included, deliberately: - Native FFM FP16 kernel (priority 100). BF16's lives in native/src/bf16_matmul.c; an FP16 peer needs new C plus build wiring. FP16 tops out at the Panama tier. - Vectorizing the BF16 dequant (IntVector.lanewise(LSHL, 16)) that PanamaVectorBf16MatmulKernel's own KDoc flags — a throughput change to an already-correct kernel, separable from making FP16 work. Refs #884 Co-Authored-By: Claude Opus 5 (1M context) --- .../backend/api/kernel/Bf16MatmulKernel.kt | 10 +- .../backend/api/kernel/KernelProvider.kt | 6 + .../api/kernel/NarrowFloatMatmulKernel.kt | 61 +++++++++ .../api/jvm/skainet-backend-cpu.api | 18 +++ .../exec/kernel/ScalarFp16MatmulKernel.kt | 49 +++++++ .../ainet/exec/kernel/ScalarKernelProvider.kt | 3 + .../kernel/PanamaVectorFp16MatmulKernel.kt | 87 +++++++++++++ .../exec/kernel/PanamaVectorKernelProvider.kt | 4 + .../ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt | 52 ++++++-- .../PanamaVectorFp16MatmulKernelParityTest.kt | 123 ++++++++++++++++++ .../exec/tensor/ops/Fp16MatmulDispatchTest.kt | 119 +++++++++++++++++ 11 files changed, 518 insertions(+), 14 deletions(-) create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/NarrowFloatMatmulKernel.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarFp16MatmulKernel.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernel.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernelParityTest.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/Fp16MatmulDispatchTest.kt diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Bf16MatmulKernel.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Bf16MatmulKernel.kt index b17b75a4..79e7f98a 100644 --- a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Bf16MatmulKernel.kt +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/Bf16MatmulKernel.kt @@ -1,5 +1,8 @@ package sk.ainet.backend.api.kernel +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.NarrowFloatCodec + /** * FP32 input × BF16-packed weight matrix multiplication: * @@ -43,7 +46,10 @@ package sk.ainet.backend.api.kernel * and byte-stride accessors on the same `ByteArray`. For a contiguous * `(k, n)` BF16 matrix `bByteStride == n * 2`. */ -public interface Bf16MatmulKernel { +public interface Bf16MatmulKernel : NarrowFloatMatmulKernel { + + override val codec: NarrowFloatCodec get() = Bf16Codec + /** * @param a left operand `(m, k)`, row-major FP32, stride [aStride] * floats per row. @@ -64,7 +70,7 @@ public interface Bf16MatmulKernel { * @param k contraction dimension (cols of A == rows of B). `k == 0` * zeros the `m × n` output block. */ - public fun matmul( + override fun matmul( a: FloatArray, aOffset: Int, aStride: Int, b: ByteArray, bByteOffset: Int, bByteStride: Int, out: FloatArray, outOffset: Int, outStride: Int, diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelProvider.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelProvider.kt index ad475190..b3ebf43a 100644 --- a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelProvider.kt +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelProvider.kt @@ -61,6 +61,12 @@ public interface KernelProvider { */ public fun matmulBf16(): Bf16MatmulKernel? = null + /** + * IEEE binary16 weight matmul, or `null` when this provider carries no FP16 kernel. + * Callers cascade to a lower-priority provider, ultimately the scalar reference. + */ + public fun matmulFp16(): Fp16MatmulKernel? = null + /** * F32 × Q8_0 matmul kernel exposed by this provider, or `null` if * this provider does not specialize Q8_0. Same fall-through pattern. diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/NarrowFloatMatmulKernel.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/NarrowFloatMatmulKernel.kt new file mode 100644 index 00000000..5c81483c --- /dev/null +++ b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/NarrowFloatMatmulKernel.kt @@ -0,0 +1,61 @@ +package sk.ainet.backend.api.kernel + +import sk.ainet.lang.types.Fp16Codec +import sk.ainet.lang.types.NarrowFloatCodec + +/** + * Matmul over a **16-bit float** right operand: `C(m,n) = A(m,k) · B(k,n)` where `B` is packed at + * 2 bytes per element and `A`/`C` are FP32. + * + * The shape, offset, aliasing and "fully overwrite the `m × n` block" obligations are identical to + * `Fp32MatmulKernel`. Only the inner-loop decode of `B` differs, which is what [codec] selects — so + * BF16 and FP16 differ by a decode step and nothing else. + * + * **Accumulation is always FP32.** The narrow format is a storage width, never an accumulate width; + * every implementation widens each `B` element before the multiply and sums in `Float`. This + * matches what PyTorch, JAX and tensor cores do, and is not a shortcut to be "optimized" away. + * + * ## Stride convention + * + * `A` and `out` use **float** strides (counting FP32 elements). `B` uses **byte** offsets and + * strides, matching the byte-packed convention of the block-quantized kernels and avoiding the + * foot-gun of mixing element- and byte-strides on the same `ByteArray`. For a contiguous `(k, n)` + * matrix, `bByteStride == n * 2`. + */ +public interface NarrowFloatMatmulKernel { + + /** The 16-bit format this kernel decodes `B` from. */ + public val codec: NarrowFloatCodec + + /** + * @param a left operand `(m, k)`, row-major FP32, [aStride] floats per row. + * @param aOffset element offset into [a] of the (0, 0) entry. + * @param aStride distance in **floats** between consecutive rows of [a]. + * @param b right operand `(k, n)`, row-major, packed little-endian 16-bit floats. + * @param bByteOffset **byte** offset into [b] of the (0, 0) entry. + * @param bByteStride **byte** distance between consecutive rows of [b]. + * @param out output `(m, n)`, row-major FP32, [outStride] floats per row. + * @param outOffset element offset into [out]. + * @param outStride distance in **floats** between consecutive rows of [out]. + * @param m rows of A and C. + * @param n columns of B and C. + * @param k contraction dimension. `k == 0` zeros the `m × n` output block. + */ + public fun matmul( + a: FloatArray, aOffset: Int, aStride: Int, + b: ByteArray, bByteOffset: Int, bByteStride: Int, + out: FloatArray, outOffset: Int, outStride: Int, + m: Int, n: Int, k: Int, + ) +} + +/** + * Matmul over an **IEEE binary16** right operand. + * + * The FP16 counterpart to [Bf16MatmulKernel]. Note that FP16 cannot use the bit-shift decode that + * makes BF16 nearly free — binary16 needs exponent rebiasing — so vectorized FP16 implementations + * lag their BF16 equivalents. Correctness first; throughput is a separate concern. + */ +public interface Fp16MatmulKernel : NarrowFloatMatmulKernel { + override val codec: NarrowFloatCodec get() = Fp16Codec +} diff --git a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api index 0ed4e5fb..2ce08b89 100644 --- a/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api +++ b/skainet-backends/skainet-backend-cpu/api/jvm/skainet-backend-cpu.api @@ -43,6 +43,13 @@ public final class sk/ainet/context/DirectCpuExecutionContext$Companion { public final class sk/ainet/exec/kernel/PanamaVectorBf16MatmulKernel : sk/ainet/backend/api/kernel/Bf16MatmulKernel { public static final field INSTANCE Lsk/ainet/exec/kernel/PanamaVectorBf16MatmulKernel; + public fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; + public fun matmul ([FII[BII[FIIIII)V +} + +public final class sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernel : sk/ainet/backend/api/kernel/Fp16MatmulKernel { + public static final field INSTANCE Lsk/ainet/exec/kernel/PanamaVectorFp16MatmulKernel; + public fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; public fun matmul ([FII[BII[FIIIII)V } @@ -52,6 +59,7 @@ public final class sk/ainet/exec/kernel/PanamaVectorKernelProvider : sk/ainet/ba public fun getPriority ()I public fun isAvailable ()Z public fun matmulBf16 ()Lsk/ainet/backend/api/kernel/Bf16MatmulKernel; + public fun matmulFp16 ()Lsk/ainet/backend/api/kernel/Fp16MatmulKernel; public fun matmulFp32 ()Lsk/ainet/backend/api/kernel/Fp32MatmulKernel; public fun matmulQ4K ()Lsk/ainet/backend/api/kernel/Q4KMatmulKernel; public fun matmulQ4_0 ()Lsk/ainet/backend/api/kernel/Q4_0MatmulKernel; @@ -69,6 +77,7 @@ public final class sk/ainet/exec/kernel/PanamaVectorKernelProviderFactory : sk/a public fun getPriority ()I public fun isAvailable ()Z public fun matmulBf16 ()Lsk/ainet/backend/api/kernel/Bf16MatmulKernel; + public fun matmulFp16 ()Lsk/ainet/backend/api/kernel/Fp16MatmulKernel; public fun matmulFp32 ()Lsk/ainet/backend/api/kernel/Fp32MatmulKernel; public fun matmulQ4K ()Lsk/ainet/backend/api/kernel/Q4KMatmulKernel; public fun matmulQ4_0 ()Lsk/ainet/backend/api/kernel/Q4_0MatmulKernel; @@ -122,6 +131,13 @@ public final class sk/ainet/exec/kernel/PanamaVectorQ8_0MatmulKernel : sk/ainet/ public final class sk/ainet/exec/kernel/ScalarBf16MatmulKernel : sk/ainet/backend/api/kernel/Bf16MatmulKernel { public static final field INSTANCE Lsk/ainet/exec/kernel/ScalarBf16MatmulKernel; + public fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; + public fun matmul ([FII[BII[FIIIII)V +} + +public final class sk/ainet/exec/kernel/ScalarFp16MatmulKernel : sk/ainet/backend/api/kernel/Fp16MatmulKernel { + public static final field INSTANCE Lsk/ainet/exec/kernel/ScalarFp16MatmulKernel; + public fun getCodec ()Lsk/ainet/lang/types/NarrowFloatCodec; public fun matmul ([FII[BII[FIIIII)V } @@ -131,6 +147,7 @@ public final class sk/ainet/exec/kernel/ScalarKernelProvider : sk/ainet/backend/ public fun getPriority ()I public fun isAvailable ()Z public fun matmulBf16 ()Lsk/ainet/backend/api/kernel/Bf16MatmulKernel; + public fun matmulFp16 ()Lsk/ainet/backend/api/kernel/Fp16MatmulKernel; public fun matmulFp32 ()Lsk/ainet/backend/api/kernel/Fp32MatmulKernel; public fun matmulQ4K ()Lsk/ainet/backend/api/kernel/Q4KMatmulKernel; public fun matmulQ4_0 ()Lsk/ainet/backend/api/kernel/Q4_0MatmulKernel; @@ -148,6 +165,7 @@ public final class sk/ainet/exec/kernel/ScalarKernelProviderFactory : sk/ainet/b public fun getPriority ()I public fun isAvailable ()Z public fun matmulBf16 ()Lsk/ainet/backend/api/kernel/Bf16MatmulKernel; + public fun matmulFp16 ()Lsk/ainet/backend/api/kernel/Fp16MatmulKernel; public fun matmulFp32 ()Lsk/ainet/backend/api/kernel/Fp32MatmulKernel; public fun matmulQ4K ()Lsk/ainet/backend/api/kernel/Q4KMatmulKernel; public fun matmulQ4_0 ()Lsk/ainet/backend/api/kernel/Q4_0MatmulKernel; diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarFp16MatmulKernel.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarFp16MatmulKernel.kt new file mode 100644 index 00000000..266da1c7 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarFp16MatmulKernel.kt @@ -0,0 +1,49 @@ +package sk.ainet.exec.kernel + +import sk.ainet.backend.api.kernel.Fp16MatmulKernel +import sk.ainet.lang.types.Fp16Codec + +/** + * Scalar reference implementation of [Fp16MatmulKernel] — three nested loops, decode inline, no + * SIMD. Always available on every KMP target, and the correctness reference an accelerated FP16 + * kernel must match within FP-ordering tolerance. + * + * Mirrors [ScalarBf16MatmulKernel]; the only difference is the decode, which for binary16 needs + * exponent rebiasing rather than BF16's bit shift. + */ +public object ScalarFp16MatmulKernel : Fp16MatmulKernel { + override fun matmul( + a: FloatArray, aOffset: Int, aStride: Int, + b: ByteArray, bByteOffset: Int, bByteStride: Int, + out: FloatArray, outOffset: Int, outStride: Int, + m: Int, n: Int, k: Int, + ) { + require(m >= 0 && n >= 0 && k >= 0) { + "ScalarFp16MatmulKernel: m, n, k must be non-negative; got m=$m n=$n k=$k" + } + if (m == 0 || n == 0) return + // k == 0 is an explicit "zero the output block" per the SPI. + if (k == 0) { + for (i in 0 until m) { + val rowOff = outOffset + i * outStride + for (j in 0 until n) out[rowOff + j] = 0f + } + return + } + for (i in 0 until m) { + val aRowOff = aOffset + i * aStride + val outRowOff = outOffset + i * outStride + for (j in 0 until n) { + // Accumulate in FP32 — the narrow format is storage only. + var sum = 0f + for (p in 0 until k) { + val bByteIdx = bByteOffset + p * bByteStride + j * 2 + val lo = b[bByteIdx].toInt() and 0xFF + val hi = b[bByteIdx + 1].toInt() and 0xFF + sum += a[aRowOff + p] * Fp16Codec.decode((hi shl 8) or lo) + } + out[outRowOff + j] = sum + } + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarKernelProvider.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarKernelProvider.kt index 31a8386b..c17bbd57 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarKernelProvider.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/kernel/ScalarKernelProvider.kt @@ -1,6 +1,7 @@ package sk.ainet.exec.kernel import sk.ainet.backend.api.kernel.Bf16MatmulKernel +import sk.ainet.backend.api.kernel.Fp16MatmulKernel import sk.ainet.backend.api.kernel.Fp32MatmulKernel import sk.ainet.backend.api.kernel.KernelProvider import sk.ainet.backend.api.kernel.Q4KMatmulKernel @@ -30,6 +31,8 @@ public object ScalarKernelProvider : KernelProvider { override fun isAvailable(): Boolean = true override fun matmulFp32(): Fp32MatmulKernel = ScalarMatmulKernel override fun matmulBf16(): Bf16MatmulKernel = ScalarBf16MatmulKernel + + override fun matmulFp16(): Fp16MatmulKernel = ScalarFp16MatmulKernel override fun matmulQ8_0(): Q8_0MatmulKernel = ScalarQ8_0MatmulKernel override fun matmulQ4_0(): Q4_0MatmulKernel = ScalarQ4_0MatmulKernel override fun matmulQ4K(): Q4KMatmulKernel = ScalarQ4_KMatmulKernel diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernel.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernel.kt new file mode 100644 index 00000000..bd68dd57 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernel.kt @@ -0,0 +1,87 @@ +package sk.ainet.exec.kernel + +import jdk.incubator.vector.FloatVector +import jdk.incubator.vector.VectorSpecies +import sk.ainet.backend.api.kernel.Fp16MatmulKernel +import sk.ainet.lang.types.Fp16Codec + +/** + * Panama Vector API implementation of [Fp16MatmulKernel] — the FP16 counterpart to + * [PanamaVectorBf16MatmulKernel], using the same i-p-j outer-product shape and FP32 FMA + * accumulation. + * + * **Why the decode stays scalar here.** BF16 can be widened lane-wise with an integer shift + * (`IntVector.lanewise(LSHL, 16).reinterpretAsFloats()`), which is why the BF16 kernel vectorizes + * its dequant. Binary16 needs exponent rebiasing and gradual-underflow handling, and mainstream + * JDKs expose no FP16 vector species, so each element is decoded scalar into a lane-width scratch + * buffer before the vectorized multiply-accumulate. The FMA over `n` is still fully vectorized — + * only the widening is not. Expect this kernel to trail the BF16 one; that is inherent to the + * format on this platform, not a defect in this implementation. + * + * Numerical parity vs [ScalarFp16MatmulKernel] is asserted by + * `PanamaVectorFp16MatmulKernelParityTest`. + */ +public object PanamaVectorFp16MatmulKernel : Fp16MatmulKernel { + + private val floatSpecies: VectorSpecies = FloatVector.SPECIES_PREFERRED + + override fun matmul( + a: FloatArray, aOffset: Int, aStride: Int, + b: ByteArray, bByteOffset: Int, bByteStride: Int, + out: FloatArray, outOffset: Int, outStride: Int, + m: Int, n: Int, k: Int, + ) { + require(m >= 0 && n >= 0 && k >= 0) { + "PanamaVectorFp16MatmulKernel: m, n, k must be non-negative; got m=$m n=$n k=$k" + } + if (m == 0 || n == 0) return + if (k == 0) { + for (i in 0 until m) { + val rowOff = outOffset + i * outStride + for (j in 0 until n) out[rowOff + j] = 0f + } + return + } + + val laneCount = floatSpecies.length() + val bound = floatSpecies.loopBound(n) + val scratch = FloatArray(laneCount) + + // Zero the output block first — the i-p-j outer product accumulates into it. + for (i in 0 until m) { + val rowOff = outOffset + i * outStride + for (j in 0 until n) out[rowOff + j] = 0f + } + + for (i in 0 until m) { + val aRowOff = aOffset + i * aStride + val outRowOff = outOffset + i * outStride + for (p in 0 until k) { + val aIp = a[aRowOff + p] + val aBcast = FloatVector.broadcast(floatSpecies, aIp) + val bRowByteOff = bByteOffset + p * bByteStride + var j = 0 + while (j < bound) { + val byteBase = bRowByteOff + j * 2 + for (lane in 0 until laneCount) { + val lo = b[byteBase + lane * 2].toInt() and 0xFF + val hi = b[byteBase + lane * 2 + 1].toInt() and 0xFF + scratch[lane] = Fp16Codec.decode((hi shl 8) or lo) + } + val bVec = FloatVector.fromArray(floatSpecies, scratch, 0) + val outVec = FloatVector.fromArray(floatSpecies, out, outRowOff + j) + aBcast.fma(bVec, outVec).intoArray(out, outRowOff + j) + j += laneCount + } + // Tail (scalar): up to laneCount - 1 trailing columns. + while (j < n) { + val bByteIdx = bRowByteOff + j * 2 + val lo = b[bByteIdx].toInt() and 0xFF + val hi = b[bByteIdx + 1].toInt() and 0xFF + out[outRowOff + j] += aIp * Fp16Codec.decode((hi shl 8) or lo) + j++ + } + } + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProvider.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProvider.kt index e9bb5b59..87218dbd 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProvider.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/PanamaVectorKernelProvider.kt @@ -1,6 +1,7 @@ package sk.ainet.exec.kernel import sk.ainet.backend.api.kernel.Bf16MatmulKernel +import sk.ainet.backend.api.kernel.Fp16MatmulKernel import sk.ainet.backend.api.kernel.Fp32MatmulKernel import sk.ainet.backend.api.kernel.KernelProvider import sk.ainet.backend.api.kernel.Q4KMatmulKernel @@ -51,6 +52,9 @@ public object PanamaVectorKernelProvider : KernelProvider { override fun matmulBf16(): Bf16MatmulKernel? = if (isAvailable()) PanamaVectorBf16MatmulKernel else null + override fun matmulFp16(): Fp16MatmulKernel? = + if (isAvailable()) PanamaVectorFp16MatmulKernel else null + override fun matmulQ8_0(): Q8_0MatmulKernel? = if (isAvailable()) PanamaVectorQ8_0MatmulKernel else null diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt index 0aa070cc..4b76fbf9 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt @@ -4,6 +4,7 @@ import jdk.incubator.vector.FloatVector import jdk.incubator.vector.VectorSpecies import jdk.incubator.vector.VectorOperators import sk.ainet.backend.api.kernel.Bf16MatmulKernel +import sk.ainet.backend.api.kernel.Fp16MatmulKernel import sk.ainet.backend.api.kernel.Fp32MatmulKernel import sk.ainet.backend.api.kernel.KernelRegistry import sk.ainet.backend.api.kernel.KernelServiceLoader @@ -12,6 +13,7 @@ import sk.ainet.backend.api.kernel.Q4KMatmulKernel import sk.ainet.backend.api.kernel.Q4_0MatmulKernel import sk.ainet.backend.api.kernel.Q8_0MatmulKernel import sk.ainet.exec.kernel.ScalarBf16MatmulKernel +import sk.ainet.exec.kernel.ScalarFp16MatmulKernel import sk.ainet.exec.kernel.ScalarMatmulKernel import sk.ainet.exec.kernel.ScalarQ4_0MatmulKernel import sk.ainet.lang.tensor.Shape @@ -23,6 +25,9 @@ import sk.ainet.lang.tensor.data.MemorySegmentTensorData import sk.ainet.lang.tensor.data.Q4MemorySegmentMarker import sk.ainet.lang.tensor.data.Q4MemorySegmentTensorData import sk.ainet.lang.tensor.data.Bf16TensorData +import sk.ainet.lang.tensor.data.NarrowFloatTensorData +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.Fp16Codec import sk.ainet.lang.tensor.data.Q4_0TensorData import sk.ainet.lang.tensor.data.Q8_0TensorData import sk.ainet.lang.tensor.data.Q8MemorySegmentMarker @@ -126,6 +131,17 @@ internal class DefaultCpuOpsJvm( ?: ScalarBf16MatmulKernel } + /** FP16 matmul kernel resolved via [KernelRegistry]; mirrors [bf16MatmulKernel]. */ + private val fp16MatmulKernel: Fp16MatmulKernel by lazy { + if (KernelRegistry.providers().isEmpty()) { + KernelServiceLoader.installAll() + } + KernelRegistry.providers() + .firstOrNull { it.isAvailable() && it.matmulFp16() != null } + ?.matmulFp16() + ?: ScalarFp16MatmulKernel + } + /** * Q4_0 matmul kernel resolved via [KernelRegistry]. Mirrors * [bf16MatmulKernel]: non-null, picks the highest-priority provider @@ -588,20 +604,32 @@ internal class DefaultCpuOpsJvm( @Suppress("UNCHECKED_CAST") CpuTensor(outData as TensorData, this, a.dtype) } - is Bf16TensorData -> { - // BF16 is dense (not block-quantized) and the kernel SPI is a + is NarrowFloatTensorData -> { + // Narrow floats are dense (not block-quantized) and the kernel SPI is a // full SGEMM with `(m, n, k)` strides — no per-batch loop needed, // unlike the matvec-shaped Q4_K / Q8_0 / Q6_K branches. - val outBuffer = FloatArray(batchSize * outputDim) - bf16MatmulKernel.matmul( - inputBuffer, 0, inputDim, - bData.packedData, 0, outputDim * Bf16TensorData.BYTES_PER_ELEMENT, - outBuffer, 0, outputDim, - batchSize, outputDim, inputDim, - ) - val outData = DenseFloatArrayTensorData(Shape(batchSize, outputDim), outBuffer) - @Suppress("UNCHECKED_CAST") - CpuTensor(outData as TensorData, this, a.dtype) + // + // The codec selects the kernel: the two formats have different bit layouts, + // so decoding F16 bytes with the BF16 kernel would yield silent garbage. + val kernel = when (bData.codec) { + Bf16Codec -> bf16MatmulKernel + Fp16Codec -> fp16MatmulKernel + else -> null + } + if (kernel == null) { + null + } else { + val outBuffer = FloatArray(batchSize * outputDim) + kernel.matmul( + inputBuffer, 0, inputDim, + bData.packedData, 0, outputDim * NarrowFloatTensorData.BYTES_PER_ELEMENT, + outBuffer, 0, outputDim, + batchSize, outputDim, inputDim, + ) + val outData = DenseFloatArrayTensorData(Shape(batchSize, outputDim), outBuffer) + @Suppress("UNCHECKED_CAST") + CpuTensor(outData as TensorData, this, a.dtype) + } } // Q6_K / Q5_1 / Q5_0 dispatch is handled in DefaultCpuOpsBase via the kernel // registry (block-major, shared with Native); not intercepted here. diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernelParityTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernelParityTest.kt new file mode 100644 index 00000000..6a8630cd --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/PanamaVectorFp16MatmulKernelParityTest.kt @@ -0,0 +1,123 @@ +package sk.ainet.exec.kernel + +import sk.ainet.lang.types.Fp16Codec +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertTrue + +/** + * [PanamaVectorFp16MatmulKernel] must agree with the scalar reference + * [ScalarFp16MatmulKernel] within FMA + reordered-reduction tolerance. + * + * Mirrors `PanamaVectorBf16MatmulKernelParityTest`. Tolerance is tighter here than the BF16 test's + * because binary16 carries three more mantissa bits, so the decode error it starts from is ~8× + * smaller. + */ +class PanamaVectorFp16MatmulKernelParityTest { + + /** Pack FP32 values as little-endian binary16 bytes. */ + private fun fp16Bytes(values: FloatArray): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = Fp16Codec.encode(values[i]) + out[i * 2] = (bits and 0xFF).toByte() + out[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return out + } + + private fun assertParity(m: Int, n: Int, k: Int, seed: Int) { + val rnd = Random(seed) + val a = FloatArray(m * k) { rnd.nextFloat() * 2f - 1f } + val bFloats = FloatArray(k * n) { rnd.nextFloat() * 2f - 1f } + val b = fp16Bytes(bFloats) + + val outScalar = FloatArray(m * n) + val outPanama = FloatArray(m * n) + + ScalarFp16MatmulKernel.matmul(a, 0, k, b, 0, n * 2, outScalar, 0, n, m, n, k) + PanamaVectorFp16MatmulKernel.matmul(a, 0, k, b, 0, n * 2, outPanama, 0, n, m, n, k) + + val tol = 1e-5f * k + for (i in outScalar.indices) { + assertTrue( + abs(outScalar[i] - outPanama[i]) <= tol, + "mismatch at $i (m=$m n=$n k=$k): scalar=${outScalar[i]} panama=${outPanama[i]}", + ) + } + } + + @Test + fun parity_on_lane_aligned_shapes() { + assertParity(m = 4, n = 32, k = 16, seed = 1) + assertParity(m = 1, n = 64, k = 64, seed = 2) + } + + @Test + fun parity_on_shapes_with_a_scalar_tail() { + // n = 7 is unlikely to be a multiple of any FloatVector lane count, exercising the tail. + assertParity(m = 3, n = 7, k = 5, seed = 3) + assertParity(m = 2, n = 13, k = 9, seed = 4) + assertParity(m = 5, n = 1, k = 3, seed = 5) + } + + @Test + fun parity_on_single_element() { + assertParity(m = 1, n = 1, k = 1, seed = 6) + } + + @Test + fun k_zero_zeroes_the_output_block_in_both() { + val m = 3 + val n = 4 + val outScalar = FloatArray(m * n) { 7f } + val outPanama = FloatArray(m * n) { 7f } + + ScalarFp16MatmulKernel.matmul(FloatArray(0), 0, 0, ByteArray(0), 0, 0, outScalar, 0, n, m, n, 0) + PanamaVectorFp16MatmulKernel.matmul(FloatArray(0), 0, 0, ByteArray(0), 0, 0, outPanama, 0, n, m, n, 0) + + assertTrue(outScalar.all { it == 0f }, "scalar must zero the block") + assertTrue(outPanama.all { it == 0f }, "panama must zero the block") + } + + @Test + fun empty_m_or_n_is_a_no_op_in_both() { + val out = FloatArray(4) { 3f } + ScalarFp16MatmulKernel.matmul(FloatArray(0), 0, 0, ByteArray(0), 0, 0, out, 0, 0, 0, 0, 0) + PanamaVectorFp16MatmulKernel.matmul(FloatArray(0), 0, 0, ByteArray(0), 0, 0, out, 0, 0, 0, 0, 0) + assertTrue(out.all { it == 3f }, "m == 0 / n == 0 must not touch out") + } + + @Test + fun both_kernels_report_the_fp16_codec() { + assertTrue(ScalarFp16MatmulKernel.codec === Fp16Codec) + assertTrue(PanamaVectorFp16MatmulKernel.codec === Fp16Codec) + } + + @Test + fun fp16_kernel_result_tracks_an_exact_fp32_reference() { + // With operands exactly representable in binary16, the kernel must reproduce a plain + // FP32 matmul exactly — proving the decode is right, not merely self-consistent. + val m = 2 + val n = 3 + val k = 4 + val a = floatArrayOf(1f, 2f, -1f, 0.5f, 0.25f, -2f, 4f, 1f) + val bFloats = FloatArray(k * n) { (it % 5) * 0.5f - 1f } + val b = fp16Bytes(bFloats) + + val out = FloatArray(m * n) + ScalarFp16MatmulKernel.matmul(a, 0, k, b, 0, n * 2, out, 0, n, m, n, k) + + for (i in 0 until m) { + for (j in 0 until n) { + var expected = 0f + for (p in 0 until k) expected += a[i * k + p] * bFloats[p * n + j] + assertTrue( + abs(expected - out[i * n + j]) <= 1e-6f, + "exact-value mismatch at ($i,$j): expected=$expected got=${out[i * n + j]}", + ) + } + } + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/Fp16MatmulDispatchTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/Fp16MatmulDispatchTest.kt new file mode 100644 index 00000000..c1353465 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/Fp16MatmulDispatchTest.kt @@ -0,0 +1,119 @@ +package sk.ainet.exec.tensor.ops + +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertTrue +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.exec.kernel.ScalarBf16MatmulKernel +import sk.ainet.exec.kernel.ScalarFp16MatmulKernel +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Fp16DenseTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec + +/** + * Integration tests for the FP32 × FP16 dispatch path in `DefaultCpuOpsJvm.matmul`. + * + * The dispatch branch matches on `NarrowFloatTensorData` and then selects the kernel by the data's + * `codec`. These tests prove an `Fp16DenseTensorData` weight reaches the FP16 kernel — and, more + * importantly, that it is NOT decoded by the BF16 kernel, which would be silent garbage rather + * than an error since both formats are 2 bytes wide. + * + * Mirrors [Bf16MatmulDispatchTest]. + */ +class Fp16MatmulDispatchTest { + + private val ctx = DirectCpuExecutionContext() + + /** binary16 carries 10 mantissa bits; accumulated error scales with `k`. */ + private val fp16TolPerK = 1.5e-3f + + private fun fp32ToFp16Bytes(values: FloatArray): ByteArray { + val out = ByteArray(values.size * 2) + for (i in values.indices) { + val bits = Fp16Codec.encode(values[i]) + out[i * 2] = (bits and 0xFF).toByte() + out[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + return out + } + + @Suppress("UNCHECKED_CAST") + private fun fp16Weight(inputDim: Int, outputDim: Int, seed: Int): Pair, ByteArray> { + val rng = Random(seed) + val values = FloatArray(inputDim * outputDim) { rng.nextFloat() - 0.5f } + val bytes = fp32ToFp16Bytes(values) + val data = Fp16DenseTensorData(Shape(inputDim, outputDim), bytes) as TensorData + return ctx.fromData(data, FP32::class) to bytes + } + + private fun assertDispatchMatchesScalar(m: Int, k: Int, n: Int, seed: Int) { + val rng = Random(seed) + val inputFloats = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val (weight, weightBytes) = fp16Weight(k, n, seed) + val input = ctx.fromFloatArray(Shape(m, k), FP32::class, inputFloats) + + val outArr = ctx.ops.matmul(input, weight).data.copyToFloatArray() + + val expected = FloatArray(m * n) + ScalarFp16MatmulKernel.matmul(inputFloats, 0, k, weightBytes, 0, n * 2, expected, 0, n, m, n, k) + + val tol = (fp16TolPerK * k.coerceAtLeast(1)).coerceAtLeast(fp16TolPerK) + for (i in expected.indices) { + assertTrue( + abs(expected[i] - outArr[i]) <= tol, + "FP16 dispatch mismatch at $i: expected=${expected[i]} got=${outArr[i]} tol=$tol", + ) + } + } + + @Test + fun single_batch_matmul_against_fp16_weight_routes_correctly() { + assertDispatchMatchesScalar(m = 1, k = 128, n = 64, seed = 1) + } + + @Test + fun multi_batch_matmul_against_fp16_weight_routes_correctly() { + assertDispatchMatchesScalar(m = 3, k = 256, n = 32, seed = 2) + } + + @Test + fun llm_typical_attention_proj_matmul_routes_correctly() { + assertDispatchMatchesScalar(m = 1, k = 512, n = 512, seed = 3) + } + + @Test + fun fp16_weights_are_not_decoded_by_the_bf16_kernel() { + // Both formats are 2 bytes wide, so a codec mix-up cannot fail loudly — it silently + // produces wrong numbers. This pins that the dispatch picks by codec, not by width: + // decoding the same bytes as BF16 must give a materially different answer, and the + // dispatch must match the FP16 reading. + val m = 1 + val k = 64 + val n = 8 + val rng = Random(42) + val inputFloats = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val (weight, bytes) = fp16Weight(k, n, seed = 42) + val input = ctx.fromFloatArray(Shape(m, k), FP32::class, inputFloats) + + val dispatched = ctx.ops.matmul(input, weight).data.copyToFloatArray() + + val asFp16 = FloatArray(m * n) + ScalarFp16MatmulKernel.matmul(inputFloats, 0, k, bytes, 0, n * 2, asFp16, 0, n, m, n, k) + val asBf16 = FloatArray(m * n) + ScalarBf16MatmulKernel.matmul(inputFloats, 0, k, bytes, 0, n * 2, asBf16, 0, n, m, n, k) + + val fp16Err = dispatched.indices.sumOf { abs(dispatched[it] - asFp16[it]).toDouble() } + val bf16Err = dispatched.indices.sumOf { abs(dispatched[it] - asBf16[it]).toDouble() } + + assertTrue(fp16Err < 1e-3, "dispatch must match the FP16 decode, err=$fp16Err") + assertTrue( + bf16Err > fp16Err * 10, + "the BF16 decode of the same bytes must differ materially " + + "(fp16Err=$fp16Err bf16Err=$bf16Err); if these are close the test proves nothing", + ) + } +} From cb1af5e7eb1a28cad245113c3733d2c13be450cf Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 26 Jul 2026 19:21:29 +0200 Subject: [PATCH 4/4] fix(lang): one dtype lattice, narrow-float completeness in ops and factory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Final stage (B4) of #884. Closes the gaps that left BF16 half-supported and repairs a type lattice that disagreed with itself. Lattice — five bugs, four of them code contradicting its own comment: - `FP32.promoteTo(Ternary)` returned Ternary: a promotion to LOWER precision, which no promotion lattice should ever produce. Comment said FP32. - `FP32.isCompatible(Ternary)` returned false; comment said Ternary promotes to FP32. - `FP16.promoteTo(Int32)` returned FP16, silently dropping ~21 bits of integer precision, while `Int32.promoteTo(FP16)` correctly returned FP32. - `Int8.promoteTo(BF16)` returned FP32 even though bf16's 8 significand bits represent every Int8 value exactly — and `BF16.promoteTo(Int8)` said BF16. Now consistent with how Int8 already treats FP16. - `commonPrecisionWith` was a SECOND hand-written lattice with no BF16 arm at all, so `BF16.commonPrecisionWith(Int8)` said FP32 while `BF16.promoteTo(Int8)` said BF16. It now delegates to `promoteTo`, so the two cannot drift apart again. `isConvertibleTo` had no BF16 arm either — every BF16 conversion reported false. Now expressed via a new `DType.isFloatingPoint()` so any float width converts to any other. Completeness: - `supportsFloatOps` admits BF16, which was excluded and therefore left BF16-tagged tensors matmul-only: every elementwise and unary op fell through to the generic scalar path. Callers still re-check for FloatArrayTensorData, so genuinely narrow-storage tensors keep falling through safely. - `convert` to FP16/BF16 now ROUNDS. It previously shared the FP32 branch, so `convert(x, FP16)` only re-tagged the dtype while keeping full FP32 precision — the tensor claimed to be FP16 while holding values no FP16 can represent, with no rounding, no overflow to Inf past the format range, and no NaN handling. BF16 was rejected as a target outright. - `DenseTensorDataFactory`: BF16 arms for `ones`/`full`/`randn`, which threw, AND for both `fromFloatArray` overloads — the latter was not in the original survey and is why the first convert-to-BF16 run still failed. Storage for a narrow `convert` target stays float-backed (values rounded, tag correct), matching how DenseTensorDataFactory already handles narrow tags. Emitting genuinely 2-byte storage from `convert` would change what every existing caller receives and is a separate decision. Tests: 9 lattice-consistency (incl. agreement across all 196 dtype pairs, float/int symmetry, and a "promotion never narrows" invariant) and 8 convert tests (rounding, fp16 overflow to Inf where bf16 stays finite, subnormal underflow, idempotence). apiDump regenerated; full engine build clean. Refs #884 Co-Authored-By: Claude Opus 5 (1M context) --- .../sk/ainet/exec/tensor/ops/DefaultCpuOps.kt | 23 +++- .../ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt | 8 +- .../exec/tensor/ops/NarrowFloatConvertTest.kt | 102 ++++++++++++++ .../api/jvm/skainet-lang-core.api | 1 + .../tensor/data/DenseTensorDataFactory.kt | 28 ++-- .../sk/ainet/lang/types/DTypeExtensions.kt | 34 +++-- .../kotlin/sk/ainet/lang/types/FP16.kt | 4 +- .../kotlin/sk/ainet/lang/types/FP32.kt | 4 +- .../kotlin/sk/ainet/lang/types/Int8.kt | 4 +- .../lang/types/DTypeLatticeConsistencyTest.kt | 127 ++++++++++++++++++ 10 files changed, 300 insertions(+), 35 deletions(-) create mode 100644 skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/NarrowFloatConvertTest.kt create mode 100644 skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/DTypeLatticeConsistencyTest.kt diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index 0c3b7c0b..f0d3aec8 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt @@ -30,8 +30,11 @@ import sk.ainet.lang.tensor.data.Q8_0BlockTensorData import sk.ainet.lang.tensor.data.TensorData import sk.ainet.lang.tensor.data.TensorDataFactory import sk.ainet.lang.tensor.ops.UpsampleMode +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.Bf16Codec import sk.ainet.lang.types.FP16 import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec import sk.ainet.lang.types.Int32 import sk.ainet.lang.types.Int8 import kotlin.math.floor @@ -2685,18 +2688,34 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory @Suppress("UNCHECKED_CAST") val outData = when (targetClass) { - FP32::class, FP16::class -> dataFactory.fromFloatArray( + FP32::class -> dataFactory.fromFloatArray( tensor.shape, targetClass, copyTensorValuesAsFloatArray(tensor) ) as TensorData + // Narrowing to 16 bits must actually ROUND. This previously shared the FP32 branch, + // so `convert(x, FP16)` only re-tagged the dtype and kept full FP32 precision — the + // resulting tensor claimed to be FP16 while holding values no FP16 can represent, with + // no rounding, no overflow-to-Inf past the format's range, and no NaN handling. + // Storage stays float-backed (matching DenseTensorDataFactory's narrow tags); it is + // the VALUES that now reflect the target format. + FP16::class, BF16::class -> { + val codec = if (targetClass == BF16::class) Bf16Codec else Fp16Codec + val src = copyTensorValuesAsFloatArray(tensor) + val rounded = FloatArray(src.size) { codec.decode(codec.encode(src[it])) } + dataFactory.fromFloatArray( + tensor.shape, + targetClass, + rounded + ) as TensorData + } Int32::class, Int8::class -> dataFactory.fromIntArray( tensor.shape, targetClass, copyTensorValuesAsIntArray(tensor) ) as TensorData else -> throw IllegalArgumentException( - "convert supports FP32, FP16, Int32, and Int8 targets, got ${targetType.name}" + "convert supports FP32, FP16, BF16, Int32, and Int8 targets, got ${targetType.name}" ) } return CpuTensor(outData, this, targetClass, GradState(requiresGrad = tensor.requiresGrad)) diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt index 4b76fbf9..587ccc5d 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt @@ -35,6 +35,7 @@ import sk.ainet.lang.tensor.data.Q8MemorySegmentTensorData import sk.ainet.lang.tensor.data.Q4_KBlockTensorData import sk.ainet.lang.tensor.data.Q4_KTensorData import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.types.BF16 import sk.ainet.lang.types.DType import sk.ainet.lang.types.FP16 import sk.ainet.lang.types.FP32 @@ -921,7 +922,12 @@ internal class DefaultCpuOpsJvm( private fun supportsFloatOps(tensor: Tensor): Boolean { val dtype = tensor.dtype - return (dtype == FP32::class || dtype == FP16::class) + // BF16 was excluded here, which left BF16-tagged tensors matmul-only — every elementwise + // and unary op fell through to the generic scalar path. All three float tags are backed by + // FloatArray when produced by DenseTensorDataFactory, so they can use the vector kernels. + // Callers re-check with `data as? FloatArrayTensorData`, so genuinely narrow-storage + // tensors (Bf16/Fp16DenseTensorData) still fall through safely rather than being misread. + return dtype == FP32::class || dtype == FP16::class || dtype == BF16::class } private fun chooseMatmul(a: Tensor, b: Tensor): Tensor? { diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/NarrowFloatConvertTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/NarrowFloatConvertTest.kt new file mode 100644 index 00000000..7d7811b7 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/NarrowFloatConvertTest.kt @@ -0,0 +1,102 @@ +package sk.ainet.exec.tensor.ops + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * `convert` to a 16-bit float target must actually round. + * + * It previously shared the FP32 branch, so `convert(x, FP16)` only re-tagged the dtype while + * keeping full FP32 precision — the tensor claimed to be FP16 while holding values no FP16 can + * represent, with no rounding, no overflow to Inf past the format's range, and no NaN handling. + * BF16 was rejected as a target outright. + */ +class NarrowFloatConvertTest { + + private val ctx = DirectCpuExecutionContext() + + private fun fp32(vararg values: Float) = + ctx.fromFloatArray(Shape(values.size), FP32::class, values) + + @Test + fun convert_to_fp16_rounds_values_to_binary16() { + // 1 + 2^-20 is far below binary16's 2^-10 step, so it must collapse onto 1.0. + val src = fp32(1.0f + 1.0f / 1048576.0f, 3.14159265f, -2.71828f) + val out = ctx.ops.convert(src, FP16).data.copyToFloatArray() + + assertEquals(1.0f, out[0], "sub-ulp detail must be rounded away, not retained") + for ((i, v) in floatArrayOf(1.0f + 1.0f / 1048576.0f, 3.14159265f, -2.71828f).withIndex()) { + assertEquals( + Fp16Codec.decode(Fp16Codec.encode(v)), out[i], + "element $i must equal the binary16 rounding of the input", + ) + } + } + + @Test + fun convert_to_bf16_is_supported_and_truncates() { + val values = floatArrayOf(1.0f, 3.14159265f, -2.71828f, 1234.5f) + val out = ctx.ops.convert(fp32(*values), BF16).data.copyToFloatArray() + for ((i, v) in values.withIndex()) { + assertEquals( + Bf16Codec.decode(Bf16Codec.encode(v)), out[i], + "element $i must equal the bf16 encoding of the input", + ) + } + } + + @Test + fun convert_to_fp16_overflows_to_infinity_past_the_ceiling() { + // The old retag path kept 70000f verbatim in a tensor claiming to be FP16, a value + // binary16 cannot represent at all. + val out = ctx.ops.convert(fp32(70000f, -70000f), FP16).data.copyToFloatArray() + assertTrue(out[0].isInfinite() && out[0] > 0, "must overflow to +Inf, got ${out[0]}") + assertTrue(out[1].isInfinite() && out[1] < 0, "must overflow to -Inf, got ${out[1]}") + } + + @Test + fun convert_to_bf16_keeps_the_same_value_finite() { + // bf16 has FP32's exponent range, so the same input that overflows fp16 stays finite. + val out = ctx.ops.convert(fp32(70000f), BF16).data.copyToFloatArray() + assertTrue(out[0].isFinite(), "bf16 must not overflow at 70000, got ${out[0]}") + } + + @Test + fun convert_to_fp16_flushes_subnormal_underflow_to_zero() { + // 2^-30 is far below binary16's smallest subnormal (2^-24). + val out = ctx.ops.convert(fp32(1.0f / 1073741824.0f), FP16).data.copyToFloatArray() + assertEquals(0.0f, out[0], "must underflow to zero rather than retain the FP32 value") + } + + @Test + fun convert_reports_the_target_dtype() { + assertEquals(FP16::class, ctx.ops.convert(fp32(1f), FP16).dtype) + assertEquals(BF16::class, ctx.ops.convert(fp32(1f), BF16).dtype) + } + + @Test + fun convert_to_fp32_is_unchanged() { + val values = floatArrayOf(1.0f, 3.14159265f, -2.71828f) + val out = ctx.ops.convert(fp32(*values), FP32).data.copyToFloatArray() + for ((i, v) in values.withIndex()) assertEquals(v, out[i], "FP32 target must be exact") + } + + @Test + fun round_tripping_through_fp16_is_idempotent() { + // Rounding an already-rounded value must not move it again. + val src = fp32(3.14159265f, -2.71828f, 1234.5f) + val once = ctx.ops.convert(src, FP16).data.copyToFloatArray() + val twice = ctx.ops.convert( + ctx.fromFloatArray(Shape(once.size), FP32::class, once), FP16, + ).data.copyToFloatArray() + for (i in once.indices) assertEquals(once[i], twice[i], "element $i must be stable") + } +} diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 2ce8cc12..84f41025 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -5807,6 +5807,7 @@ public final class sk/ainet/lang/types/DType$Companion { public final class sk/ainet/lang/types/DTypeExtensionsKt { public static final fun commonPrecisionWith (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/types/DType; public static final fun isConvertibleTo (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;)Z + public static final fun isFloatingPoint (Lsk/ainet/lang/types/DType;)Z public static final fun kotlinClass (Lsk/ainet/lang/types/DType;)Lkotlin/reflect/KClass; } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt index 66e1f577..96fb66fe 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt @@ -314,7 +314,7 @@ public class DenseTensorDataFactory: TensorDataFactory { @Suppress("UNCHECKED_CAST") public fun fromFloatArray(data: FloatArray, dtype: T): TensorData { return when (dtype) { - is FP32, FP16 -> { + is FP32, is FP16, is BF16 -> { createFloatTensorData(Shape(data.size), data, dtype) as TensorData } @@ -383,6 +383,12 @@ public class DenseTensorDataFactory: TensorDataFactory { } } + /** + * The [DType] instance for a 16-bit float tag. Both formats are stored float-backed here; only + * the tag differs, and the tag is what the trace and StableHLO export read. + */ + private fun narrowTag(dtype: KClass<*>): DType = if (dtype == BF16::class) BF16 else FP16 + override fun ones(shape: Shape, dtype: KClass): TensorData { @Suppress("UNCHECKED_CAST") return when (dtype) { @@ -390,9 +396,11 @@ public class DenseTensorDataFactory: TensorDataFactory { val data = FloatArray(shape.volume) { 1.0f } createFloatTensorData(shape, data, FP32 as T) as TensorData } - FP16::class -> { + FP16::class, BF16::class -> { + // Float-backed, narrow-tagged (mirrors zeros): the dtype tag is what the DSL + // trace / StableHLO export reads to pick the MLIR element type. val data = FloatArray(shape.volume) { 1.0f } - createFloatTensorData(shape, data, FP16 as T) as TensorData + createFloatTensorData(shape, data, narrowTag(dtype) as T) as TensorData } Int32::class -> { val data = IntArray(shape.volume) { 1 } @@ -414,10 +422,12 @@ public class DenseTensorDataFactory: TensorDataFactory { val data = FloatArray(shape.volume) { floatValue } createFloatTensorData(shape, data, FP32 as T) as TensorData } - FP16::class -> { + FP16::class, BF16::class -> { + // Float-backed, narrow-tagged (mirrors zeros): the dtype tag is what the DSL + // trace / StableHLO export reads to pick the MLIR element type. val floatValue = value.toFloat() val data = FloatArray(shape.volume) { floatValue } - createFloatTensorData(shape, data, FP16 as T) as TensorData + createFloatTensorData(shape, data, narrowTag(dtype) as T) as TensorData } Int32::class -> { val intValue = value.toInt() @@ -466,7 +476,7 @@ public class DenseTensorDataFactory: TensorDataFactory { } createFloatTensorData(shape, data, FP32 as T) as TensorData } - FP16::class -> { + FP16::class, BF16::class -> { val data = FloatArray(shape.volume) var hasSpare = false var spare = 0.0f @@ -488,7 +498,7 @@ public class DenseTensorDataFactory: TensorDataFactory { hasSpare = true } } - createFloatTensorData(shape, data, FP16 as T) as TensorData + createFloatTensorData(shape, data, narrowTag(dtype) as T) as TensorData } else -> throw IllegalArgumentException("randn only supports floating point types: $dtype") } @@ -631,8 +641,8 @@ public class DenseTensorDataFactory: TensorDataFactory { FP32::class -> { createFloatTensorData(shape, data, FP32 as T) as TensorData } - FP16::class -> { - createFloatTensorData(shape, data, FP16 as T) as TensorData + FP16::class, BF16::class -> { + createFloatTensorData(shape, data, narrowTag(dtype) as T) as TensorData } else -> throw IllegalArgumentException("fromFloatArray only supports floating point types: $dtype") } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/DTypeExtensions.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/DTypeExtensions.kt index f382e16c..a5b24274 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/DTypeExtensions.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/DTypeExtensions.kt @@ -24,36 +24,32 @@ public fun DType.kotlinClass(): KClass<*> = when (this) { */ public fun DType.isConvertibleTo(target: DType): Boolean = when { this == target -> true - // Floating point conversions - this is FP32 && target is FP16 -> true - this is FP16 && target is FP32 -> true + // Any float width converts to any other — FP32 is a strict superset of both 16-bit + // formats, and narrowing is a well-defined (lossy) rounding. + this.isFloatingPoint() && target.isFloatingPoint() -> true // Integer conversions (with potential precision loss warnings) this is Int32 && target is Int8 -> true this is Int8 && target is Int32 -> true this is Int8 && target is Int4 -> true this is Int4 && target is Int8 -> true // Mixed float-int conversions - this is FP32 && target is Int32 -> true - this is FP16 && target is Int32 -> true - this is Int32 && target is FP32 -> true - this is Int32 && target is FP16 -> true + this.isFloatingPoint() && target is Int32 -> true + this is Int32 && target.isFloatingPoint() -> true // Ternary conversions this is Ternary && target is Int8 -> true this is Int8 && target is Ternary -> true else -> false } +/** True for the IEEE-style float types: [FP16], [BF16], [FP32], [FP64]. */ +public fun DType.isFloatingPoint(): Boolean = this is FP16 || this is BF16 || this is FP32 || this is FP64 + /** - * Returns the common precision type for mixed operations + * Returns the common precision type for mixed operations. + * + * Delegates to [DType.promoteTo], which is the exhaustive per-type lattice. This used to be a + * second, hand-written lattice that disagreed with it — notably it had no BF16 arm at all, so + * `BF16.commonPrecisionWith(Int8)` fell through to FP32 while `BF16.promoteTo(Int8)` said BF16. + * One lattice, one answer. */ -public fun DType.commonPrecisionWith(other: DType): DType = when { - this == other -> this - // Floating point takes precedence - this is FP32 || other is FP32 -> FP32 - this is FP16 || other is FP16 -> FP16 - // Higher precision integer takes precedence - this is Int32 || other is Int32 -> Int32 - this is Int8 || other is Int8 -> Int8 - this is Int4 || other is Int4 -> Int4 - else -> FP32 // Default fallback -} \ No newline at end of file +public fun DType.commonPrecisionWith(other: DType): DType = promoteTo(other) \ No newline at end of file diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP16.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP16.kt index b3e46f89..e72d2ca1 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP16.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP16.kt @@ -21,7 +21,9 @@ public object FP16 : DType { is Ternary -> FP16 // FP16 + Ternary → FP16 is Int4 -> FP16 // FP16 + Int4 → FP16 is Int8 -> FP16 // FP16 + Int8 → FP16 - is Int32 -> FP16 // FP16 + Int32 → FP16 + // Int32 needs 32 bits of integer precision; FP16 carries 11 mantissa bits, so + // promoting to FP16 would silently drop ~21 of them. Matches Int32.promoteTo(FP16). + is Int32 -> FP32 // FP16 + Int32 → FP32 is FP16 -> FP16 // FP16 + FP16 → FP16 is FP32 -> FP32 // FP16 + FP32 → FP32 is FP64 -> FP64 // FP16 + FP64 → FP64 diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP32.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP32.kt index af83d270..0e170e1e 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP32.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/FP32.kt @@ -6,7 +6,7 @@ public object FP32 : DType { override fun isCompatible(other: DType): Boolean { return when (other) { - is Ternary -> false // Ternary can promote to FP32 + is Ternary -> true // Ternary can promote to FP32 is Int4 -> true // Int4 can promote to FP32 is Int8 -> true // Int8 can promote to FP32 is Int32 -> true // Int32 can promote to FP32 @@ -18,7 +18,7 @@ public object FP32 : DType { override fun promoteTo(other: DType): DType { return when (other) { - is Ternary -> Ternary // FP32 + Ternary → FP32 + is Ternary -> FP32 // FP32 + Ternary → FP32 is Int4 -> FP32 // FP32 + Int4 → FP32 is Int8 -> FP32 // FP32 + Int8 → FP32 is Int32 -> FP32 // FP32 + Int32 → FP32 diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int8.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int8.kt index dc3b09cd..e4254607 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int8.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/types/Int8.kt @@ -29,7 +29,9 @@ public object Int8 : DType { is UInt32 -> Int64 // Int8 + UInt32 → Int64 is UInt64 -> FP64 // Int8 + UInt64 → FP64 (no integer type covers both) is FP16 -> FP16 // Int8 + FP16 → FP16 - is BF16 -> FP32 // Int8 + BF16 → FP32 + // bf16 carries 8 significand bits, so it represents every Int8 value (-128..127) + // exactly — same reason Int8 + FP16 stays FP16. Matches BF16.promoteTo(Int8). + is BF16 -> BF16 // Int8 + BF16 → BF16 is FP32 -> FP32 // Int8 + FP32 → FP32 is FP64 -> FP64 // Int8 + FP64 → FP64 } diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/DTypeLatticeConsistencyTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/DTypeLatticeConsistencyTest.kt new file mode 100644 index 00000000..f0fb10d6 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/types/DTypeLatticeConsistencyTest.kt @@ -0,0 +1,127 @@ +package sk.ainet.lang.types + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * The type lattice must have exactly one answer per question. + * + * `commonPrecisionWith` used to be a second, hand-written lattice that disagreed with + * [DType.promoteTo] — it had no BF16 arm at all, so `BF16.commonPrecisionWith(Int8)` fell through + * to FP32 while `BF16.promoteTo(Int8)` said BF16. It now delegates, and these tests pin that plus + * the symmetry and precision-safety properties the lattice is supposed to have. + */ +class DTypeLatticeConsistencyTest { + + private val floats = listOf(FP16, BF16, FP32, FP64) + private val all = listOf( + Ternary, Int4, Int8, Int16, Int32, Int64, + UInt8, UInt16, UInt32, UInt64, + FP16, BF16, FP32, FP64, + ) + + @Test + fun commonPrecisionWith_agrees_with_promoteTo_everywhere() { + for (a in all) { + for (b in all) { + assertEquals( + a.promoteTo(b), a.commonPrecisionWith(b), + "lattices disagree for ${a.name} + ${b.name}", + ) + } + } + } + + @Test + fun bf16_promotion_is_now_reachable_through_commonPrecisionWith() { + // The specific contradiction that motivated this: small ints stay in BF16. + assertEquals(BF16, BF16.commonPrecisionWith(Int8)) + assertEquals(BF16, BF16.commonPrecisionWith(Int4)) + assertEquals(BF16, BF16.commonPrecisionWith(Ternary)) + // ...while anything needing more range or precision escalates. + assertEquals(FP32, BF16.commonPrecisionWith(Int32)) + assertEquals(FP32, BF16.commonPrecisionWith(FP32)) + assertEquals(FP64, BF16.commonPrecisionWith(FP64)) + } + + @Test + fun the_two_narrow_floats_promote_to_fp32_in_both_directions() { + // Neither subsumes the other: fp16 has more mantissa, bf16 more exponent. + assertEquals(FP32, FP16.promoteTo(BF16)) + assertEquals(FP32, BF16.promoteTo(FP16)) + assertEquals(FP32, FP16.commonPrecisionWith(BF16)) + assertEquals(FP32, BF16.commonPrecisionWith(FP16)) + } + + @Test + fun float_promotion_is_symmetric() { + for (a in floats) { + for (b in floats) { + assertEquals( + a.promoteTo(b), b.promoteTo(a), + "asymmetric promotion: ${a.name}+${b.name} vs ${b.name}+${a.name}", + ) + } + } + } + + @Test + fun float_int_promotion_is_symmetric_for_the_common_types() { + // Regression: FP16.promoteTo(Int32) used to return FP16 while Int32.promoteTo(FP16) + // returned FP32 — the FP16 direction silently dropped ~21 bits of integer precision. + val ints = listOf(Int8, Int32) + for (f in floats) { + for (i in ints) { + assertEquals( + f.promoteTo(i), i.promoteTo(f), + "asymmetric promotion: ${f.name}+${i.name} vs ${i.name}+${f.name}", + ) + } + } + assertEquals(FP32, FP16.promoteTo(Int32), "Int32 does not fit in FP16's 11 mantissa bits") + } + + @Test + fun fp32_ternary_promotes_to_fp32_not_ternary() { + // The code returned Ternary while its own comment said FP32 — a promotion to a LOWER + // precision, which no promotion lattice should ever produce. + assertEquals(FP32, FP32.promoteTo(Ternary)) + assertTrue(FP32.isCompatible(Ternary), "Ternary promotes to FP32, so it is compatible") + } + + @Test + fun promotion_never_narrows_a_float() { + // A promotion result must be at least as wide as either operand. + for (a in floats) { + for (b in floats) { + val r = a.promoteTo(b) + assertTrue( + r.sizeInBits >= maxOf(a.sizeInBits, b.sizeInBits), + "${a.name}+${b.name} promoted to ${r.name}, narrower than an operand", + ) + } + } + } + + @Test + fun narrow_floats_are_convertible_in_every_float_direction() { + for (a in floats) { + for (b in floats) { + assertTrue(a.isConvertibleTo(b), "${a.name} -> ${b.name} should be convertible") + } + } + // BF16 was entirely absent from isConvertibleTo, so these all returned false. + assertTrue(BF16.isConvertibleTo(FP32)) + assertTrue(FP32.isConvertibleTo(BF16)) + assertTrue(BF16.isConvertibleTo(Int32)) + assertTrue(Int32.isConvertibleTo(BF16)) + } + + @Test + fun isFloatingPoint_covers_exactly_the_float_types() { + for (t in all) { + assertEquals(t in floats, t.isFloatingPoint(), "isFloatingPoint wrong for ${t.name}") + } + } +}