Skip to content

feat(io-onnx): import LeakyRelu, Silu, Clip, BatchNormalization and Pad as real ops #1466

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 31, 2026
Merged

feat(io-onnx): import LeakyRelu, Silu, Clip, BatchNormalization and Pad as real ops #1466

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
40 changes: 39 additions & 1 deletion ...kainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,42 @@ public class ComputeGraphExecutor(
"exp" -> listOf(ops.exp(inputs[0]))
"expm1" -> listOf(ops.expm1(inputs[0]))
"sqrt" -> listOf(ops.sqrt(inputs[0]))
"leakyRelu", "leaky_relu" -> {
val negativeSlope = (params["negativeSlope"] as? Number)?.toFloat() ?: 0.01f
listOf(ops.leakyRelu(inputs[0], negativeSlope))
}
"clamp" -> {
val minVal = (params["minVal"] as? Number)?.toFloat() ?: -Float.MAX_VALUE
val maxVal = (params["maxVal"] as? Number)?.toFloat() ?: Float.MAX_VALUE
listOf(ops.clamp(inputs[0], minVal, maxVal))
}
"pad2d" -> {
val padLeft = (params["padLeft"] as? Number)?.toInt() ?: 0
val padRight = (params["padRight"] as? Number)?.toInt() ?: 0
val padTop = (params["padTop"] as? Number)?.toInt() ?: 0
val padBottom = (params["padBottom"] as? Number)?.toInt() ?: 0
listOf(ops.pad2d(inputs[0], padLeft, padRight, padTop, padBottom))
}
"batchNorm", "batchNormalization" -> {
// Inference-mode batch norm over NC*: y = scale·(x − mean)/√(var + ε) + bias,
// with the per-channel [C] operands reshaped to (1, C, 1, …) to broadcast.
require(inputs.size >= 5) {
"batchNorm requires 5 inputs (x, scale, bias, mean, variance), got ${inputs.size}"
}
val x = inputs[0]
val epsilon = (params["epsilon"] as? Number)?.toFloat() ?: 1e-5f
val channels = inputs[1].shape[0]
val broadcastShape = Shape(IntArray(x.rank) { idx -> if (idx == 1) channels else 1 })
val scale = ops.reshape(inputs[1], broadcastShape)
val bias = ops.reshape(inputs[2], broadcastShape)
val mean = ops.reshape(inputs[3], broadcastShape)
val variance = ops.reshape(inputs[4], broadcastShape)
val normalized = ops.divide(
ops.subtract(x, mean),
ops.sqrt(ops.addScalar(variance, epsilon))
)
listOf(ops.add(ops.multiply(normalized, scale), bias))
}

// Binary math ops
"add" -> listOf(ops.add(inputs[0], inputs[1]))
Expand Down Expand Up @@ -394,7 +430,9 @@ public class ComputeGraphExecutor(
)

public companion object {
private val INPUT_OPS = setOf("input", "weight", "parameter", "constant")
// "onnx_initializer" marks imported ONNX initializer nodes: like inputs, their values
// come from the caller (keyed by node id or tensor name), never from dispatching.
private val INPUT_OPS = setOf("input", "weight", "parameter", "constant", "onnx_initializer")

private val fusedOpHandlers = mutableMapOf<String, FusedOpHandler<*, *>>()

Expand Down
2 changes: 2 additions & 0 deletions skainet-io/skainet-io-onnx/build.gradle.kts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,8 @@ kotlin {
dependencies {
implementation(libs.junit)
implementation(libs.kotlinx.coroutines.test)
// Real CPU ops for import-and-execute tests of the mapped ONNX operations.
implementation(project(":skainet-backends:skainet-backend-cpu"))
}
}

Expand Down
189 changes: 173 additions & 16 deletions skainet-io/skainet-io-onnx/src/commonMain/kotlin/sk/ainet/io/onnx/OnnxImport.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,9 +118,17 @@ public class OnnxToComputeGraphImporter {

// Map each ONNX node to a ComputeGraph node and wire edges.
graph.nodes.forEach { nodeView ->
val op = nodeView.toOperation()
val op = nodeView.toOperation(graph.initializers)

// Ops like Clip/Pad fold their trailing initializer inputs (min/max, pads) into
// parameters at import time — only the data input stays wired.
val nodeInputs = if (op.parameters[ONNX_INPUTS_CONSUMED_AT_IMPORT] == true) {
nodeView.inputs.take(1)
} else {
nodeView.inputs
}

val inputSpecs = nodeView.inputs.mapNotNull { inputName ->
val inputSpecs = nodeInputs.mapNotNull { inputName ->
if (inputName.isEmpty()) return@mapNotNull null
val produced = ensureProducer(inputName, tensorInfo, computeGraph, producers)
produced.tensorSpec
Expand All @@ -130,7 +138,7 @@ public class OnnxToComputeGraphImporter {
resolveTensorSpec(outputName, tensorInfo, inputSpecs, op, idx)
}

val initMetadata = graph.initializers.filterKeys { it in nodeView.inputs }
val initMetadata = graph.initializers.filterKeys { it in nodeInputs }
val graphNode = GraphNode(
id = nodeView.name,
operation = op,
Expand All @@ -144,7 +152,7 @@ public class OnnxToComputeGraphImporter {
computeGraph.addNode(graphNode)

// Wire edges now that we have a concrete destination.
nodeView.inputs.forEachIndexed { inputIdx, inputName ->
nodeInputs.forEachIndexed { inputIdx, inputName ->
if (inputName.isEmpty()) return@forEachIndexed
val produced = ensureProducer(inputName, tensorInfo, computeGraph, producers)
computeGraph.addEdge(
Expand Down Expand Up @@ -223,7 +231,7 @@ public class OnnxToComputeGraphImporter {
)
}

private fun OnnxNodeView.toOperation(): Operation {
private fun OnnxNodeView.toOperation(initializers: Map<String, OnnxTensorView>): Operation {
val upper = opType.uppercase()
val attrInts = { key: String -> attributes[key]?.ints?.map { it.toSafeInt() } }
return when (upper) {
Expand All @@ -238,16 +246,67 @@ public class OnnxToComputeGraphImporter {
)
"RELU" -> ReluOperation<DType, Any>()
"SIGMOID" -> SigmoidOperation<DType, Any>()
"LEAKYRELU" -> OnnxPlaceholderOperation(
opType = opType,
type = "activation",
parameters = attributes["alpha"]?.floats?.firstOrNull()?.let { mapOf("alpha" to it) } ?: emptyMap()
)
"SILU" -> OnnxPlaceholderOperation(
opType = opType,
"LEAKYRELU" -> OnnxMappedOperation(
opName = "leakyRelu",
type = "activation",
parameters = emptyMap()
// ONNX attribute `alpha` (default 0.01) is SKaiNET's negativeSlope.
parameters = mapOf("negativeSlope" to (attributes["alpha"]?.floats?.firstOrNull() ?: 0.01f))
)
"SILU" -> OnnxMappedOperation(opName = "silu", type = "activation")
"CLIP" -> {
// Opset < 11 carries min/max as attributes; opset >= 11 as optional initializer
// inputs. A min/max fed by a *dynamic* tensor cannot fold to a clamp — keep the
// placeholder so execution fails loudly instead of silently not clipping.
val minResolved = attributes["min"]?.floats?.firstOrNull()
?: resolveScalarFloatInitializer(inputs.getOrNull(1), initializers)
val maxResolved = attributes["max"]?.floats?.firstOrNull()
?: resolveScalarFloatInitializer(inputs.getOrNull(2), initializers)
val hasDynamicBound = listOfNotNull(inputs.getOrNull(1), inputs.getOrNull(2))
.any { it.isNotEmpty() && it !in initializers }
if (hasDynamicBound) {
OnnxPlaceholderOperation(
opType = opType,
type = "onnx",
parameters = mapOf("reason" to "Clip with dynamic min/max inputs is not supported")
)
} else {
OnnxMappedOperation(
opName = "clamp",
type = "activation",
parameters = mapOf(
"minVal" to (minResolved ?: -Float.MAX_VALUE),
"maxVal" to (maxResolved ?: Float.MAX_VALUE),
ONNX_INPUTS_CONSUMED_AT_IMPORT to true,
)
)
}
}
"PAD" -> {
// Only the form pad2d covers: constant mode, zero fill, 4D input with no
// batch/channel padding. pads order (4D): [Nb, Cb, Hb, Wb, Ne, Ce, He, We].
val mode = attributes["mode"]?.strings?.firstOrNull() ?: "constant"
val pads = attrInts("pads")
?: inputs.getOrNull(1)?.takeIf { it.isNotEmpty() }?.let { initializers[it]?.intListOrNull() }
val constantValue = attributes["value"]?.floats?.firstOrNull()
?: resolveScalarFloatInitializer(inputs.getOrNull(2), initializers)
?: 0f
if (mode == "constant" && constantValue == 0f &&
pads != null && pads.size == 8 &&
pads[0] == 0 && pads[1] == 0 && pads[4] == 0 && pads[5] == 0
) {
OnnxPadOperation(padTop = pads[2], padLeft = pads[3], padBottom = pads[6], padRight = pads[7])
} else {
OnnxPlaceholderOperation(
opType = opType,
type = "onnx",
parameters = mapOf(
"reason" to "Pad supports constant mode with zero value and 4D spatial-only pads",
"mode" to mode,
"pads" to (pads ?: emptyList<Int>()),
)
)
}
}
"ADD" -> AddOperation<DType, Any>()
"MUL" -> MultiplyOperation<DType, Any>()
"DIV" -> DivideOperation<DType, Any>()
Expand All @@ -266,10 +325,12 @@ public class OnnxToComputeGraphImporter {
"TRANSPOSE" -> TransposeOperation<DType, Any>(
parameters = attrInts("perm")?.let { mapOf("perm" to it) } ?: emptyMap()
)
"BATCHNORMALIZATION" -> OnnxPlaceholderOperation(
opType = opType,
"BATCHNORMALIZATION" -> OnnxMappedOperation(
opName = "batchNorm",
type = "normalization",
parameters = attributes.mapValues { (_, v) -> v.ints.ifEmpty { v.floats } }
// Inference-mode BN: inputs are (X, scale, B, mean, var); only epsilon matters
// (momentum/spatial are training-time attributes).
parameters = mapOf("epsilon" to (attributes["epsilon"]?.floats?.firstOrNull() ?: 1e-5f))
)
"CONCAT" -> OnnxConcatOperation(
axis = attributes["axis"]?.ints?.firstOrNull()?.toSafeInt() ?: 0
Expand Down Expand Up @@ -353,6 +414,7 @@ private fun opShapeInference(op: Operation, inputs: List<TensorSpec>, outputInde
is OnnxConcatOperation -> op.inferOutputs(inputs).getOrNull(outputIndex)?.shape
is OnnxSliceOperation -> op.inferOutputs(inputs).getOrNull(outputIndex)?.shape
is OnnxSplitOperation -> op.inferOutputs(inputs).getOrNull(outputIndex)?.shape
is OnnxPadOperation -> op.inferOutputs(inputs).getOrNull(outputIndex)?.shape
is Conv2dOperation<*, *> -> inputs.firstOrNull()?.shape
is MaxPool2dOperation<*, *> -> inputs.firstOrNull()?.shape
is Upsample2dOperation<*, *> -> op.inferOutputs(inputs).getOrNull(outputIndex)?.shape
Expand Down Expand Up @@ -407,10 +469,105 @@ private fun Long.toSafeInt(): Int = when {
else -> toInt()
}

/**
* Marks an operation whose trailing ONNX inputs (e.g. Clip's min/max, Pad's pads tensor) were
* folded into [Operation.parameters] at import time: only the first (data) input is wired.
*/
internal const val ONNX_INPUTS_CONSUMED_AT_IMPORT: String = "onnxInputsConsumedAtImport"

/** Read a scalar FLOAT initializer (float_data or little-endian raw_data), or null. */
private fun resolveScalarFloatInitializer(
inputName: String?,
initializers: Map<String, OnnxTensorView>
): Float? {
if (inputName.isNullOrEmpty()) return null
val proto = initializers[inputName]?.proto ?: return null
proto.floatData.firstOrNull()?.let { return it }
val raw = proto.rawData.array
if (proto.dataType == TensorProto.DataType.FLOAT.value && raw.size >= 4) {
val bits = (raw[0].toInt() and 0xFF) or
((raw[1].toInt() and 0xFF) shl 8) or
((raw[2].toInt() and 0xFF) shl 16) or
((raw[3].toInt() and 0xFF) shl 24)
return Float.fromBits(bits)
}
return null
}

/** Read an INT64 initializer as an Int list (int64_data or little-endian raw_data), or null. */
private fun OnnxTensorView.intListOrNull(): List<Int>? {
if (proto.int64Data.isNotEmpty()) return proto.int64Data.map { it.toSafeInt() }
val raw = proto.rawData.array
if (proto.dataType == TensorProto.DataType.INT64.value && raw.isNotEmpty() && raw.size % 8 == 0) {
return List(raw.size / 8) { i ->
var v = 0L
for (b in 7 downTo 0) {
v = (v shl 8) or (raw[i * 8 + b].toLong() and 0xFF)
}
v.toSafeInt()
}
}
return null
}

private class InitializerOperation(
tensorName: String
) : BasePlaceholderOperation("onnx_initializer", mapOf("name" to tensorName))

/**
* An ONNX node mapped to a real SKaiNET operation: the name matches the ComputeGraphExecutor
* dispatch name and [parameters] carries the attributes it reads. `execute` is never called
* directly — the executor dispatches on the operation name, the same way OnnxConcatOperation
* and friends run.
*/
private class OnnxMappedOperation(
opName: String,
type: String,
parameters: Map<String, Any> = emptyMap(),
) : BasePlaceholderOperation(opName, parameters, type)

/**
* ONNX `Pad` folded to SKaiNET `pad2d` (constant mode, zero fill, 4D, spatial-only pads).
* Carries its own shape inference: the output grows by the pads on H and W.
*/
private class OnnxPadOperation(
private val padTop: Int,
private val padLeft: Int,
private val padBottom: Int,
private val padRight: Int,
) : BasePlaceholderOperation(
"pad2d",
mapOf(
"padTop" to padTop,
"padLeft" to padLeft,
"padBottom" to padBottom,
"padRight" to padRight,
ONNX_INPUTS_CONSUMED_AT_IMPORT to true,
),
"shape",
) {
override fun inferOutputs(inputs: List<TensorSpec>): List<TensorSpec> {
val inputShape = inputs.firstOrNull()?.shape
val outputShape = if (inputShape != null && inputShape.size == 4) {
listOf(
inputShape[0],
inputShape[1],
inputShape[2] + padTop + padBottom,
inputShape[3] + padLeft + padRight,
)
} else {
inputShape
}
return listOf(
TensorSpec(
name = "pad2d_output",
shape = outputShape,
dtype = inputs.firstOrNull()?.dtype ?: "unknown"
)
)
}
}

private class OnnxPlaceholderOperation(
opType: String,
type: String = "onnx",
Expand Down
Loading
Loading