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 ee31af11..8fe1cf91 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 @@ -50,7 +50,10 @@ internal class DefaultCpuOpsJvm( override fun withSchedule(schedule: sk.ainet.context.schedule.Schedule): sk.ainet.lang.tensor.ops.TensorOps = if (schedule === this.schedule) this else DefaultCpuOpsJvm(dataFactory, schedule) - private val floatSpecies: VectorSpecies = FloatVector.SPECIES_PREFERRED + // Read through JvmVectorKernels (an object → static final field): C2 constant-folds static + // finals but not instance finals, and the Vector API only intrinsifies with a constant species. + // As an instance field every FloatVector op here ran the allocating Java fallback. + private inline val floatSpecies: VectorSpecies get() = JvmVectorKernels.floatSpecies /** * On the JVM, auto-install ServiceLoader-discovered providers (Panama Vector, @@ -917,7 +920,9 @@ internal class DefaultCpuOpsJvm( return super.mean(tensor, dim) } - private fun vectorFloatBinary( + // inline for the same reason as JvmVectorKernels.binaryFloat: the op lambdas must reach + // the vector loops as code, not as Function2 objects, or the Vector API is not intrinsified. + private inline fun vectorFloatBinary( a: Tensor, b: Tensor, vectorOp: (FloatVector, FloatVector) -> FloatVector, @@ -950,37 +955,11 @@ internal class DefaultCpuOpsJvm( // Case 2: scalar broadcast if (aVol == 1) { - 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, bBuf, bBase + idx) - vectorOp(va, vb).intoArray(outBuffer, idx) - idx += speciesLen - } - while (idx < outVolume) { - outBuffer[idx] = scalarOp(aval, bBuf[bBase + idx]) - idx++ - } + JvmVectorKernels.binaryFloatScalarLeft(aBuf[aBase], bBuf, bBase, outBuffer, outVolume, vectorOp, scalarOp) return floatResult(Shape(outShape.dimensions.copyOf()), a.dtype, outBuffer) } if (bVol == 1) { - val bval = bBuf[bBase] - val speciesLen = floatSpecies.length() - var idx = 0 - val loopBound = floatSpecies.loopBound(outVolume) - while (idx < loopBound) { - 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(aBuf[aBase + idx], bval) - idx++ - } + JvmVectorKernels.binaryFloatScalarRight(aBuf, aBase, bBuf[bBase], outBuffer, outVolume, vectorOp, scalarOp) return floatResult(Shape(outShape.dimensions.copyOf()), a.dtype, outBuffer) } @@ -998,40 +977,22 @@ internal class DefaultCpuOpsJvm( a.shape.rank >= 1 && aLast == outLast && a.shape.dimensions.dropLast(1).all { it == 1 } ) if (bIsBias && aVol == outVolume) { - val step = floatSpecies.length() - val loopBoundTail = floatSpecies.loopBound(outLast) for (g in 0 until groups) { - val aOff = g * outLast - var idx = 0 - while (idx < loopBoundTail) { - 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(aBuf[aBase + aOff + idx], bBuf[bBase + idx]) - idx++ - } + val off = g * outLast + JvmVectorKernels.binaryFloat( + aBuf, bBuf, outBuffer, outLast, vectorOp, scalarOp, + aOffset = aBase + off, bOffset = bBase, outOffset = off, + ) } return floatResult(Shape(outShape.dimensions.copyOf()), a.dtype, outBuffer) } if (aIsBias && bVol == outVolume) { - val step = floatSpecies.length() - val loopBoundTail = floatSpecies.loopBound(outLast) for (g in 0 until groups) { - val bOff = g * outLast - var idx = 0 - while (idx < loopBoundTail) { - 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(aBuf[aBase + idx], bBuf[bBase + bOff + idx]) - idx++ - } + val off = g * outLast + JvmVectorKernels.binaryFloat( + aBuf, bBuf, outBuffer, outLast, vectorOp, scalarOp, + aOffset = aBase, bOffset = bBase + off, outOffset = off, + ) } return floatResult(Shape(outShape.dimensions.copyOf()), a.dtype, outBuffer) } @@ -1041,7 +1002,7 @@ internal class DefaultCpuOpsJvm( return null } - private fun vectorFloatUnary( + private inline fun vectorFloatUnary( tensor: Tensor, vectorOp: (FloatVector) -> FloatVector, scalarOp: (Float) -> Float 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 37c82539..d9749ab5 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 @@ -14,10 +14,14 @@ import java.nio.ByteOrder * and future API changes are easier to adapt. */ internal object JvmVectorKernels { - private val floatSpecies: VectorSpecies = FloatVector.SPECIES_PREFERRED + @PublishedApi + internal val floatSpecies: VectorSpecies = FloatVector.SPECIES_PREFERRED // region Elementwise - fun binaryFloat( + // inline: each call site gets its own loop with the lambda body in place, so C2 sees a + // monomorphic FloatVector op it can intrinsify. Passed through a Function1/2 the vectors + // escape into the generic invoke and the Vector API runs its allocating Java fallback. + inline fun binaryFloat( a: FloatArray, b: FloatArray, out: FloatArray, @@ -26,6 +30,7 @@ internal object JvmVectorKernels { scalarOp: (Float, Float) -> Float, aOffset: Int = 0, bOffset: Int = 0, + outOffset: Int = 0, ) { var index = 0 val step = floatSpecies.length() @@ -33,16 +38,64 @@ internal object JvmVectorKernels { while (index < loopBound) { val va = FloatVector.fromArray(floatSpecies, a, aOffset + index) val vb = FloatVector.fromArray(floatSpecies, b, bOffset + index) - op(va, vb).intoArray(out, index) + op(va, vb).intoArray(out, outOffset + index) index += step } while (index < length) { - out[index] = scalarOp(a[aOffset + index], b[bOffset + index]) + out[outOffset + index] = scalarOp(a[aOffset + index], b[bOffset + index]) index++ } } - fun unaryFloat( + /** `out[i] = op(a[aOffset + i], s)` — scalar right operand, same loop shape as [binaryFloat]. */ + inline fun binaryFloatScalarRight( + a: FloatArray, + aOffset: Int, + s: Float, + out: FloatArray, + length: Int, + op: (FloatVector, FloatVector) -> FloatVector, + scalarOp: (Float, Float) -> Float, + ) { + var index = 0 + val step = floatSpecies.length() + val loopBound = floatSpecies.loopBound(length) + val vs = FloatVector.broadcast(floatSpecies, s) + while (index < loopBound) { + op(FloatVector.fromArray(floatSpecies, a, aOffset + index), vs).intoArray(out, index) + index += step + } + while (index < length) { + out[index] = scalarOp(a[aOffset + index], s) + index++ + } + } + + /** `out[i] = op(s, b[bOffset + i])` — scalar left operand, same loop shape as [binaryFloat]. */ + inline fun binaryFloatScalarLeft( + s: Float, + b: FloatArray, + bOffset: Int, + out: FloatArray, + length: Int, + op: (FloatVector, FloatVector) -> FloatVector, + scalarOp: (Float, Float) -> Float, + ) { + var index = 0 + val step = floatSpecies.length() + val loopBound = floatSpecies.loopBound(length) + val vs = FloatVector.broadcast(floatSpecies, s) + while (index < loopBound) { + op(vs, FloatVector.fromArray(floatSpecies, b, bOffset + index)).intoArray(out, index) + index += step + } + while (index < length) { + out[index] = scalarOp(s, b[bOffset + index]) + index++ + } + } + + inline fun unaryFloat( input: FloatArray, out: FloatArray, length: Int,