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 b3ebf43a..d1126a82 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 @@ -134,6 +134,7 @@ public interface KernelProvider { return when (weight) { "Float32" -> matmulFp32() != null "BFloat16" -> matmulBf16() != null + "Float16" -> matmulFp16() != null "Q4_K" -> matmulQ4K() != null "Q8_0" -> matmulQ8_0() != null "Q4_0" -> matmulQ4_0() != null 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 index bd68dd57..56c64300 100644 --- 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 @@ -1,6 +1,8 @@ package sk.ainet.exec.kernel import jdk.incubator.vector.FloatVector +import jdk.incubator.vector.IntVector +import jdk.incubator.vector.VectorOperators import jdk.incubator.vector.VectorSpecies import sk.ainet.backend.api.kernel.Fp16MatmulKernel import sk.ainet.lang.types.Fp16Codec @@ -15,16 +17,35 @@ import sk.ainet.lang.types.Fp16Codec * 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. + * only the widening is not. * - * Numerical parity vs [ScalarFp16MatmulKernel] is asserted by + * **Why not [Fp16Codec] (#887).** The codec is portable integer bit math, and calling it once per + * weight element made this kernel run at a flat ~0.5 GFLOP/s — 2-18x *slower* than the FP32 SGEMM + * it replaces, while the structurally identical BF16 kernel was 1.5-2.1x faster. `float16ToFloat` + * is a HotSpot intrinsic (JDK 20+) that lowers to a single `vcvtph2ps` on F16C hardware and to a + * compact branch-free sequence elsewhere, so the scratch fill stops dominating the inner loop. + * The kernel is JVM-only and this provider already gates on JDK 21+, so the intrinsic is always + * available where this code runs. + * + * The substitution is exact: the JDK conversion and [Fp16Codec.decode] agree bit-for-bit on all + * 65536 inputs, which `Fp16CodecIntrinsicParityTest` asserts exhaustively. Reaching that agreement + * is why the codec now quiets NaN — the hardware conversion does, so the codec was aligned with it + * in the same change rather than the kernel being held back. Numerical parity vs + * [ScalarFp16MatmulKernel] — which still goes through the codec — is asserted by * `PanamaVectorFp16MatmulKernelParityTest`. */ public object PanamaVectorFp16MatmulKernel : Fp16MatmulKernel { private val floatSpecies: VectorSpecies = FloatVector.SPECIES_PREFERRED + /** + * Derived from [floatSpecies]' shape rather than taken as `SPECIES_PREFERRED`, so the two + * always have the same lane count — [widen] reinterprets between them lanewise. + */ + private val intSpecies: VectorSpecies = IntVector.SPECIES_PREFERRED.withShape( + floatSpecies.vectorShape(), + ) as VectorSpecies + override fun matmul( a: FloatArray, aOffset: Int, aStride: Int, b: ByteArray, bByteOffset: Int, bByteStride: Int, @@ -45,7 +66,8 @@ public object PanamaVectorFp16MatmulKernel : Fp16MatmulKernel { val laneCount = floatSpecies.length() val bound = floatSpecies.loopBound(n) - val scratch = FloatArray(laneCount) + // Raw 16-bit patterns, not decoded floats: the widening happens in the vector domain. + val scratch = IntArray(laneCount) // Zero the output block first — the i-p-j outer product accumulates into it. for (i in 0 until m) { @@ -66,9 +88,9 @@ public object PanamaVectorFp16MatmulKernel : Fp16MatmulKernel { 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) + scratch[lane] = (hi shl 8) or lo } - val bVec = FloatVector.fromArray(floatSpecies, scratch, 0) + val bVec = widen(IntVector.fromArray(intSpecies, scratch, 0)) val outVec = FloatVector.fromArray(floatSpecies, out, outRowOff + j) aBcast.fma(bVec, outVec).intoArray(out, outRowOff + j) j += laneCount @@ -78,10 +100,67 @@ public object PanamaVectorFp16MatmulKernel : Fp16MatmulKernel { 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) + out[outRowOff + j] += aIp * halfToFloat(lo, hi) j++ } } } } + + /** + * Widen one little-endian binary16 element to FP32, for the scalar tail. + * + * `toShort()` keeps the low 16 bits, which is exactly the packed element; the sign extension + * that produces is what `float16ToFloat` expects. The tail runs at most `laneCount - 1` times + * per row, so the intrinsic is enough here and the vector path is reserved for [widen]. + */ + private fun halfToFloat(lo: Int, hi: Int): Float = + java.lang.Float.float16ToFloat((((hi shl 8) or lo).toShort())) + + /** + * Widen a whole vector of binary16 patterns to FP32, branch-free. + * + * The classic shift-and-rebias conversion, done lanewise. Shifting the sign-free pattern left + * by 13 lands binary16's exponent and mantissa in FP32's positions; adding `(127 - 15) << 23` + * rebiases the exponent. Two cases need a correction on top, and both are applied under a mask + * rather than a branch: + * + * - **Inf/NaN** (exponent all ones) needs a second `(128 - 16) << 23`, which saturates the + * FP32 exponent to all ones. + * - **Zero and subnormals** (exponent zero) are renormalized by the FPU instead of by a loop: + * bump the exponent by one and subtract `2⁻¹⁴`. For a subnormal `m * 2⁻²⁴` the bumped value + * is `2⁻¹⁴ * (1 + m * 2⁻¹⁰)`, so the subtraction leaves exactly `m * 2⁻²⁴`; for zero it + * leaves `+0`, and the sign is reapplied afterwards either way. + * + * A signaling NaN stays signaling here, where [Fp16Codec.decode] would quiet it. That is not + * observable: every lane feeds the FMA below, and the FMA quiets it. The exhaustive kernel + * sweep in `PanamaVectorFp16MatmulKernelParityTest` asserts bit equality with the codec on + * every non-NaN pattern and NaN-ness on the rest, which is exactly this contract. + */ + private fun widen(h: IntVector): FloatVector { + val shifted = h.and(0x7FFF).lanewise(VectorOperators.LSHL, 13) + val expField = shifted.and(EXP_FIELD) + + var biased = shifted.add(EXP_REBIAS) + biased = biased.add(EXP_REBIAS, expField.compare(VectorOperators.EQ, EXP_FIELD)) + + val renormalized = biased.add(SUBNORMAL_BUMP).reinterpretAsFloats().sub(SUBNORMAL_MAGIC) + val subnormal = expField.compare(VectorOperators.EQ, 0).cast(floatSpecies) + + val magnitude = biased.reinterpretAsFloats().blend(renormalized, subnormal) + val sign = h.and(0x8000).lanewise(VectorOperators.LSHL, 16) + return magnitude.reinterpretAsInts().or(sign).reinterpretAsFloats() + } + + /** binary16's exponent field once shifted into FP32 position: `0x7C00 shl 13`. */ + private const val EXP_FIELD = 0x0F80_0000 + + /** `(127 - 15) shl 23` — the FP32/binary16 exponent bias difference. */ + private const val EXP_REBIAS = 0x3800_0000 + + /** `1 shl 23` — one exponent step, to lift a subnormal into the magic constant's binade. */ + private const val SUBNORMAL_BUMP = 0x0080_0000 + + /** `2⁻¹⁴` (bits `113 shl 23`) — binary16's smallest normal, subtracted to renormalize. */ + private val SUBNORMAL_MAGIC: Float = Float.fromBits(0x3880_0000) } 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 index 6a8630cd..a7129dd7 100644 --- 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 @@ -4,6 +4,7 @@ import sk.ainet.lang.types.Fp16Codec import kotlin.math.abs import kotlin.random.Random import kotlin.test.Test +import kotlin.test.assertEquals import kotlin.test.assertTrue /** @@ -95,6 +96,49 @@ class PanamaVectorFp16MatmulKernelParityTest { assertTrue(PanamaVectorFp16MatmulKernel.codec === Fp16Codec) } + @Test + fun panama_decode_matches_the_codec_on_every_bit_pattern() { + // The Panama kernel widens through Float.float16ToFloat while the scalar kernel goes + // through Fp16Codec (#887). A 1xN matmul with a = [1] and a zeroed accumulator makes + // out[j] the decoded weight exactly — 1*x + 0 is exact — so this compares the two decode + // paths directly over the whole domain rather than sampling. + // + // Chunked at 999 columns so every chunk has a vectorized body plus a scalar tail for any + // lane count in {2, 4, 8, 16}, and the tail lands on different patterns in each chunk. + val chunk = 999 + val a = floatArrayOf(1f) + var base = 0 + while (base <= 0xFFFF) { + val n = minOf(chunk, 0x1_0000 - base) + val b = ByteArray(n * 2) + for (j in 0 until n) { + val bits = base + j + b[j * 2] = (bits and 0xFF).toByte() + b[j * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + val out = FloatArray(n) + PanamaVectorFp16MatmulKernel.matmul(a, 0, 1, b, 0, n * 2, out, 0, n, 1, n, 1) + + for (j in 0 until n) { + val bits = base + j + val expected = Fp16Codec.decode(bits) + val actual = out[j] + when { + // FMA quiets signaling NaNs and may rewrite the payload, so only NaN-ness is + // meaningful here; the payload itself is pinned by Fp16CodecIntrinsicParityTest. + expected.isNaN() -> assertTrue(actual.isNaN(), "expected NaN at 0x${bits.toString(16)}") + // Accumulating -0 into +0 yields +0, so the sign of zero cannot survive a matmul. + expected == 0f -> assertTrue(actual == 0f, "expected zero at 0x${bits.toString(16)}") + else -> assertEquals( + expected.toRawBits(), actual.toRawBits(), + "decode diverged at 0x${bits.toString(16)}: codec=$expected panama=$actual", + ) + } + } + base += n + } + } + @Test fun fp16_kernel_result_tracks_an_exact_fp32_reference() { // With operands exactly representable in binary16, the kernel must reproduce a plain diff --git a/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt b/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt index cc42e955..88d3bfd3 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt +++ b/skainet-backends/skainet-backend-native-cpu/native/CMakeLists.txt @@ -16,6 +16,7 @@ set(SKAINET_KERNEL_SOURCES src/q6k_matmul.c src/fp32_matmul.c src/bf16_matmul.c + src/fp16_matmul.c src/q8_0_matmul.c src/q4_0_matmul.c ) diff --git a/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h b/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h index 6d58e7eb..54424ea8 100644 --- a/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h +++ b/skainet-backends/skainet-backend-native-cpu/native/include/skainet_kernels.h @@ -147,6 +147,23 @@ SKAINET_API void skainet_bf16_matmul( int32_t m, int32_t n, int32_t k ); +/* + * Row-major FP32 × FP16 matmul: C(m, n) = A(m, k) * B(k, n). + * + * Identical contract to skainet_bf16_matmul, with B packed as IEEE + * binary16 little-endian (2 bytes per element) instead of BF16. + * + * FP16 → FP32 needs exponent rebiasing and gradual-underflow handling + * rather than BF16's single shift; it is done branch-free so the inner + * loop still vectorizes. See src/fp16_matmul.c. + */ +SKAINET_API void skainet_fp16_matmul( + const float* a, int32_t a_offset, int32_t a_stride, + const uint8_t* b, int32_t b_byte_offset, int32_t b_byte_stride, + float* c, int32_t c_offset, int32_t c_stride, + int32_t m, int32_t n, int32_t k +); + /* * Q8_0 matrix-vector multiply. * diff --git a/skainet-backends/skainet-backend-native-cpu/native/src/fp16_matmul.c b/skainet-backends/skainet-backend-native-cpu/native/src/fp16_matmul.c new file mode 100644 index 00000000..54297d6b --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/native/src/fp16_matmul.c @@ -0,0 +1,179 @@ +#include "skainet_kernels.h" + +#include +#include +#include + +/* + * Native row-major FP32 × FP16 matmul matching the + * sk.ainet.backend.api.kernel.Fp16MatmulKernel SPI: + * + * C(m, n) = A(m, k) * B(k, n) + * + * A, C : FP32 (FloatArray on the JVM side; float* here). + * Strides in floats. `a_stride == k` for a contiguous parent. + * B : packed IEEE binary16 (ByteArray on the JVM side; uint8_t* here). + * Strides in *bytes*. `b_byte_stride == n * 2` for a contiguous + * parent. Each FP16 value is 2 bytes little-endian. + * + * The BF16 sibling gets its dequant for free — BF16 is the high half of an + * FP32, so the conversion is one shift. Binary16 has a narrower exponent and + * a wider mantissa, so it needs rebiasing, and subnormals need renormalizing. + * Doing that with branches would cost more than the multiply it feeds, so + * fp16_to_float below is branch-free: the two special cases are folded in + * with arithmetic masks, which keeps the inner loop a straight-line sequence + * the vectorizer can widen (unpack to 32-bit lanes, integer ops, one FMA). + * + * Deliberately *not* using _Float16 or F16C intrinsics. The x86_64 build + * carries no -march flag, so F16C cannot be assumed, and _Float16 without it + * lowers to libgcc helper calls that are slower than this and block + * vectorization outright. AArch64 does build with +fp16, but a second code + * path would double the surface to test for a conversion that is already a + * handful of integer ops. Runtime ISA dispatch is the place for that, if it + * ever pays for itself. + * + * Iteration order depends on m. + * + * At m == 1 it is plain i-p-j: each B element is used exactly once, so there + * is nothing to reuse and the best thing to do is stream B sequentially. + * Tiling here is pure overhead — measured 15% slower on ffn_up 8B — and m == 1 + * is the decode step of inference. + * + * At m > 1, i-p-j would walk the whole of B once per row of A: for ffn_up 8B + * at m=16 that is 16 passes over 90 MiB. Instead j is tiled, and within a tile + * each B row is decoded once into a small stack buffer and multiplied into all + * m rows of C. B is then read once in total rather than m times, and the + * decode count drops from m*k*n to k*n. The tile is sized so the decoded row + * and the m C rows it feeds stay resident together. + * + * Accumulation order into any given C element is p ascending on both paths, so + * the two are bit-identical to each other and to the original i-p-j + * formulation, not merely close. + * + * Caller contract (mirrors skainet_bf16_matmul): + * - C is FULLY OVERWRITTEN in the m×n block. + * - k == 0 zeros the m×n block. + * - m == 0 || n == 0 is a no-op. + * - Negative m / n / k are caller errors; defensively treated as no-op. + * + * NaN note: a signaling binary16 NaN stays signaling here, matching the JVM + * Panama kernel. It is not observable — the value is immediately multiplied, + * and that quiets it. + */ + +/* binary16's exponent field once shifted into FP32 position (0x7C00 << 13). */ +#define SKAINET_FP16_EXP_FIELD 0x0F800000u +/* (127 - 15) << 23 — the FP32/binary16 exponent bias difference. */ +#define SKAINET_FP16_REBIAS 0x38000000u +/* 1 << 23 — one exponent step, lifting a subnormal into the magic binade. */ +#define SKAINET_FP16_SUBNORMAL_BUMP 0x00800000u + +/* + * Columns decoded per pass. 512 floats is a 2 KiB stack buffer — small enough + * to leave the C rows it feeds resident alongside it, large enough that the + * per-tile loop overhead disappears against the k*m inner work. + */ +#define SKAINET_FP16_TILE 512 + +static inline float skainet_fp16_to_float(uint16_t h) { + const uint32_t bits = (uint32_t) h; + const uint32_t sign = (bits & 0x8000u) << 16; + const uint32_t shifted = (bits & 0x7FFFu) << 13; + const uint32_t exp_field = shifted & SKAINET_FP16_EXP_FIELD; + + /* 0 or 0xFFFFFFFF, so the corrections below are masks rather than jumps. */ + const uint32_t is_inf_nan = + (uint32_t) -(int32_t) (exp_field == SKAINET_FP16_EXP_FIELD); + const uint32_t is_subnormal = (uint32_t) -(int32_t) (exp_field == 0u); + + /* Rebias; Inf/NaN takes a second rebias, which saturates the exponent. */ + const uint32_t biased = + shifted + SKAINET_FP16_REBIAS + (SKAINET_FP16_REBIAS & is_inf_nan); + + /* Zero and subnormals: bump one exponent step and subtract 2^-14, which + * makes the FPU renormalize. For a subnormal m * 2^-24 the bumped value is + * 2^-14 * (1 + m * 2^-10), so the subtraction leaves exactly m * 2^-24; + * for zero it leaves +0. The sign is reapplied afterwards either way. */ + const uint32_t bumped = biased + SKAINET_FP16_SUBNORMAL_BUMP; + float renormalized; + memcpy(&renormalized, &bumped, sizeof(float)); + renormalized -= 0x1p-14f; + uint32_t renormalized_bits; + memcpy(&renormalized_bits, &renormalized, sizeof(uint32_t)); + + const uint32_t magnitude = + (renormalized_bits & is_subnormal) | (biased & ~is_subnormal); + const uint32_t out_bits = sign | magnitude; + + float out; + memcpy(&out, &out_bits, sizeof(float)); + return out; +} + +SKAINET_API void skainet_fp16_matmul( + const float* SKAINET_RESTRICT a, int32_t a_offset, int32_t a_stride, + const uint8_t* SKAINET_RESTRICT b, int32_t b_byte_offset, int32_t b_byte_stride, + float* SKAINET_RESTRICT c, int32_t c_offset, int32_t c_stride, + int32_t m, int32_t n, int32_t k +) { + if (m <= 0 || n <= 0) return; + + /* Zero the output block. Required by the SPI for k == 0 AND a + * prerequisite for the i-p-j accumulator below. */ + for (int32_t i = 0; i < m; ++i) { + float* SKAINET_RESTRICT c_row = c + c_offset + (size_t) i * c_stride; + for (int32_t j = 0; j < n; ++j) { + c_row[j] = 0.0f; + } + } + if (k <= 0) return; + + if (m == 1) { + const float* SKAINET_RESTRICT a_row = a + a_offset; + float* SKAINET_RESTRICT c_row = c + c_offset; + for (int32_t p = 0; p < k; ++p) { + const float a_ip = a_row[p]; + const uint8_t* SKAINET_RESTRICT b_row = + b + b_byte_offset + (size_t) p * b_byte_stride; + for (int32_t j = 0; j < n; ++j) { + /* Read 2 bytes LE. memcpy is strict-aliasing safe and the + * compiler folds it to a single 16-bit load. */ + uint16_t bits; + memcpy(&bits, b_row + (size_t) j * 2, sizeof(uint16_t)); + c_row[j] += a_ip * skainet_fp16_to_float(bits); + } + } + return; + } + + float decoded[SKAINET_FP16_TILE]; + + for (int32_t j0 = 0; j0 < n; j0 += SKAINET_FP16_TILE) { + const int32_t tile = + (n - j0) < SKAINET_FP16_TILE ? (n - j0) : SKAINET_FP16_TILE; + + for (int32_t p = 0; p < k; ++p) { + const uint8_t* SKAINET_RESTRICT b_row = + b + b_byte_offset + (size_t) p * b_byte_stride + (size_t) j0 * 2; + + /* Decode this row's tile once, for all m rows of C below. */ + for (int32_t j = 0; j < tile; ++j) { + /* Read 2 bytes LE. memcpy is strict-aliasing safe and the + * compiler folds it to a single 16-bit load. */ + uint16_t bits; + memcpy(&bits, b_row + (size_t) j * 2, sizeof(uint16_t)); + decoded[j] = skainet_fp16_to_float(bits); + } + + for (int32_t i = 0; i < m; ++i) { + const float a_ip = + a[a_offset + (size_t) i * a_stride + p]; + float* SKAINET_RESTRICT c_row = + c + c_offset + (size_t) i * c_stride + j0; + for (int32_t j = 0; j < tile; ++j) { + c_row[j] += a_ip * decoded[j]; + } + } + } + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp16MatmulKernel.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp16MatmulKernel.kt new file mode 100644 index 00000000..a499f048 --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeFp16MatmulKernel.kt @@ -0,0 +1,112 @@ +package sk.ainet.exec.kernel + +import java.lang.foreign.Arena +import java.lang.foreign.FunctionDescriptor +import java.lang.foreign.Linker +import java.lang.foreign.MemorySegment +import java.lang.foreign.ValueLayout +import java.lang.invoke.MethodHandle +import sk.ainet.backend.api.kernel.Fp16MatmulKernel + +/** + * Native (FFM) implementation of [Fp16MatmulKernel]. + * + * Wraps the bundled C symbol + * + * void skainet_fp16_matmul( + * const float* a, int32_t a_offset, int32_t a_stride, + * const uint8_t* b, int32_t b_byte_offset, int32_t b_byte_stride, + * float* c, int32_t c_offset, int32_t c_stride, + * int32_t m, int32_t n, int32_t k); + * + * Mirrors [NativeBf16MatmulKernel] exactly apart from the symbol and the + * dequant the C side performs: binary16 needs exponent rebiasing and + * gradual-underflow handling where BF16 needs one shift, so the C kernel does + * it branch-free to keep the inner loop vectorizable. + * + * **Why this exists (#887).** Until it did, [NativeKernelProvider] carried + * `matmulBf16` but not `matmulFp16`, so BF16 resolved to a native kernel while + * FP16 fell through to `PanamaVectorFp16MatmulKernel` at priority 50. That + * asymmetry — not the cost of the decode, which is within ~15% between the two + * Panama kernels — is what made FP16 measure 2-18x slower than the FP32 SGEMM + * while BF16 measured faster than it. + * + * Numerical parity vs `ScalarFp16MatmulKernel` is asserted by + * `NativeFp16MatmulKernelParityTest` within FMA + reordered-reduction + * tolerance — the same bar the BF16 and FP32 native parity tests use. + */ +internal object NativeFp16MatmulKernel : Fp16MatmulKernel { + + fun isAvailable(): Boolean = handle != null + + 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) { + "NativeFp16MatmulKernel: m, n, k must be non-negative; got m=$m n=$n k=$k" + } + if (m == 0 || n == 0) return + + val mh = handle + ?: error("NativeFp16MatmulKernel.matmul invoked while native library unavailable") + + // Reach calculations. For non-contiguous strides we may skip past + // unused elements; allocating to the full reach keeps the kernel's + // pointer arithmetic simple and matches the Kotlin-side bounds. + val aReachFloats = if (m == 0 || k == 0) 0 else aOffset + (m - 1) * aStride + k + val bReachBytes = if (k == 0 || n == 0) 0 + else bByteOffset + (k - 1) * bByteStride + n * 2 + val cReachFloats = outOffset + (m - 1) * outStride + n + + Arena.ofConfined().use { arena -> + val aBytes = aReachFloats.toLong() * java.lang.Float.BYTES + val bBytes = bReachBytes.toLong() + val cBytes = cReachFloats.toLong() * java.lang.Float.BYTES + val fAlign = ValueLayout.JAVA_FLOAT.byteAlignment() + val bAlign = ValueLayout.JAVA_BYTE.byteAlignment() + + val aSeg: MemorySegment = if (aBytes > 0) arena.allocate(aBytes, fAlign) else MemorySegment.NULL + val bSeg: MemorySegment = if (bBytes > 0) arena.allocate(bBytes, bAlign) else MemorySegment.NULL + val cSeg: MemorySegment = arena.allocate(cBytes, fAlign) + + if (aReachFloats > 0) { + MemorySegment.copy(a, 0, aSeg, ValueLayout.JAVA_FLOAT, 0L, aReachFloats) + } + if (bReachBytes > 0) { + MemorySegment.copy(b, 0, bSeg, ValueLayout.JAVA_BYTE, 0L, bReachBytes) + } + + mh.invoke( + aSeg, aOffset, aStride, + bSeg, bByteOffset, bByteStride, + cSeg, outOffset, outStride, + m, n, k, + ) + + MemorySegment.copy(cSeg, ValueLayout.JAVA_FLOAT, 0L, out, 0, cReachFloats) + } + } + + private val handle: MethodHandle? by lazy { + val lookup = NativeLibraryLoader.lookup() ?: return@lazy null + val symbol = lookup.find("skainet_fp16_matmul").orElse(null) ?: return@lazy null + val descriptor = FunctionDescriptor.ofVoid( + ValueLayout.ADDRESS, // a + ValueLayout.JAVA_INT, // a_offset + ValueLayout.JAVA_INT, // a_stride + ValueLayout.ADDRESS, // b + ValueLayout.JAVA_INT, // b_byte_offset + ValueLayout.JAVA_INT, // b_byte_stride + ValueLayout.ADDRESS, // c + ValueLayout.JAVA_INT, // c_offset + ValueLayout.JAVA_INT, // c_stride + ValueLayout.JAVA_INT, // m + ValueLayout.JAVA_INT, // n + ValueLayout.JAVA_INT, // k + ) + runCatching { Linker.nativeLinker().downcallHandle(symbol, descriptor) }.getOrNull() + } +} diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeKernelProvider.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeKernelProvider.kt index f52221f8..d15b37a7 100644 --- a/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeKernelProvider.kt +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmMain/kotlin/sk/ainet/exec/kernel/NativeKernelProvider.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.MemSegKernelProvider @@ -75,6 +76,10 @@ import sk.ainet.backend.api.kernel.Q8_0MatmulKernel * - PR 3: MemSeg-input zero-copy sibling. * - PR 5: native FP32 matmul wired into [matmulFp32]. * - Now: native `matmulQ5K`, `matmulQ6K`, `matmulQ8_0`, `matmulQ4_0` all wired. + * - Now: native `matmulFp16`, closing the gap against `matmulBf16` (#887). + * Every narrow-float accessor the SPI declares is wired here; a format + * served natively on one side and by the JVM fallback on the other looks + * like a slow kernel rather than a missing one. */ public object NativeKernelProvider : KernelProvider, MemSegKernelProvider { override val name: String = "native-ffm" @@ -94,6 +99,9 @@ public object NativeKernelProvider : KernelProvider, MemSegKernelProvider { override fun matmulBf16(): Bf16MatmulKernel? = if (NativeBf16MatmulKernel.isAvailable()) NativeBf16MatmulKernel else null + override fun matmulFp16(): Fp16MatmulKernel? = + if (NativeFp16MatmulKernel.isAvailable()) NativeFp16MatmulKernel else null + override fun matmulQ8_0(): Q8_0MatmulKernel? = if (NativeQ8_0MatmulKernel.isAvailable()) NativeQ8_0MatmulKernel else null diff --git a/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFp16MatmulKernelParityTest.kt b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFp16MatmulKernelParityTest.kt new file mode 100644 index 00000000..9a1d170c --- /dev/null +++ b/skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/NativeFp16MatmulKernelParityTest.kt @@ -0,0 +1,335 @@ +package sk.ainet.exec.kernel + +import sk.ainet.lang.types.Fp16Codec +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * Numerical parity tests for [NativeFp16MatmulKernel] against + * [PanamaVectorFp16MatmulKernel]. Mirrors `NativeBf16MatmulKernelParityTest`, + * same `1e-5 * k` bar clamped to a `1e-5` floor. + * + * The two kernels reach the same values by different routes — the C side + * folds the special cases with integer masks, the JVM side with vector masks + * — so [decode_matches_the_codec_on_every_bit_pattern] pins the conversion + * itself exhaustively rather than trusting the sampled shapes below to have + * covered a subnormal or an infinity. + */ +class NativeFp16MatmulKernelParityTest { + + @BeforeTest + fun checkAvailable() { + assertTrue( + NativeFp16MatmulKernel.isAvailable(), + "Native FP16 kernel must be available — bundled libskainet_kernels missing " + + "or skainet_fp16_matmul symbol unresolved", + ) + } + + /** Encode FP32 into binary16, store little-endian in a byte buffer. */ + 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, + a: FloatArray, aOffset: Int, aStride: Int, + b: ByteArray, bByteOffset: Int, bByteStride: Int, + outStride: Int, + tolScale: Float = 1e-5f, + ) { + val outPanama = FloatArray(m * outStride) + val outNative = FloatArray(m * outStride) + PanamaVectorFp16MatmulKernel.matmul( + a, aOffset, aStride, + b, bByteOffset, bByteStride, + outPanama, 0, outStride, + m, n, k, + ) + NativeFp16MatmulKernel.matmul( + a, aOffset, aStride, + b, bByteOffset, bByteStride, + outNative, 0, outStride, + m, n, k, + ) + val tol = (tolScale * k.coerceAtLeast(1)).coerceAtLeast(tolScale) + assertEquals(outPanama.size, outNative.size, "length mismatch") + for (i in outPanama.indices) { + val diff = abs(outPanama[i] - outNative[i]) + assertTrue( + diff <= tol, + "mismatch at $i: panama=${outPanama[i]} native=${outNative[i]} diff=$diff tol=$tol", + ) + } + } + + @Test + fun decode_matches_the_codec_on_every_bit_pattern() { + // A 1xN matmul with a = [1] into a zeroed accumulator makes out[j] the + // decoded weight exactly (1*x + 0 is exact), so this compares the C + // conversion against Fp16Codec over the whole 16-bit domain. + val n = 0x1_0000 + val b = ByteArray(n * 2) + for (bits in 0 until n) { + b[bits * 2] = (bits and 0xFF).toByte() + b[bits * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + val out = FloatArray(n) + NativeFp16MatmulKernel.matmul( + floatArrayOf(1f), 0, 1, + b, 0, n * 2, + out, 0, n, + 1, n, 1, + ) + + for (bits in 0 until n) { + val expected = Fp16Codec.decode(bits) + val actual = out[bits] + when { + // The multiply quiets a signaling NaN, so only NaN-ness is + // meaningful — same contract the Panama kernel sweep pins. + expected.isNaN() -> assertTrue(actual.isNaN(), "expected NaN at 0x${bits.toString(16)}") + // Accumulating -0 into +0 yields +0; the sign of zero cannot survive. + expected == 0f -> assertTrue(actual == 0f, "expected zero at 0x${bits.toString(16)}") + else -> assertEquals( + expected.toRawBits(), actual.toRawBits(), + "decode diverged at 0x${bits.toString(16)}: codec=$expected native=$actual", + ) + } + } + } + + @Test + fun decode_matches_the_codec_on_every_bit_pattern_through_the_tiled_path() { + // The sweep above uses m = 1, which takes the straight i-p-j path. At + // m > 1 the kernel tiles j and decodes through a different loop, so + // sweep the domain again with a second row to cover it. Row 1's a is 0, + // so row 0 still reads out as the decoded weight exactly. + val n = 0x1_0000 + val b = ByteArray(n * 2) + for (bits in 0 until n) { + b[bits * 2] = (bits and 0xFF).toByte() + b[bits * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + val out = FloatArray(2 * n) + NativeFp16MatmulKernel.matmul( + floatArrayOf(1f, 0f), 0, 1, + b, 0, n * 2, + out, 0, n, + 2, n, 1, + ) + + for (bits in 0 until n) { + val expected = Fp16Codec.decode(bits) + val actual = out[bits] + when { + expected.isNaN() -> assertTrue(actual.isNaN(), "expected NaN at 0x${bits.toString(16)}") + expected == 0f -> assertTrue(actual == 0f, "expected zero at 0x${bits.toString(16)}") + else -> assertEquals( + expected.toRawBits(), actual.toRawBits(), + "tiled decode diverged at 0x${bits.toString(16)}: codec=$expected native=$actual", + ) + } + } + } + + @Test + fun multi_tile_n_with_partial_last_tile_matches_panama() { + // At m > 1 the kernel tiles j at 512 columns. The other shapes here are + // either m == 1 or n <= 256, so without this case the tiled path runs as + // a single full tile and the tile-boundary arithmetic is never + // exercised. n = 1100 is two full tiles plus a 76-column remainder. + val rng = Random(7) + val m = 3; val n = 1100; val k = 17 + val a = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val bFloats = FloatArray(k * n) { rng.nextFloat() - 0.5f } + val b = fp16Bytes(bFloats) + assertParity( + m = m, n = n, k = k, + a = a, aOffset = 0, aStride = k, + b = b, bByteOffset = 0, bByteStride = n * 2, + outStride = n, + ) + } + + @Test + fun tiled_and_single_row_paths_agree_on_the_same_weights() { + // m == 1 and m > 1 take different loop orders. Accumulation stays p + // ascending in both, so row 0 of a multi-row call must be bit-identical + // to the same row computed on its own — not merely within tolerance. + val rng = Random(31) + val n = 700; val k = 9 + val bFloats = FloatArray(k * n) { rng.nextFloat() - 0.5f } + val b = fp16Bytes(bFloats) + val aRow = FloatArray(k) { rng.nextFloat() - 0.5f } + + val single = FloatArray(n) + NativeFp16MatmulKernel.matmul(aRow, 0, k, b, 0, n * 2, single, 0, n, 1, n, k) + + val a2 = FloatArray(2 * k) + aRow.copyInto(a2, 0) + aRow.copyInto(a2, k) + val pair = FloatArray(2 * n) + NativeFp16MatmulKernel.matmul(a2, 0, k, b, 0, n * 2, pair, 0, n, 2, n, k) + + for (j in 0 until n) { + assertEquals( + single[j].toRawBits(), pair[j].toRawBits(), + "tiled path diverged from the single-row path at column $j", + ) + } + } + + @Test + fun small_2x3x4_contiguous_matches_panama() { + val a = floatArrayOf(1f, 2f, 3f, 4f, 5f, 6f, 7f, 8f) // [2, 4] + val bFloats = FloatArray(4 * 3) { it.toFloat() } // [4, 3] + val b = fp16Bytes(bFloats) + assertParity( + m = 2, n = 3, k = 4, + a = a, aOffset = 0, aStride = 4, + b = b, bByteOffset = 0, bByteStride = 3 * 2, + outStride = 3, + ) + } + + @Test + fun random_8x16x32_matches_panama() { + val rng = Random(42) + val a = FloatArray(8 * 32) { rng.nextFloat() - 0.5f } + val bFloats = FloatArray(32 * 16) { rng.nextFloat() - 0.5f } + val b = fp16Bytes(bFloats) + assertParity( + m = 8, n = 16, k = 32, + a = a, aOffset = 0, aStride = 32, + b = b, bByteOffset = 0, bByteStride = 16 * 2, + outStride = 16, + ) + } + + @Test + fun non_aligned_n_exercises_tail_loop() { + val rng = Random(1234) + val m = 5; val n = 7; val k = 23 + val a = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val bFloats = FloatArray(k * n) { rng.nextFloat() - 0.5f } + val b = fp16Bytes(bFloats) + assertParity( + m = m, n = n, k = k, + a = a, aOffset = 0, aStride = k, + b = b, bByteOffset = 0, bByteStride = n * 2, + outStride = n, + ) + } + + @Test + fun strided_a_sub_block_matches_panama() { + val parentA = FloatArray(4 * 8) { it.toFloat() } + val bFloats = FloatArray(8 * 3) { (it + 1).toFloat() } + val b = fp16Bytes(bFloats) + assertParity( + m = 2, n = 3, k = 8, + a = parentA, aOffset = 1 * 8, aStride = 8, + b = b, bByteOffset = 0, bByteStride = 3 * 2, + outStride = 3, + ) + } + + @Test + fun subnormal_and_extreme_weights_match_panama() { + // The branch-free special cases are the part most likely to diverge + // between the C and JVM formulations, and random weights never hit them. + val specials = floatArrayOf( + 0f, -0f, 5.9604645e-8f, -5.9604645e-8f, // zero, smallest subnormals + 6.0975552e-5f, -6.0975552e-5f, // largest subnormal + 6.1035156e-5f, -6.1035156e-5f, // smallest normal + 65504f, -65504f, 1f, -1f, 0.5f, -0.25f, + ) + val k = specials.size + val n = 4 + val bFloats = FloatArray(k * n) { specials[it / n] } + val b = fp16Bytes(bFloats) + val a = FloatArray(2 * k) { if (it % 3 == 0) 1f else 0.5f } + assertParity( + m = 2, n = n, k = k, + a = a, aOffset = 0, aStride = k, + b = b, bByteOffset = 0, bByteStride = n * 2, + outStride = n, + ) + } + + @Test + fun llm_typical_256_squared_matches_panama() { + val rng = Random(99) + val m = 256; val n = 256; val k = 256 + val a = FloatArray(m * k) { rng.nextFloat() - 0.5f } + val bFloats = FloatArray(k * n) { rng.nextFloat() - 0.5f } + val b = fp16Bytes(bFloats) + assertParity( + m = m, n = n, k = k, + a = a, aOffset = 0, aStride = k, + b = b, bByteOffset = 0, bByteStride = n * 2, + outStride = n, + tolScale = 5e-5f, + ) + } + + @Test + fun zero_m_or_n_no_op() { + val out = FloatArray(5) { 7f } + NativeFp16MatmulKernel.matmul( + FloatArray(0), 0, 0, + ByteArray(0), 0, 0, + out, 0, 0, + m = 0, n = 5, k = 0, + ) + for (v in out) assertEquals(7f, v, "out should be unchanged when m == 0") + } + + @Test + fun zero_k_zeros_output() { + val out = FloatArray(2 * 3) { 9f } + NativeFp16MatmulKernel.matmul( + FloatArray(0), 0, 0, + ByteArray(0), 0, 0, + out, 0, 3, + m = 2, n = 3, k = 0, + ) + for (v in out) assertEquals(0f, v, "out block should be zeroed when k == 0") + } + + @Test + fun rejects_negative_dimensions() { + assertFailsWith { + NativeFp16MatmulKernel.matmul( + FloatArray(0), 0, 0, + ByteArray(0), 0, 0, + FloatArray(0), 0, 0, + m = -1, n = 1, k = 1, + ) + } + } + + @Test + fun provider_returns_native_fp16_when_available() { + // The regression this whole change exists to prevent: the provider + // carried matmulBf16 but not matmulFp16, so FP16 silently cascaded to + // the JVM kernel and looked like a slow kernel rather than a missing one. + val kernel = NativeKernelProvider.matmulFp16() + assertTrue( + kernel === NativeFp16MatmulKernel, + "Provider must hand out the native FP16 kernel when bundled lib is loaded", + ) + } +} 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 index 512152ec..d32f5791 100644 --- 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 @@ -36,7 +36,11 @@ public interface NarrowFloatCodec { */ public fun encode(value: Float): Int - /** Decode a 16-bit pattern (low 16 bits used) back to FP32. Always exact — f32 is a superset. */ + /** + * Decode a 16-bit pattern (low 16 bits used) back to FP32. Exact for every finite value and + * for the infinities — f32 is a superset of both formats. NaN handling is per-implementation: + * [Bf16Codec] reproduces the pattern verbatim, [Fp16Codec] forces it quiet. + */ public fun decode(bits: Int): Float } @@ -123,30 +127,48 @@ public object Fp16Codec : NarrowFloatCodec { return sign or (q + if (roundUp) 1 else 0) } - /** FP16 → FP32 — always exact, including subnormals (which renormalize into FP32 normals). */ + /** + * FP16 → FP32 — always exact, including subnormals (which renormalize into FP32 normals). + * + * Straight-line: the magnitude is computed by a three-way select on the exponent field with no + * data-dependent loop. This runs once per weight element inside `ScalarFp16MatmulKernel`, so a + * renormalization loop in the subnormal arm would sit on the hot path of every FP16 matmul on + * targets without an intrinsic (see #887; the JVM kernel calls `Float.float16ToFloat` instead). + * + * Subnormals need no loop because binary16 defines them as `mant * 2⁻²⁴` for `mant` in + * `[1, 1023]`. Both factors are exact in FP32 and the product is a normal FP32 (the smallest, + * `2⁻²⁴`, is far above FP32's `2⁻¹²⁶` normal floor), so the multiply is exact and the hardware + * does the renormalization. `mant == 0` falls out of the same expression as ±0. + * + * A NaN is **quieted**, keeping its payload — matching [encode], which never emits a signaling + * binary16 NaN either, and matching the hardware conversion (`vcvtph2ps` quiets, so the JVM + * intrinsic does too). Preserving the signaling bit instead would buy nothing, since the first + * arithmetic use quiets it anyway, and would make the JVM and non-JVM kernels disagree on those + * 1022 patterns. Infinities are untouched — quieting is applied only when the mantissa is + * non-zero, or ±Inf would decode as NaN. + * + * With that, this is bit-for-bit identical to the JDK's `Float.float16ToFloat` across all 65536 + * inputs; `Fp16CodecIntrinsicParityTest` asserts that exhaustively. + */ 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)) + val magnitude = when { + exp == 0 -> (mant * TWO_POW_MINUS_24).toRawBits() // ±0 and subnormals + exp != 0x1F -> ((exp - 15 + 127) shl 23) or (mant shl 13) // normals + mant == 0 -> 0x7F80_0000 // ±Inf + else -> 0x7FC0_0000 or (mant shl 13) // NaN, forced quiet } + return Float.fromBits(sign or magnitude) } + + /** + * 2⁻²⁴ — the value of binary16's smallest subnormal, and the scale that turns a subnormal + * mantissa into its FP32 value. Written as a bit pattern rather than a decimal literal so the + * constant is exactly the power of two the multiply relies on. + */ + private val TWO_POW_MINUS_24: Float = Float.fromBits(0x3380_0000) } 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 index 865fddb2..bf38253d 100644 --- 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 @@ -238,4 +238,112 @@ class NarrowFloatCodecTest { assertEquals(Fp16Codec.decode(0x3C00), Fp16Codec.decode(0x3C00 or 0xFFFF_0000.toInt())) assertEquals(Bf16Codec.decode(0x3F80), Bf16Codec.decode(0x3F80 or 0xFFFF_0000.toInt())) } + + // ------------------------------------------------- FP16 decode: loop-free equivalence (#887) + + /** + * The pre-#887 [Fp16Codec.decode], kept verbatim as the oracle. + * + * Its subnormal arm renormalizes with a data-dependent `do/while`, which is what made the FP16 + * matmul kernels slow. The replacement computes the same value straight-line; this reference + * exists so that equivalence is asserted against the actual old code rather than re-derived, + * and it runs on every target, unlike the JDK-intrinsic comparison in + * `Fp16CodecIntrinsicParityTest`. + */ + private fun decodeByRenormalizationLoop(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) + } else { + 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)) + else -> Float.fromBits(sign or ((exp - 15 + 127) shl 23) or (mant shl 13)) + } + } + + @Test + fun fp16_loop_free_decode_matches_the_renormalization_loop_on_every_non_nan_input() { + // 63490 patterns — every normal, subnormal, zero and infinity — compared as raw bits so + // the sign of zero counts. NaN is deliberately excluded: it is the one place #887 changed + // the result, pinned separately below. + var checked = 0 + for (bits in 0..0xFFFF) { + if (isNaNPattern(bits)) continue + assertEquals( + decodeByRenormalizationLoop(bits).toRawBits(), + Fp16Codec.decode(bits).toRawBits(), + "decode diverged at 0x${bits.toString(16)}", + ) + checked++ + } + assertEquals(0x1_0000 - 2 * 1023, checked, "sweep must cover every non-NaN pattern") + } + + @Test + fun fp16_decode_quiets_signaling_nans_and_keeps_the_payload() { + // Quieting is the deliberate divergence from the old decode (#887): it makes the codec + // agree with the hardware/JDK conversion the JVM kernel now uses, so no target disagrees. + // The payload below the quiet bit must still survive, or a NaN could decode as Inf. + // + // How many patterns this *changed* relative to the old decode is deliberately not asserted + // here. Kotlin/JS quiets a signaling NaN itself when a Float crosses float32/double, so on + // that target the old implementation is observationally identical to this one and the + // count is 0 rather than 1022. The count is pinned in `Fp16CodecIntrinsicParityTest`, + // where no platform sits in between. Everything below is portable and holds everywhere. + for (bits in 0..0xFFFF) { + if (!isNaNPattern(bits)) continue + val decoded = Fp16Codec.decode(bits) + assertTrue(decoded.isNaN(), "0x${bits.toString(16)} must stay NaN") + + val raw = decoded.toRawBits() + assertTrue(raw and 0x0040_0000 != 0, "0x${bits.toString(16)} must decode quiet") + // The binary16 mantissa lands at bit 13; the quiet bit (0x0040_0000) is bit 9 of that + // mantissa, so forcing it quiet is exactly an OR onto the shifted payload. + assertEquals( + ((bits and 0x03FF) shl 13) or 0x0040_0000, raw and 0x007F_FFFF, + "payload must be preserved for 0x${bits.toString(16)}", + ) + assertEquals( + (bits and 0x8000) shl 16, raw and 0x8000_0000.toInt(), + "sign must be preserved for 0x${bits.toString(16)}", + ) + } + } + + @Test + fun fp16_infinities_are_not_caught_by_the_nan_quieting() { + // The quiet bit is applied only when the mantissa is non-zero; otherwise Inf becomes NaN. + assertEquals(Float.POSITIVE_INFINITY, Fp16Codec.decode(0x7C00)) + assertEquals(Float.NEGATIVE_INFINITY, Fp16Codec.decode(0xFC00)) + } + + private fun isNaNPattern(bits: Int): Boolean = + ((bits ushr 10) and 0x1F) == 0x1F && (bits and 0x03FF) != 0 + + @Test + fun fp16_decode_preserves_the_sign_of_zero_and_of_subnormals() { + // The subnormal arm now multiplies rather than assembling bits, so pin the cases where a + // multiply could plausibly lose the sign. + assertEquals(0x8000_0000.toInt(), Fp16Codec.decode(0x8000).toRawBits(), "-0") + assertEquals(0x0000_0000, Fp16Codec.decode(0x0000).toRawBits(), "+0") + assertTrue(Fp16Codec.decode(0x8001) < 0f, "smallest negative subnormal must stay negative") + assertEquals( + -Fp16Codec.decode(0x03FF), Fp16Codec.decode(0x83FF), + "negative subnormals must mirror positive ones exactly", + ) + } } diff --git a/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/types/Fp16CodecIntrinsicParityTest.kt b/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/types/Fp16CodecIntrinsicParityTest.kt new file mode 100644 index 00000000..483b0379 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/jvmTest/kotlin/sk/ainet/lang/types/Fp16CodecIntrinsicParityTest.kt @@ -0,0 +1,136 @@ +package sk.ainet.lang.types + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * [Fp16Codec.decode] must agree bit-for-bit with the JDK's `Float.float16ToFloat` on every input. + * + * This is the licence for `PanamaVectorFp16MatmulKernel` to call the JDK conversion instead of the + * codec (#887): the intrinsic lowers to `vcvtph2ps` where the codec was a per-element function + * call, and the swap is only safe because the two are the same function. The codec stays the + * reference for targets without an intrinsic, so this test also guards the codec against drift in + * either direction. + * + * JVM-only by necessity — `Float.float16ToFloat` is JDK 20+ and has no `commonMain` equivalent, + * which is why the codec exists in the first place. The portable equivalence check lives in + * `NarrowFloatCodecTest`. + */ +class Fp16CodecIntrinsicParityTest { + + private fun intrinsic(bits: Int): Float = java.lang.Float.float16ToFloat(bits.toShort()) + + @Test + fun decode_matches_the_jdk_intrinsic_on_all_65536_inputs() { + for (bits in 0..0xFFFF) { + assertEquals( + intrinsic(bits).toRawBits(), + Fp16Codec.decode(bits).toRawBits(), + "decode diverged from Float.float16ToFloat at 0x${bits.toString(16)}", + ) + } + } + + @Test + fun decode_matches_the_intrinsic_across_every_exponent_class() { + // The sweep above would still pass if one whole class were somehow skipped, so assert the + // classes are actually populated: normals, subnormals, zeros, infinities and NaNs. + var normals = 0 + var subnormals = 0 + var zeros = 0 + var infinities = 0 + var nans = 0 + for (bits in 0..0xFFFF) { + val exp = (bits ushr 10) and 0x1F + val mant = bits and 0x03FF + when { + exp == 0x1F && mant == 0 -> infinities++ + exp == 0x1F -> nans++ + exp == 0 && mant == 0 -> zeros++ + exp == 0 -> subnormals++ + else -> normals++ + } + } + assertEquals(2, zeros) + assertEquals(2, infinities) + assertEquals(2 * 1023, subnormals) + assertEquals(2 * 1023, nans) + assertEquals(2 * 30 * 1024, normals) + } + + /** + * The pre-#887 [Fp16Codec.decode], kept verbatim, so the scope of the behaviour change can be + * asserted rather than described. + * + * This lives in the JVM test rather than beside the portable one in `NarrowFloatCodecTest` + * because it is the only target that can observe the difference: Kotlin/JS quiets a signaling + * NaN itself whenever a Float crosses float32/double, so there the old implementation and the + * new one produce identical bits and the count below would be 0. + */ + private fun decodeByRenormalizationLoop(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) + } else { + 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)) + else -> Float.fromBits(sign or ((exp - 15 + 127) shl 23) or (mant shl 13)) + } + } + + @Test + fun exactly_the_signaling_nans_changed_relative_to_the_old_decode() { + // 1022 patterns — mantissas 1..0x1FF, both signs — and nothing else. A change anywhere + // outside that set would be a regression, not the intended quieting. + var changed = 0 + for (bits in 0..0xFFFF) { + if (decodeByRenormalizationLoop(bits).toRawBits() == Fp16Codec.decode(bits).toRawBits()) { + continue + } + changed++ + val exp = (bits ushr 10) and 0x1F + val mant = bits and 0x03FF + assertTrue( + exp == 0x1F && mant != 0 && mant < 0x0200, + "0x${bits.toString(16)} changed but is not a signaling NaN", + ) + } + assertEquals(2 * 511, changed, "only the signaling NaNs may differ from the old decode") + } + + @Test + fun the_intrinsic_round_trips_the_codecs_own_encode() { + // Cross-check the other direction too: whatever encode produces must decode identically + // under both implementations, so an encode change cannot silently split them. + val samples = floatArrayOf( + 0.0f, -0.0f, 1.0f, -1.0f, 0.5f, 65504f, -65504f, 6.1e-5f, -6.1e-5f, + 5.96e-8f, 1e-8f, 3.14159265f, -2.71828f, 1234.5f, + Float.POSITIVE_INFINITY, Float.NEGATIVE_INFINITY, + ) + for (v in samples) { + val bits = Fp16Codec.encode(v) + assertEquals( + intrinsic(bits).toRawBits(), + Fp16Codec.decode(bits).toRawBits(), + "round-trip of $v (bits=0x${bits.toString(16)})", + ) + } + assertTrue(Fp16Codec.decode(Fp16Codec.encode(Float.NaN)).isNaN()) + assertTrue(intrinsic(Fp16Codec.encode(Float.NaN)).isNaN()) + } +}