Skip to content

perf(backend-cpu): let C2 intrinsify the elementwise Vector API paths #1360

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
michalharakal merged 1 commit into from
Oct 17, 2026
Merged

perf(backend-cpu): let C2 intrinsify the elementwise Vector API paths #1360

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 20 additions & 59 deletions ...kends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOpsJvm.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Float> = 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<Float> get() = JvmVectorKernels.floatSpecies

/**
* On the JVM, auto-install ServiceLoader-discovered providers (Panama Vector,
Expand Down Expand Up @@ -917,7 +920,9 @@ internal class DefaultCpuOpsJvm(
return super.mean(tensor, dim)
}

private fun <T : DType, V> 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 <T : DType, V> vectorFloatBinary(
a: Tensor<T, V>,
b: Tensor<T, V>,
vectorOp: (FloatVector, FloatVector) -> FloatVector,
Expand Down Expand Up @@ -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)
}

Expand All @@ -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)
}
Expand All @@ -1041,7 +1002,7 @@ internal class DefaultCpuOpsJvm(
return null
}

private fun <T : DType, V> vectorFloatUnary(
private inline fun <T : DType, V> vectorFloatUnary(
tensor: Tensor<T, V>,
vectorOp: (FloatVector) -> FloatVector,
scalarOp: (Float) -> Float
Expand Down
63 changes: 58 additions & 5 deletions ...kends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/tensor/ops/JvmVectorKernels.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,14 @@ import java.nio.ByteOrder
* and future API changes are easier to adapt.
*/
internal object JvmVectorKernels {
private val floatSpecies: VectorSpecies<Float> = FloatVector.SPECIES_PREFERRED
@PublishedApi
internal val floatSpecies: VectorSpecies<Float> = 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,
Expand All @@ -26,23 +30,72 @@ internal object JvmVectorKernels {
scalarOp: (Float, Float) -> Float,
aOffset: Int = 0,
bOffset: Int = 0,
outOffset: Int = 0,
) {
var index = 0
val step = floatSpecies.length()
val loopBound = floatSpecies.loopBound(length)
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,
Expand Down
Loading