LinalgMicrobench.kt
src/jvmMain/kotlin/afield/examples/LinalgMicrobench.kt · 219 lines
package afield.examples
import afield.math.linalg.Vector3d
import afield.math.linalg.Vector3f
import afield.math.linalg.arena.FrameArena
import afield.math.linalg.storage.SegmentVec3Buffer
import afield.math.linalg.storage.Vec3StorageLayout
import afield.math.linalg.vec
import java.lang.management.ManagementFactory
import java.nio.ByteBuffer
object LinalgMicrobench
Screening microbenchmark for the generated linear-algebra backends.
This is deliberately NOT a rigorous JMH harness. A short in-process loop cannot reliably defeat HotSpot escape analysis, so absolute nanoseconds are indicative only. The purpose is to screen for gross mistakes: a backend that is an order of magnitude off its siblings, an allocation storm hiding behind a "destination" API, or arena validation that costs far more than the arithmetic it guards. For real numbers, promote the winning cases to a JMH module.
Run: run-jvm-main --timeout 120 afield.examples.LinalgMicrobench or set FIELD_MICROBENCH_ITERS / FIELD_MICROBENCH_WARMUP to override the loop sizes.
FIELD_TEST id: examples.linalg-microbench runner: jvm-main entrypoint: afield.examples.LinalgMicrobench tags: [bench, linalg] timeoutSec: 180 expect: exitCode: 0 stdoutContains: - "Generated linalg microbenchmark completed"
object LinalgMicrobench {
private const val SCALE = 1.0001f
private const val SCALE_D = 1.0001
fun main(args: Array<String>)
body — 64 lines, from line 40
@JvmStatic
fun main(args: Array<String>) {
val iters = (System.getenv("FIELD_MICROBENCH_ITERS")?.toIntOrNull() ?: 5_000_000)
val warmup = (System.getenv("FIELD_MICROBENCH_WARMUP")?.toIntOrNull() ?: 2_000_000)
println("Linalg microbenchmark — warmup=$warmup iters=$iters")
println("arena checks compiled in: ${FrameArena.checksCompiledIn}")
println()
val results = mutableListOf<Row>()
results += measure("heap Vector3f (alloc per op)", warmup, iters, ::heapAllocKernel)
results += measure("heap Vector3f (reused dest)", warmup, iters, ::heapReuseKernel)
results += measure("arena ArenaVec3f (checks ON)", warmup, iters) { n -> arenaKernel(n, debugChecks = true) }
results += measure("arena ArenaVec3f (checks OFF)", warmup, iters) { n -> arenaKernel(n, debugChecks = false) }
results += measure("segment PACKED_12", warmup, iters) { n -> segmentKernel(n, Vec3StorageLayout.PACKED_12) }
results += measure("segment ALIGNED_16", warmup, iters) { n -> segmentKernel(n, Vec3StorageLayout.ALIGNED_16) }
results += measure("heap Vector3d (reused dest)", warmup, iters, ::doubleReuseKernel)
report("Vector3 kernel (add, scalar-mul, cross, dot, normalize)", results, iters)
// The float backends run identical arithmetic; their checksums must agree bit-for-bit.
val floatChecks = results.dropLast(1).map { it.checksum }
val agree = floatChecks.all { it == floatChecks.first() }
println("float backends checksum agreement: ${if (agree) "OK (${floatChecks.first()})" else "MISMATCH $floatChecks"}")
println()
val ctor = mutableListOf<Row>()
ctor += measure("Vector3f(x, y, z)", warmup, iters, ::ctorRawKernel)
ctor += measure("vec(x, y, z)", warmup, iters, ::ctorVecKernel)
ctor += measure("Vector3f.of(x, y, z) [literal]", warmup, iters, ::ctorLiteralKernel)
report("Construction", ctor, iters)
check(agree) { "Float backends disagree — a storage lowering is wrong, not just slow." }
println("Generated linalg microbenchmark completed")
}
private data class Row(val name: String, val nanos: Long, val gcCount: Long, val checksum: Double)
private fun measure(name: String, warmup: Int, iters: Int, kernel: (Int) -> Double): Row {
kernel(warmup)
val gcBefore = gcCount()
val start = System.nanoTime()
val checksum = kernel(iters)
val elapsed = System.nanoTime() - start
return Row(name, elapsed, gcCount() - gcBefore, checksum)
}
private fun report(title: String, rows: List<Row>, iters: Int) {
val baseline = rows.first().nanos.toDouble()
println(title)
println(" %-34s %12s %10s %8s %s".format("backend", "ns/op", "vs base", "GCs", "checksum"))
rows.forEach { row ->
val nsPerOp = row.nanos.toDouble() / iters
val ratio = row.nanos / baseline
println(
" %-34s %12.3f %9.2fx %8d %s".format(
row.name, nsPerOp, ratio, row.gcCount, "%.4g".format(row.checksum)
)
)
}
println()
}
private fun gcCount(): Long =
ManagementFactory.getGarbageCollectorMXBeans().sumOf { maxOf(it.collectionCount, 0) }
kernels: identical math, different storage
source — 82 lines, from line 107
private fun heapAllocKernel(n: Int): Double {
val a = Vector3f(1f, 2f, 3f)
val b = Vector3f(4f, 5f, 6f)
var acc = 0.0
for (i in 0 until n) {
a.x += 1e-3f * (i and 1023)
val c = a.add(b, Vector3f()).mul(SCALE, Vector3f())
val d = a.cross(b, Vector3f())
val cn = c.normalize(Vector3f())
acc += (cn.dot(b) + d.x).toDouble()
}
return acc
}
private fun heapReuseKernel(n: Int): Double {
val a = Vector3f(1f, 2f, 3f)
val b = Vector3f(4f, 5f, 6f)
val c = Vector3f()
val d = Vector3f()
var acc = 0.0
for (i in 0 until n) {
a.x += 1e-3f * (i and 1023)
a.add(b, c); c.mul(SCALE, c)
a.cross(b, d)
c.normalize(c)
acc += (c.dot(b) + d.x).toDouble()
}
return acc
}
private fun doubleReuseKernel(n: Int): Double {
val a = Vector3d(1.0, 2.0, 3.0)
val b = Vector3d(4.0, 5.0, 6.0)
val c = Vector3d()
val d = Vector3d()
var acc = 0.0
for (i in 0 until n) {
a.x += 1e-3 * (i and 1023)
a.add(b, c); c.mul(SCALE_D, c)
a.cross(b, d)
c.normalize(c)
acc += c.dot(b) + d.x
}
return acc
}
private fun arenaKernel(n: Int, debugChecks: Boolean): Double {
var acc = 0.0
FrameArena(256, debugChecks = debugChecks).use { arena ->
arena.scoped {
val a = vec3(1f, 2f, 3f)
val b = vec3(4f, 5f, 6f)
val c = vec3()
val d = vec3()
for (i in 0 until n) {
a.x += 1e-3f * (i and 1023)
a.add(b, c); c.mul(SCALE, c)
a.cross(b, d)
c.normalize(c)
acc += (c.dot(b) + d.x).toDouble()
}
}
}
return acc
}
private fun segmentKernel(n: Int, layout: Vec3StorageLayout): Double {
var acc = 0.0
val buffer = SegmentVec3Buffer.of(ByteBuffer.allocateDirect(4 * 16), 4, layout)
with(buffer) {
val a = buffer[0]; val b = buffer[1]; val c = buffer[2]; val d = buffer[3]
a.set(1f, 2f, 3f); b.set(4f, 5f, 6f)
for (i in 0 until n) {
a.x += 1e-3f * (i and 1023)
a.add(b, c); c.mul(SCALE, c)
a.cross(b, d)
c.normalize(c)
acc += (c.dot(b) + d.x).toDouble()
}
}
return acc
}
construction kernels
private fun ctorRawKernel(n: Int): Double {
var acc = 0.0
for (i in 0 until n) {
val v = Vector3f(i.toFloat(), (i + 1).toFloat(), (i + 2).toFloat())
acc += v.x.toDouble()
}
return acc
}
private fun ctorVecKernel(n: Int): Double {
var acc = 0.0
for (i in 0 until n) {
val v = vec(i, i + 1, i + 2)
acc += v.x.toDouble()
}
return acc
}
private fun ctorLiteralKernel(n: Int): Double {
var acc = 0.0
for (i in 0 until n) {
val v = Vector3f.of(i, i + 1, i + 2)
acc += v.x.toDouble()
}
return acc
}
}