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 f0d3aec8..dde63a47 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 @@ -13,6 +13,7 @@ import sk.ainet.backend.api.kernel.KernelRegistry import sk.ainet.lang.tensor.data.RowDequantSource import sk.ainet.lang.tensor.data.FloatArrayTensorData import sk.ainet.lang.tensor.data.IntArrayTensorData +import sk.ainet.lang.tensor.data.NarrowFloatInputMajorTensorData import sk.ainet.lang.tensor.data.Q4_0TensorData import sk.ainet.lang.tensor.data.Q8_0TensorData import sk.ainet.lang.tensor.data.Q4_KTensorData @@ -624,6 +625,17 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory // packed type that can be a matmul weight. See transformers #178. is Q8_0TensorData -> return newTensor(Q8_0BlockTensorData(Shape(cols, rows), d.packedData) as TensorData, tensor.dtype, tensor) is Q4_0TensorData -> return newTensor(Q4_0BlockTensorData(Shape(cols, rows), d.packedData) as TensorData, tensor.dtype, tensor) + // Narrow floats (FP16/BF16) relaid input-major at load: the transpose is the + // same buffer read with the other shape's strides, so hand back an ordinary + // dense narrow tensor over it. This is what lets a KEEP_NATIVE weight survive + // `Linear.onForward`'s `weight.t()` and reach the narrow matmul kernel — see + // issue #888. Note the asymmetry with the block-quant arms above: only the + // *input-major* type is safe to reinterpret. A row-major narrow buffer falls + // through to the generic path on purpose, because swapping its shape would + // silently yield a different matrix rather than the transpose. + // Lives only here: DefaultCpuOpsJvm.transpose intercepts nothing that would + // shadow this case, so the JVM falls through to this arm too. + is NarrowFloatInputMajorTensorData -> return newTensor(d.transposedView() as TensorData, tensor.dtype, tensor) else -> {} } } diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/NarrowFloatTransposeDispatchTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/NarrowFloatTransposeDispatchTest.kt new file mode 100644 index 00000000..6676d319 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/NarrowFloatTransposeDispatchTest.kt @@ -0,0 +1,175 @@ +package sk.ainet.exec.tensor.ops + +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertSame +import kotlin.test.assertTrue +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.NarrowFloatDenseTensorData +import sk.ainet.lang.tensor.data.NarrowFloatInputMajorTensorData +import sk.ainet.lang.tensor.data.NarrowFloatTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.matmul +import sk.ainet.lang.tensor.t +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Fp16Codec +import sk.ainet.lang.types.NarrowFloatCodec + +/** + * Proves a KEEP_NATIVE weight survives `ctx.ops.matmul(x, ops.transpose(W))` still packed, so it + * reaches the narrow matmul kernel instead of being widened elementwise (issue #888). + * + * Sibling of [PackedMatmulDispatchTest], and deliberately in `commonTest` for the same reason: + * `DefaultCpuOps` and `DefaultCpuOpsJvm` intercept transpose separately, so only running on both + * jvmTest and linuxX64Test proves both arms are wired. + * + * Weights are stored `[out, in]` as they are on disk, then transposed — the exact shape of what + * `Linear.onForward` does on every forward pass. + */ +class NarrowFloatTransposeDispatchTest { + + private val ctx = DirectCpuExecutionContext() + + private val outDim = 3 + private val inDim = 4 + + /** `[out, in]`, all exactly representable in binary16 and bfloat16 alike. */ + private val weightRowMajor = floatArrayOf( + 1.0f, 2.0f, -1.0f, 0.5f, + 4.0f, -0.5f, 2.0f, 1.0f, + -2.0f, 1.0f, 0.25f, 8.0f, + ) + + private val input = floatArrayOf(1.0f, -1.0f, 2.0f, 0.5f) + + private fun pack(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 + } + + /** `y[j] = sum_i x[i] * W[j][i]` — what `x.matmul(W.t())` must produce. */ + private fun reference(): FloatArray = FloatArray(outDim) { j -> + var acc = 0.0f + for (i in 0 until inDim) acc += input[i] * weightRowMajor[j * inDim + i] + acc + } + + private fun inputTensor(): Tensor = + ctx.fromFloatArray(Shape(1, inDim), FP32::class, input) + + @Suppress("UNCHECKED_CAST") + private fun inputMajorWeight(codec: NarrowFloatCodec): Tensor { + val data = NarrowFloatInputMajorTensorData.fromRowMajor( + Shape(outDim, inDim), pack(weightRowMajor, codec), codec, + ) + return ctx.fromData(data as TensorData, FP32::class) + } + + private fun assertClose(expected: FloatArray, actual: FloatArray, label: String) { + assertEquals(expected.size, actual.size, "$label: length") + for (i in expected.indices) { + assertTrue( + abs(expected[i] - actual[i]) < 1e-4f, + "$label: element $i expected ${expected[i]} but was ${actual[i]}", + ) + } + } + + @Test + fun `transpose of an input-major weight stays narrow and shares the buffer`() { + for (codec in listOf(Fp16Codec, Bf16Codec)) { + val w = inputMajorWeight(codec) + val original = (w.data as NarrowFloatTensorData).packedData + + val wt = w.t() + + assertTrue( + wt.data is NarrowFloatTensorData, + "${codec.dtype.name}: transpose widened the weight — the kernel is unreachable", + ) + assertTrue( + wt.data is NarrowFloatDenseTensorData, + "${codec.dtype.name}: expected a plain dense narrow tensor after transpose", + ) + assertEquals(Shape(inDim, outDim), wt.shape, "${codec.dtype.name}: shape must swap") + assertEquals( + codec, (wt.data as NarrowFloatTensorData).codec, + "${codec.dtype.name}: codec must survive the transpose", + ) + assertSame( + original, (wt.data as NarrowFloatTensorData).packedData, + "${codec.dtype.name}: transpose must not copy — copying per forward is the bug", + ) + } + } + + @Test + fun `matmul through the transpose matches the fp32 reference`() { + val expected = reference() + for (codec in listOf(Fp16Codec, Bf16Codec)) { + val y = inputTensor().matmul(inputMajorWeight(codec).t()) + assertEquals(Shape(1, outDim), y.shape, "${codec.dtype.name}: output shape") + assertClose(expected, y.data.copyToFloatArray(), codec.dtype.name) + } + } + + @Test + fun `a row-major narrow weight is left alone by the lazy transpose`() { + // The safety property. Only the input-major type may be reinterpreted; swapping the shape + // of a row-major narrow buffer would silently produce a different matrix. The generic path + // is slow, but it is correct — and correctness is what must not regress here. + @Suppress("UNCHECKED_CAST") + val rowMajor = ctx.fromData( + NarrowFloatDenseTensorData( + Shape(outDim, inDim), pack(weightRowMajor, Fp16Codec), Fp16Codec, + ) as TensorData, + FP32::class, + ) + + val wt = rowMajor.t() + assertEquals(Shape(inDim, outDim), wt.shape) + + // Whatever representation the generic path chose, the values must be the true transpose. + val got = wt.data.copyToFloatArray() + for (j in 0 until outDim) { + for (i in 0 until inDim) { + assertEquals( + weightRowMajor[j * inDim + i], got[i * outDim + j], + "row-major transpose wrong at [$i, $j]", + ) + } + } + assertClose(reference(), inputTensor().matmul(wt).data.copyToFloatArray(), "row-major") + } + + @Test + fun `the two codecs disagree on identical bytes`() { + // Vacuity guard: both formats are 2 bytes per element, so nothing above would catch a + // dispatch that picked the kernel by byte width instead of by codec. + val bytes = pack(weightRowMajor, Fp16Codec) + + @Suppress("UNCHECKED_CAST") + fun matmulAs(codec: NarrowFloatCodec): FloatArray { + val data = NarrowFloatInputMajorTensorData.fromRowMajor(Shape(outDim, inDim), bytes, codec) + val w = ctx.fromData(data as TensorData, FP32::class) + return inputTensor().matmul(w.t()).data.copyToFloatArray() + } + + val asFp16 = matmulAs(Fp16Codec) + val asBf16 = matmulAs(Bf16Codec) + assertTrue( + asFp16.indices.any { abs(asFp16[it] - asBf16[it]) > 1e-3f }, + "reading the same bytes under both codecs produced the same result — " + + "the codec assertions in this class would prove nothing", + ) + } +} 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 587ccc5d..4364b89b 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 @@ -239,6 +239,8 @@ internal class DefaultCpuOpsJvm( @Suppress("UNCHECKED_CAST") return newTensor(transposed as TensorData, tensor.dtype, tensor) } + // Narrow-float input-major lazy transpose is handled in DefaultCpuOpsBase too — + // nothing above intercepts it, so it falls through. Issue #888. // Q6_K / Q5_1 / Q5_0 lazy transpose is handled in DefaultCpuOpsBase // (block-major, shared with Native); the JVM ops don't intercept them here. // MemorySegment FP32 fast path: physical transpose via SIMD. 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 index 55d9febc..62395e4d 100644 --- 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 @@ -35,6 +35,15 @@ public enum class NarrowFloatLoadPolicy { * * **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. + * + * **Layout caveat (issue #888).** This loader emits row-major bytes, matching the file. But + * projections are stored `[out, in]` while the narrow matmul dispatch needs `[in, out]`, so a + * consumer that goes through `Linear` — which transposes the weight on every forward — will + * see that transpose fall to the generic elementwise path and widen the tensor anyway, at a + * cost far exceeding what KEEP_NATIVE saves. Consumers running such a hot path should relay + * the bytes once with `NarrowFloatInputMajorTensorData.fromRowMajor`, which makes the + * transpose a zero-copy view. Tensors that are gathered rather than multiplied — embedding + * tables above all — should stay row-major, since input-major storage strides their row reads. */ KEEP_NATIVE, } 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 84f41025..d31b2273 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 @@ -3095,6 +3095,24 @@ public final class sk/ainet/lang/tensor/data/NarrowFloatDenseTensorData$Companio public final fun fromFloatArray (Lsk/ainet/lang/tensor/Shape;[FLsk/ainet/lang/types/NarrowFloatCodec;)Lsk/ainet/lang/tensor/data/NarrowFloatDenseTensorData; } +public final class sk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorData : sk/ainet/lang/tensor/data/NarrowFloatTensorData { + public static final field Companion Lsk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorData$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 fun transposedView ()Lsk/ainet/lang/tensor/data/NarrowFloatDenseTensorData; +} + +public final class sk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorData$Companion { + public final fun fromRowMajor (Lsk/ainet/lang/tensor/Shape;[BLsk/ainet/lang/types/NarrowFloatCodec;)Lsk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorData; +} + 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; 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 index 5f95a53b..1f8f916d 100644 --- 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 @@ -161,5 +161,136 @@ public class Fp16DenseTensorData( } } +/** + * A rank-2 narrow-float weight whose bytes are stored **input-major**: the element at logical + * `[row, col]` sits at flat byte index `(col * rows + row) * 2`, the transpose of the usual + * row-major order. + * + * ### Why this exists + * + * Projections are stored `[out, in]` on disk, but `chooseQuantizedMatmul` needs `[in, out]`, so + * every `Linear.onForward` calls `weight.t()` first. Transposing a row-major narrow tensor has no + * fast path — it walks the tensor elementwise through boxed `get()` and widens to FP32, which at + * real projection sizes costs hundreds of milliseconds to seconds *per weight, per token*. That + * made KEEP_NATIVE slower than not using it at all. + * + * Storing the bytes input-major once, at load, makes that transpose free: input-major storage of + * `[out, in]` **is** row-major storage of `[in, out]`, so [transposedView] hands back an ordinary + * [NarrowFloatDenseTensorData] over the very same buffer. No copy, and — unlike the lazy transpose + * used for the K-quants — element access stays correct on both sides, because each type indexes + * the shared bytes with the strides its own shape implies. + * + * The layout is carried in the type rather than a flag so that a bare shape swap cannot be applied + * to a row-major buffer by accident. Doing that would not throw; it would silently reinterpret the + * weight as a different matrix. + * + * Build these with [fromRowMajor], which performs the one-off relayout. + */ +public class NarrowFloatInputMajorTensorData( + initialShape: Shape, + private val data: ByteArray, + override val codec: NarrowFloatCodec, +) : NarrowFloatTensorData { + + override val shape: Shape = Shape(initialShape.dimensions.copyOf()) + override val packedData: ByteArray get() = data + + private val rows: Int = shape.dimensions[0] + private val cols: Int = shape.dimensions[1] + + init { + require(shape.dimensions.size == 2) { + "Input-major layout is only defined for rank-2 weights, got shape $shape" + } + 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" + } + } + + /** Flat index of logical `[row, col]` under input-major storage. */ + private fun flatIndex(indices: IntArray): Int { + require(indices.size == 2) { + "Number of indices (${indices.size}) must match tensor dimensions (2)" + } + val row = indices[0] + val col = indices[1] + require(row in 0 until rows) { "Index $row out of bounds for dimension 0 with size $rows" } + require(col in 0 until cols) { "Index $col out of bounds for dimension 1 with size $cols" } + return col * rows + row + } + + override fun get(vararg indices: Int): Float { + val byteIdx = flatIndex(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 = flatIndex(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() + } + + /** Decodes in logical row-major order, so callers see the same values a dense tensor would. */ + override fun copyToFloatArray(): FloatArray { + val out = FloatArray(shape.volume) + var dst = 0 + for (row in 0 until rows) { + for (col in 0 until cols) { + val byteIdx = (col * rows + row) * NarrowFloatTensorData.BYTES_PER_ELEMENT + val lo = data[byteIdx].toInt() and 0xFF + val hi = data[byteIdx + 1].toInt() and 0xFF + out[dst++] = codec.decode((hi shl 8) or lo) + } + } + return out + } + + /** + * The transpose, sharing this instance's buffer — no copy. Input-major `[rows, cols]` is + * row-major `[cols, rows]`, so the result is an ordinary dense narrow tensor that + * `chooseQuantizedMatmul` accepts directly. + */ + public fun transposedView(): NarrowFloatDenseTensorData = + NarrowFloatDenseTensorData(Shape(cols, rows), data, codec) + + public companion object { + /** + * Relayout `rowMajorBytes` (logical `[rows, cols]`, row-major) into input-major order. + * This is the one-off cost that buys a free transpose on every later forward pass. + */ + public fun fromRowMajor( + shape: Shape, + rowMajorBytes: ByteArray, + codec: NarrowFloatCodec, + ): NarrowFloatInputMajorTensorData { + require(shape.dimensions.size == 2) { + "Input-major layout is only defined for rank-2 weights, got shape $shape" + } + val rows = shape.dimensions[0] + val cols = shape.dimensions[1] + val required = shape.volume * NarrowFloatTensorData.BYTES_PER_ELEMENT + require(rowMajorBytes.size >= required) { + "Data size ${rowMajorBytes.size} is less than required $required bytes " + + "for ${shape.volume} ${codec.dtype.name} elements" + } + val out = ByteArray(required) + for (row in 0 until rows) { + for (col in 0 until cols) { + val src = (row * cols + col) * NarrowFloatTensorData.BYTES_PER_ELEMENT + val dst = (col * rows + row) * NarrowFloatTensorData.BYTES_PER_ELEMENT + out[dst] = rowMajorBytes[src] + out[dst + 1] = rowMajorBytes[src + 1] + } + } + return NarrowFloatInputMajorTensorData(shape, out, codec) + } + } +} + /** Decode a [NarrowFloatTensorData] to a fresh FloatArray. */ public fun NarrowFloatTensorData.toFloatArray(): FloatArray = copyToFloatArray() diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorDataTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorDataTest.kt new file mode 100644 index 00000000..08da6522 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/data/NarrowFloatInputMajorTensorDataTest.kt @@ -0,0 +1,168 @@ +package sk.ainet.lang.tensor.data + +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.Bf16Codec +import sk.ainet.lang.types.Fp16Codec +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** + * Pins [NarrowFloatInputMajorTensorData] — the relaid layout that makes transposing a KEEP_NATIVE + * weight free (issue #888). + * + * The invariant under test throughout: input-major storage of `[rows, cols]` **is** row-major + * storage of `[cols, rows]`, so [NarrowFloatInputMajorTensorData.transposedView] can share the + * buffer outright. Element access has to stay correct on both sides of that view, which is what + * distinguishes this from the K-quant lazy transpose (where `get()` on a transposed tensor is + * meaningless and only the kernel's direct `packedData` read is valid). + */ +class NarrowFloatInputMajorTensorDataTest { + + /** 2x3, all exactly representable in both narrow formats so comparisons can be exact. */ + private val rows = 2 + private val cols = 3 + private val values = floatArrayOf( + 1.0f, 2.0f, 4.0f, + 8.0f, 0.5f, -2.0f, + ) + + private fun rowMajorBytes(codec: sk.ainet.lang.types.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 + } + + @Test + fun `element access matches the row-major original`() { + val data = NarrowFloatInputMajorTensorData.fromRowMajor( + Shape(rows, cols), rowMajorBytes(Fp16Codec), Fp16Codec, + ) + for (r in 0 until rows) { + for (c in 0 until cols) { + assertEquals( + values[r * cols + c], data.get(r, c), + "get($r, $c) must see the logical element, not the relaid byte order", + ) + } + } + } + + @Test + fun `copyToFloatArray decodes in logical row-major order`() { + val data = NarrowFloatInputMajorTensorData.fromRowMajor( + Shape(rows, cols), rowMajorBytes(Bf16Codec), Bf16Codec, + ) + assertContentEquals( + values, data.copyToFloatArray(), + "consumers that do not care about storage must see the ordinary row-major sequence", + ) + } + + @Test + fun `transposedView shares the buffer and reads as the transpose`() { + val data = NarrowFloatInputMajorTensorData.fromRowMajor( + Shape(rows, cols), rowMajorBytes(Fp16Codec), Fp16Codec, + ) + val view = data.transposedView() + + assertSame( + data.packedData, view.packedData, + "the whole point is that no copy happens — a copy per forward is the bug being fixed", + ) + assertEquals(Shape(cols, rows), view.shape) + assertEquals(Fp16Codec, view.codec, "codec must survive; FP16 and BF16 are both 2 bytes") + + for (r in 0 until rows) { + for (c in 0 until cols) { + assertEquals( + data.get(r, c), view.get(c, r), + "view[$c, $r] must be the transpose of data[$r, $c]", + ) + } + } + } + + @Test + fun `relayout round-trips back to the original bytes`() { + // Relaying twice (with the shape flipped in between) is the identity — the cheapest + // statement of "this is a transpose and not some other permutation". + val original = rowMajorBytes(Fp16Codec) + val once = NarrowFloatInputMajorTensorData.fromRowMajor(Shape(rows, cols), original, Fp16Codec) + val twice = NarrowFloatInputMajorTensorData.fromRowMajor( + Shape(cols, rows), once.packedData, Fp16Codec, + ) + assertContentEquals(original, twice.packedData) + } + + @Test + fun `a square weight is still genuinely transposed`() { + // Square shapes are where an off-by-one in the relayout hides: the byte count matches + // either way, so only the values reveal a wrong permutation. + val square = floatArrayOf(1.0f, 2.0f, 4.0f, 8.0f) + val bytes = ByteArray(square.size * 2) + for (i in square.indices) { + val bits = Fp16Codec.encode(square[i]) + bytes[i * 2] = (bits and 0xFF).toByte() + bytes[i * 2 + 1] = ((bits ushr 8) and 0xFF).toByte() + } + val view = NarrowFloatInputMajorTensorData + .fromRowMajor(Shape(2, 2), bytes, Fp16Codec) + .transposedView() + + assertContentEquals( + floatArrayOf(1.0f, 4.0f, 2.0f, 8.0f), view.copyToFloatArray(), + "transposedView of [[1,2],[4,8]] must be [[1,4],[2,8]]", + ) + } + + @Test + fun `set writes through to the logical element`() { + val data = NarrowFloatInputMajorTensorData.fromRowMajor( + Shape(rows, cols), rowMajorBytes(Fp16Codec), Fp16Codec, + ) + data.set(1, 2, value = 16.0f) + assertEquals(16.0f, data.get(1, 2)) + assertEquals(16.0f, data.transposedView().get(2, 1), "the shared buffer must see it too") + assertEquals(values[0], data.get(0, 0), "neighbouring elements must be untouched") + } + + @Test + fun `rank other than two is rejected`() { + // Norms are rank-1 and embeddings are gathered, not matmul'd; neither should ever be + // relaid. Failing loudly beats silently mis-indexing them. + assertFailsWith { + NarrowFloatInputMajorTensorData.fromRowMajor(Shape(4), ByteArray(8), Fp16Codec) + } + assertFailsWith { + NarrowFloatInputMajorTensorData(Shape(2, 2, 2), ByteArray(16), Fp16Codec) + } + } + + @Test + fun `a short buffer is rejected`() { + assertFailsWith { + NarrowFloatInputMajorTensorData.fromRowMajor(Shape(4, 4), ByteArray(8), Fp16Codec) + } + } + + @Test + fun `the two codecs disagree on identical bytes`() { + // Vacuity guard for every codec assertion above: both formats are 2 bytes per element, + // so a codec mix-up cannot be caught by shape or size checks alone. + val bytes = rowMajorBytes(Fp16Codec) + val asFp16 = NarrowFloatInputMajorTensorData.fromRowMajor(Shape(rows, cols), bytes, Fp16Codec) + val asBf16 = NarrowFloatInputMajorTensorData.fromRowMajor(Shape(rows, cols), bytes, Bf16Codec) + assertTrue( + !asFp16.copyToFloatArray().contentEquals(asBf16.copyToFloatArray()), + "if these agreed, the codec-preservation assertions would prove nothing", + ) + } +}