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 53d0fa9c0..9ec156e7c 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 @@ -780,11 +780,13 @@ internal class DefaultCpuOpsJvm( } override fun silu(tensor: Tensor): Tensor { - val data = tensor.data as? FloatArrayTensorData ?: return super.silu(tensor) - val buf = data.buffer - val out = FloatArray(buf.size) - for (i in buf.indices) { - val x = buf[i] + val win = floatWinOf(tensor.data) ?: return super.silu(tensor) + val buf = win.arr + val off = win.off + val n = tensor.shape.volume + val out = FloatArray(n) + for (i in 0 until n) { + val x = buf[off + i] out[i] = x / (1f + kotlin.math.exp(-x)) } return floatResult(Shape(tensor.shape.dimensions.copyOf()), tensor.dtype, out) @@ -822,8 +824,12 @@ internal class DefaultCpuOpsJvm( if (!supportsFloatOps(a) || !supportsFloatOps(b)) return null if (a.dtype != b.dtype) return null - val aData = a.data as? FloatArrayTensorData ?: return null - val bData = b.data as? FloatArrayTensorData ?: return null + val aWin = floatWinOf(a.data) ?: return null + val bWin = floatWinOf(b.data) ?: return null + val aBuf = aWin.arr + val aBase = aWin.off + val bBuf = bWin.arr + val bBase = bWin.off // Determine broadcasted output shape val outShape = try { broadcastShapes(a.shape, b.shape) } catch (e: IllegalArgumentException) { return null } @@ -835,41 +841,41 @@ internal class DefaultCpuOpsJvm( // Case 1: exact shape match (fast path) if (a.shape == b.shape) { - JvmVectorKernels.binaryFloat(aData.buffer, bData.buffer, outBuffer, outVolume, vectorOp, scalarOp) + JvmVectorKernels.binaryFloat(aBuf, bBuf, outBuffer, outVolume, vectorOp, scalarOp, aOffset = aBase, bOffset = bBase) return floatResult(Shape(outShape.dimensions.copyOf()), a.dtype, outBuffer) } // Case 2: scalar broadcast if (aVol == 1) { - val aval = aData.buffer[0] + val aval = aBuf[aBase] val speciesLen = floatSpecies.length() var idx = 0 val loopBound = floatSpecies.loopBound(outVolume) while (idx < loopBound) { val va = FloatVector.broadcast(floatSpecies, aval) - val vb = FloatVector.fromArray(floatSpecies, bData.buffer, idx) + val vb = FloatVector.fromArray(floatSpecies, bBuf, bBase + idx) vectorOp(va, vb).intoArray(outBuffer, idx) idx += speciesLen } while (idx < outVolume) { - outBuffer[idx] = scalarOp(aval, bData.buffer[idx]) + outBuffer[idx] = scalarOp(aval, bBuf[bBase + idx]) idx++ } return floatResult(Shape(outShape.dimensions.copyOf()), a.dtype, outBuffer) } if (bVol == 1) { - val bval = bData.buffer[0] + val bval = bBuf[bBase] val speciesLen = floatSpecies.length() var idx = 0 val loopBound = floatSpecies.loopBound(outVolume) while (idx < loopBound) { - val va = FloatVector.fromArray(floatSpecies, aData.buffer, idx) + val va = FloatVector.fromArray(floatSpecies, aBuf, aBase + idx) val vb = FloatVector.broadcast(floatSpecies, bval) vectorOp(va, vb).intoArray(outBuffer, idx) idx += speciesLen } while (idx < outVolume) { - outBuffer[idx] = scalarOp(aData.buffer[idx], bval) + outBuffer[idx] = scalarOp(aBuf[aBase + idx], bval) idx++ } return floatResult(Shape(outShape.dimensions.copyOf()), a.dtype, outBuffer) @@ -895,13 +901,13 @@ internal class DefaultCpuOpsJvm( val aOff = g * outLast var idx = 0 while (idx < loopBoundTail) { - val va = FloatVector.fromArray(floatSpecies, aData.buffer, aOff + idx) - val vb = FloatVector.fromArray(floatSpecies, bData.buffer, idx) + val va = FloatVector.fromArray(floatSpecies, aBuf, aBase + aOff + idx) + val vb = FloatVector.fromArray(floatSpecies, bBuf, bBase + idx) vectorOp(va, vb).intoArray(outBuffer, aOff + idx) idx += step } while (idx < outLast) { - outBuffer[aOff + idx] = scalarOp(aData.buffer[aOff + idx], bData.buffer[idx]) + outBuffer[aOff + idx] = scalarOp(aBuf[aBase + aOff + idx], bBuf[bBase + idx]) idx++ } } @@ -914,13 +920,13 @@ internal class DefaultCpuOpsJvm( val bOff = g * outLast var idx = 0 while (idx < loopBoundTail) { - val va = FloatVector.fromArray(floatSpecies, aData.buffer, idx) - val vb = FloatVector.fromArray(floatSpecies, bData.buffer, bOff + idx) + val va = FloatVector.fromArray(floatSpecies, aBuf, aBase + idx) + val vb = FloatVector.fromArray(floatSpecies, bBuf, bBase + bOff + idx) vectorOp(va, vb).intoArray(outBuffer, bOff + idx) idx += step } while (idx < outLast) { - outBuffer[bOff + idx] = scalarOp(aData.buffer[idx], bData.buffer[bOff + idx]) + outBuffer[bOff + idx] = scalarOp(aBuf[aBase + idx], bBuf[bBase + bOff + idx]) idx++ } } @@ -938,13 +944,27 @@ internal class DefaultCpuOpsJvm( scalarOp: (Float) -> Float ): Tensor? { if (!supportsFloatOps(tensor)) return null - val tensorData = tensor.data as? FloatArrayTensorData ?: return null + val win = floatWinOf(tensor.data) ?: return null val volume = tensor.shape.volume val outBuffer = FloatArray(volume) - JvmVectorKernels.unaryFloat(tensorData.buffer, outBuffer, volume, vectorOp, scalarOp) + JvmVectorKernels.unaryFloat(win.arr, outBuffer, volume, vectorOp, scalarOp, inputOffset = win.off) return floatResult(Shape(tensor.shape.dimensions.copyOf()), tensor.dtype, outBuffer) } + /** Dense-FP32 window (array + base offset): plain array data at 0, slab-backed data (#1173) at its arrayOffset. */ + private class FloatWin(@JvmField val arr: FloatArray, @JvmField val off: Int) + + @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) + private fun floatWinOf(d: sk.ainet.lang.tensor.data.TensorData<*, *>): FloatWin? = when (d) { + is FloatArrayTensorData<*> -> FloatWin(d.buffer, 0) + is sk.ainet.lang.tensor.data.StorageFloatTensorData<*> -> { + val st = d.storage + st.checkAlive() + FloatWin(st.floats!!, st.arrayOffset) + } + else -> null + } + private fun supportsFloatOps(a: Tensor, b: Tensor): Boolean { return supportsFloatOps(a) && a.dtype == b.dtype && @@ -1006,18 +1026,20 @@ internal class DefaultCpuOpsJvm( return CpuTensor(result as TensorData, this, a.dtype) } - // ---- FloatArray path ---- - val aData = a.data as? FloatArrayTensorData ?: return null - val bData = b.data as? FloatArrayTensorData ?: return null + // ---- FloatArray path (plain arrays at offset 0, slab windows at their base — #1173) ---- + val aWin = floatWinOf(a.data) ?: return null + val bWin = floatWinOf(b.data) ?: return null val work = m.toLong() * n.toLong() * k.toLong() val outBuffer = FloatArray(m * n) - // Try BLAS for large sizes if enabled and available - if (JvmCpuBackendConfig.blasEnabled && JvmBlas.isAvailable()) { + // Try BLAS for large sizes if enabled and available. The JNI shim takes whole arrays, so + // it only serves offset-0 operands; slab windows go through the kernel SPI below, which + // takes offsets natively. + if (aWin.off == 0 && bWin.off == 0 && JvmCpuBackendConfig.blasEnabled && JvmBlas.isAvailable()) { val blasThreshold = 512L * 512L * 256L // tuneable if (work >= blasThreshold) { - val ok = JvmBlas.sgemmRowMajorNN(m, n, k, 1f, aData.buffer, bData.buffer, outBuffer) + val ok = JvmBlas.sgemmRowMajorNN(m, n, k, 1f, aWin.arr, bWin.arr, outBuffer) if (ok) { return floatResult(Shape(m, n), a.dtype, outBuffer) } @@ -1029,8 +1051,8 @@ internal class DefaultCpuOpsJvm( // handles small + large inputs in one path, so the previous // simple-vs-blocked fork is no longer needed. fp32MatmulKernel.matmul( - aData.buffer, 0, k, - bData.buffer, 0, n, + aWin.arr, aWin.off, k, + bWin.arr, bWin.off, n, outBuffer, 0, n, m, n, k, ) @@ -1039,16 +1061,17 @@ internal class DefaultCpuOpsJvm( private fun vectorFloatReduceAllSum(tensor: Tensor): Tensor? { if (!supportsFloatOps(tensor)) return null - val data = tensor.data as? FloatArrayTensorData ?: return null - val buffer = data.buffer - val n = buffer.size + val win = floatWinOf(tensor.data) ?: return null + val buffer = win.arr + val base = win.off + val n = tensor.shape.volume if (n == 0) return null // NOTE: For numerical reproducibility with Kotlin's FloatArray.sum(), // perform strict left-to-right scalar accumulation. var acc = 0.0f var idx = 0 while (idx < n) { - acc += buffer[idx] + acc += buffer[base + idx] idx++ } val outData = DenseFloatArrayTensorData(Shape(), floatArrayOf(acc)) diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmVectorKernels.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmVectorKernels.kt index bde53ce38..d252c1fbb 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmVectorKernels.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmVectorKernels.kt @@ -21,18 +21,20 @@ internal object JvmVectorKernels { length: Int, op: (FloatVector, FloatVector) -> FloatVector, scalarOp: (Float, Float) -> Float, + aOffset: Int = 0, + bOffset: Int = 0, ) { var index = 0 val step = floatSpecies.length() val loopBound = floatSpecies.loopBound(length) while (index < loopBound) { - val va = FloatVector.fromArray(floatSpecies, a, index) - val vb = FloatVector.fromArray(floatSpecies, b, index) + val va = FloatVector.fromArray(floatSpecies, a, aOffset + index) + val vb = FloatVector.fromArray(floatSpecies, b, bOffset + index) op(va, vb).intoArray(out, index) index += step } while (index < length) { - out[index] = scalarOp(a[index], b[index]) + out[index] = scalarOp(a[aOffset + index], b[bOffset + index]) index++ } } @@ -43,17 +45,18 @@ internal object JvmVectorKernels { length: Int, op: (FloatVector) -> FloatVector, scalarOp: (Float) -> Float, + inputOffset: Int = 0, ) { var index = 0 val step = floatSpecies.length() val loopBound = floatSpecies.loopBound(length) while (index < loopBound) { - val v = FloatVector.fromArray(floatSpecies, input, index) + val v = FloatVector.fromArray(floatSpecies, input, inputOffset + index) op(v).intoArray(out, index) index += step } while (index < length) { - out[index] = scalarOp(input[index]) + out[index] = scalarOp(input[inputOffset + index]) index++ } } diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/SlabOperandVectorPathTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/SlabOperandVectorPathTest.kt new file mode 100644 index 000000000..96d59c9c8 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/tensor/ops/SlabOperandVectorPathTest.kt @@ -0,0 +1,118 @@ +package sk.ainet.exec.tensor.ops + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.forwardScope +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.StorageFloatTensorData +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertTrue + +/** + * #1173: the JVM Panama vector paths accept slab-backed operands at nonzero offsets instead of + * silently falling back to the scalar loops. Every vector-eligible shape is driven with both + * operands sliced from a `ForwardScope` slab (leading pad ⇒ nonzero `arrayOffset`) and compared + * bit-for-bit against the Ambient result — the tightest possible guard against off-by-offset + * reads, which produce plausible garbage rather than crashes. + */ +@OptIn(ExperimentalMemoryApi::class) +class SlabOperandVectorPathTest { + + private val n = 67 // deliberately not a multiple of any vector species length + private val aVals = FloatArray(n) { (it % 13 - 6) * 0.25f } + private val bVals = FloatArray(n) { (it % 7 - 3) * 0.5f } + private val biasVals = FloatArray(n) { (it % 5 - 2) * 0.125f } + + private fun ambient(op: (ctx: DirectCpuExecutionContext) -> Tensor): FloatArray = + op(DirectCpuExecutionContext()).data.copyToFloatArray() + + /** Runs [op] with all operands created inside a scope, after a pad allocation forcing nonzero offsets. */ + private fun slab(op: (ctx: sk.ainet.context.ExecutionContext) -> Tensor): FloatArray { + lateinit var result: FloatArray + DirectCpuExecutionContext().forwardScope(slabFloats = 4096) { scoped, _ -> + scoped.zeros(Shape(5), FP32::class) // pad: everything after this has offset > 0 + val out = op(scoped) + assertTrue( + (out.data is StorageFloatTensorData<*>), + "output should be slab-backed — otherwise this test is not testing the vector paths", + ) + result = out.data.copyToFloatArray() + } + return result + } + + private fun t(ctx: sk.ainet.context.ExecutionContext, shape: Shape, v: FloatArray): Tensor { + val slice = FloatArray(shape.volume) { v[it] } + return ctx.fromFloatArray(shape, FP32::class, slice) + } + + private fun check(name: String, op: (ctx: sk.ainet.context.ExecutionContext) -> Tensor) { + val expected = ambient { op(it) } + val actual = slab(op) + assertContentEquals(expected, actual, name) + } + + @Test + fun exactShapeBinary() = check("add exact") { ctx -> + val a = t(ctx, Shape(n), aVals) + val b = t(ctx, Shape(n), bVals) + a.ops.add(a, b) + } + + @Test + fun scalarBroadcastBothSides() { + check("a scalar") { ctx -> + val a = t(ctx, Shape(1), floatArrayOf(1.5f)) + val b = t(ctx, Shape(n), bVals) + a.ops.multiply(a, b) + } + check("b scalar") { ctx -> + val a = t(ctx, Shape(n), aVals) + val b = t(ctx, Shape(1), floatArrayOf(-0.75f)) + a.ops.subtract(a, b) + } + } + + @Test + fun biasBroadcastBothSides() { + check("bias on b") { ctx -> + val a = t(ctx, Shape(3, n), FloatArray(3 * n) { aVals[it % n] }) + val b = t(ctx, Shape(n), biasVals) + a.ops.add(a, b) + } + check("bias on a") { ctx -> + val a = t(ctx, Shape(n), biasVals) + val b = t(ctx, Shape(3, n), FloatArray(3 * n) { bVals[it % n] }) + a.ops.add(a, b) + } + } + + @Test + fun unaryActivations() { + check("relu") { ctx -> t(ctx, Shape(n), aVals).let { it.ops.relu(it) } } + check("silu") { ctx -> t(ctx, Shape(n), aVals).let { it.ops.silu(it) } } + } + + @Test + fun reduceAllMatchesAmbientBitForBit() { + // sum(null) returns a rank-0 scalar; both paths must accumulate in the same order. + val expected = ambient { ctx -> t(ctx, Shape(n), aVals).let { it.ops.sum(it, null) } } + lateinit var actual: FloatArray + DirectCpuExecutionContext().forwardScope(slabFloats = 4096) { scoped, _ -> + scoped.zeros(Shape(5), FP32::class) + val x = t(scoped, Shape(n), aVals) + actual = x.ops.sum(x, null).data.copyToFloatArray() + } + assertContentEquals(expected, actual, "reduce-all sum") + } + + @Test + fun matmul2dThroughTheKernelSpi() = check("matmul") { ctx -> + val a = t(ctx, Shape(7, 9), FloatArray(63) { aVals[it % n] }) + val b = t(ctx, Shape(9, 5), FloatArray(45) { bVals[it % n] }) + a.ops.matmul(a, b) + } +} 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 481b89b8a..e687e0d46 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 @@ -1327,14 +1327,21 @@ public final class sk/ainet/lang/memory/TernaryCodec { public static synthetic fun decode$default (Lsk/ainet/lang/memory/TernaryCodec;Lsk/ainet/lang/tensor/storage/TensorEncoding;[BIIILjava/lang/Object;)[F public final fun decodeBitNet ([BII)[F public static synthetic fun decodeBitNet$default (Lsk/ainet/lang/memory/TernaryCodec;[BIIILjava/lang/Object;)[F + public final fun decodeBitNetPlanes ([BIII)[F + public static synthetic fun decodeBitNetPlanes$default (Lsk/ainet/lang/memory/TernaryCodec;[BIIIILjava/lang/Object;)[F + public final fun decodeBitNetPlanesRow ([BIII[FII)V + public static synthetic fun decodeBitNetPlanesRow$default (Lsk/ainet/lang/memory/TernaryCodec;[BIII[FIIILjava/lang/Object;)V public final fun decodeTq1_0 ([BII)[F public static synthetic fun decodeTq1_0$default (Lsk/ainet/lang/memory/TernaryCodec;[BIIILjava/lang/Object;)[F public final fun decodeTq2_0 ([BII)[F public static synthetic fun decodeTq2_0$default (Lsk/ainet/lang/memory/TernaryCodec;[BIIILjava/lang/Object;)[F public final fun encode (Lsk/ainet/lang/tensor/storage/TensorEncoding;[F)[B public final fun encodeBitNet ([F)[B + public final fun encodeBitNetPlanes ([FII)[B public final fun encodeTq1_0 ([F)[B public final fun encodeTq2_0 ([F)[B + public final fun planesRowScale ([BIIII)F + public static synthetic fun planesRowScale$default (Lsk/ainet/lang/memory/TernaryCodec;[BIIIIILjava/lang/Object;)F } public final class sk/ainet/lang/memory/ViewsKt { @@ -4690,6 +4697,38 @@ public final class sk/ainet/lang/tensor/data/BitNetB158TensorData$Companion { public final fun fromRawBytes (Lsk/ainet/lang/tensor/Shape;[B)Lsk/ainet/lang/tensor/data/BitNetB158TensorData; } +public final class sk/ainet/lang/tensor/data/BitNetPlanesTensorData : sk/ainet/lang/tensor/data/TensorData, sk/ainet/lang/tensor/storage/PackedBlockStorage { + public static final field Companion Lsk/ainet/lang/tensor/data/BitNetPlanesTensorData$Companion; + public fun (Lsk/ainet/lang/tensor/Shape;[B)V + public fun copyToFloatArray ()[F + public fun dequantizeBlock (I[FI)V + public fun get ([I)Ljava/lang/Float; + public synthetic fun get ([I)Ljava/lang/Object; + public fun getBlockCount ()I + public fun getBlockOrder ()Lsk/ainet/lang/memory/BlockOrder; + public fun getBlockSize ()I + public final fun getCols ()I + public fun getElementCount ()J + public fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; + public fun getPackedData ()[B + public fun getPackedView ()Lsk/ainet/lang/memory/TensorView; + public fun getPhysicalBytes ()J + public final fun getRows ()I + public fun getShape ()Lsk/ainet/lang/tensor/Shape; + public fun getView ()Lsk/ainet/lang/memory/TensorView; + public final fun rowScale (I)F + public fun set ([IF)V + public synthetic fun set ([ILjava/lang/Object;)V + public fun toFloatArray ()[F + public fun toTensorStorage (Lsk/ainet/lang/tensor/storage/LogicalDType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; + public fun toTensorStorage (Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/storage/Placement;)Lsk/ainet/lang/tensor/storage/TensorStorage; +} + +public final class sk/ainet/lang/tensor/data/BitNetPlanesTensorData$Companion { + public final fun fromFloats (Lsk/ainet/lang/tensor/Shape;[F)Lsk/ainet/lang/tensor/data/BitNetPlanesTensorData; + public final fun fromRawBytes (Lsk/ainet/lang/tensor/Shape;[B)Lsk/ainet/lang/tensor/data/BitNetPlanesTensorData; +} + public final class sk/ainet/lang/tensor/data/DenseFloatArrayTensorData : sk/ainet/lang/tensor/data/FloatArrayTensorData { public fun (Lsk/ainet/lang/tensor/Shape;[F)V public fun copyToFloatArray ()[F @@ -7417,6 +7456,20 @@ public final class sk/ainet/lang/tensor/storage/TensorEncoding$BITNET_B1_58 : sk public fun toString ()Ljava/lang/String; } +public final class sk/ainet/lang/tensor/storage/TensorEncoding$BITNET_PLANES : sk/ainet/lang/tensor/storage/TensorEncoding { + public static final field INSTANCE Lsk/ainet/lang/tensor/storage/TensorEncoding$BITNET_PLANES; + public static final field PLANES I + public static final field ROW_SCALE_BYTES I + public final fun bufferBytes (II)I + public fun equals (Ljava/lang/Object;)Z + public fun getName ()Ljava/lang/String; + public fun hashCode ()I + public fun physicalBytes (J)Ljava/lang/Long; + public final fun planeStrideBytes (II)I + public final fun rowScalesByteOffset (II)I + public fun toString ()Ljava/lang/String; +} + public final class sk/ainet/lang/tensor/storage/TensorEncoding$DENSE_I8_ABSMAX : sk/ainet/lang/tensor/storage/TensorEncoding { public static final field CODE_RANGE I public static final field INSTANCE Lsk/ainet/lang/tensor/storage/TensorEncoding$DENSE_I8_ABSMAX;