diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt index bca42eaa..a99b6e99 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt @@ -202,6 +202,32 @@ public class ComputeGraphExecutor( listOf(ops.add(ops.multiply(normalized, scale), bias)) } + // Convolution ops + "convTranspose1d" -> listOf( + ops.convTranspose1d( + input = inputs[0], + weight = inputs[1], + bias = inputs.getOrNull(2), + stride = (params["stride"] as? Number)?.toInt() ?: 1, + padding = (params["padding"] as? Number)?.toInt() ?: 0, + outputPadding = (params["outputPadding"] as? Number)?.toInt() ?: 0, + dilation = (params["dilation"] as? Number)?.toInt() ?: 1, + groups = (params["groups"] as? Number)?.toInt() ?: 1, + ) + ) + "convTranspose2d" -> listOf( + ops.convTranspose2d( + input = inputs[0], + weight = inputs[1], + bias = inputs.getOrNull(2), + stride = params.intPair("stride", 1 to 1), + padding = params.intPair("padding", 0 to 0), + outputPadding = params.intPair("outputPadding", 0 to 0), + dilation = params.intPair("dilation", 1 to 1), + groups = (params["groups"] as? Number)?.toInt() ?: 1, + ) + ) + // Binary math ops "add" -> listOf(ops.add(inputs[0], inputs[1])) "subtract" -> listOf(ops.subtract(inputs[0], inputs[1])) @@ -451,6 +477,23 @@ public class ComputeGraphExecutor( } } +/** + * Reads a per-axis Int pair recorded either as a `Pair` (direct op parameters), a two-element + * list (KSP traces and the ONNX importer), or a single number applied to both axes. + */ +private fun Map.intPair(key: String, default: Pair): Pair = + when (val value = this[key]) { + is Pair<*, *> -> ((value.first as? Number)?.toInt() ?: default.first) to + ((value.second as? Number)?.toInt() ?: default.second) + is List<*> -> if (value.size == 2 && value.all { it is Number }) { + (value[0] as Number).toInt() to (value[1] as Number).toInt() + } else { + default + } + is Number -> value.toInt() to value.toInt() + else -> default + } + /** * Handler for executing a fused operation. * diff --git a/skainet-io/skainet-io-onnx/src/commonMain/kotlin/sk/ainet/io/onnx/OnnxImport.kt b/skainet-io/skainet-io-onnx/src/commonMain/kotlin/sk/ainet/io/onnx/OnnxImport.kt index fbb345a9..55b8be99 100644 --- a/skainet-io/skainet-io-onnx/src/commonMain/kotlin/sk/ainet/io/onnx/OnnxImport.kt +++ b/skainet-io/skainet-io-onnx/src/commonMain/kotlin/sk/ainet/io/onnx/OnnxImport.kt @@ -244,6 +244,54 @@ public class OnnxToComputeGraphImporter { attrInts("kernel_shape")?.let { put("kernelShape", it) } } ) + "CONVTRANSPOSE" -> { + // ONNX weight layout equals TensorOps.convTranspose{1,2}d ([C_in, C_out/groups, k...]). + // ONNX pads are [*_begins..., *_ends...] per spatial axis while SKaiNET takes one + // symmetric padding per axis, so asymmetric pads (and auto_pad/output_shape, which + // derive pads at runtime) keep the placeholder and fail loudly on execute. + val strides = attrInts("strides") + val pads = attrInts("pads") + val dilations = attrInts("dilations") + val outputPadding = attrInts("output_padding") + val group = attributes["group"]?.ints?.firstOrNull()?.toSafeInt() ?: 1 + val autoPad = attributes["auto_pad"]?.strings?.firstOrNull() ?: "NOTSET" + val spatialRank = attrInts("kernel_shape")?.size + ?: strides?.size + ?: dilations?.size + ?: pads?.let { it.size / 2 } + ?: inputs.getOrNull(1)?.takeIf { it.isNotEmpty() } + ?.let { initializers[it]?.shape?.size?.minus(2) } + val symmetricPads = pads == null || + (pads.size == 2 * (spatialRank ?: 0) && + (0 until pads.size / 2).all { pads[it] == pads[it + pads.size / 2] }) + val supported = autoPad == "NOTSET" && "output_shape" !in attributes && + symmetricPads && (spatialRank == 1 || spatialRank == 2) + when { + !supported -> OnnxPlaceholderOperation( + opType = opType, + type = "onnx", + parameters = mapOf( + "reason" to "ConvTranspose supports 1D/2D with auto_pad=NOTSET, " + + "symmetric pads and no output_shape", + "autoPad" to autoPad, + ) + ) + spatialRank == 1 -> OnnxConvTranspose1dOperation( + stride = strides?.get(0) ?: 1, + padding = pads?.get(0) ?: 0, + outputPadding = outputPadding?.get(0) ?: 0, + dilation = dilations?.get(0) ?: 1, + groups = group, + ) + else -> OnnxConvTranspose2dOperation( + stride = strides ?: listOf(1, 1), + padding = pads?.take(2) ?: listOf(0, 0), + outputPadding = outputPadding ?: listOf(0, 0), + dilation = dilations ?: listOf(1, 1), + groups = group, + ) + } + } "RELU" -> ReluOperation() "SIGMOID" -> SigmoidOperation() "LEAKYRELU" -> OnnxMappedOperation( @@ -415,6 +463,8 @@ private fun opShapeInference(op: Operation, inputs: List, outputInde 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 OnnxConvTranspose1dOperation -> op.inferOutputs(inputs).getOrNull(outputIndex)?.shape + is OnnxConvTranspose2dOperation -> op.inferOutputs(inputs).getOrNull(outputIndex)?.shape is Conv2dOperation<*, *> -> inputs.firstOrNull()?.shape is MaxPool2dOperation<*, *> -> inputs.firstOrNull()?.shape is Upsample2dOperation<*, *> -> op.inferOutputs(inputs).getOrNull(outputIndex)?.shape @@ -568,6 +618,90 @@ private class OnnxPadOperation( } } +/** + * ONNX `ConvTranspose` (1D) mapped to SKaiNET `convTranspose1d`. Output length is + * `(in - 1) * stride - 2 * padding + dilation * (k - 1) + outputPadding + 1`. + */ +private class OnnxConvTranspose1dOperation( + private val stride: Int, + private val padding: Int, + private val outputPadding: Int, + private val dilation: Int, + private val groups: Int, +) : BasePlaceholderOperation( + "convTranspose1d", + mapOf( + "stride" to stride, + "padding" to padding, + "outputPadding" to outputPadding, + "dilation" to dilation, + "groups" to groups, + ), + "convolution", +) { + override fun inferOutputs(inputs: List): List { + val x = inputs.getOrNull(0)?.shape + val w = inputs.getOrNull(1)?.shape + val outputShape = if (x != null && w != null && x.size == 3 && w.size == 3) { + listOf( + x[0], + w[1] * groups, + (x[2] - 1) * stride - 2 * padding + dilation * (w[2] - 1) + outputPadding + 1, + ) + } else { + x + } + return listOf( + TensorSpec( + name = "convTranspose1d_output", + shape = outputShape, + dtype = inputs.firstOrNull()?.dtype ?: "unknown" + ) + ) + } +} + +/** ONNX `ConvTranspose` (2D) mapped to SKaiNET `convTranspose2d`; same output formula per H/W axis. */ +private class OnnxConvTranspose2dOperation( + private val stride: List, + private val padding: List, + private val outputPadding: List, + private val dilation: List, + private val groups: Int, +) : BasePlaceholderOperation( + "convTranspose2d", + mapOf( + "stride" to stride, + "padding" to padding, + "outputPadding" to outputPadding, + "dilation" to dilation, + "groups" to groups, + ), + "convolution", +) { + override fun inferOutputs(inputs: List): List { + val x = inputs.getOrNull(0)?.shape + val w = inputs.getOrNull(1)?.shape + val outputShape = if (x != null && w != null && x.size == 4 && w.size == 4) { + listOf( + x[0], + w[1] * groups, + (x[2] - 1) * stride[0] - 2 * padding[0] + dilation[0] * (w[2] - 1) + outputPadding[0] + 1, + (x[3] - 1) * stride[1] - 2 * padding[1] + dilation[1] * (w[3] - 1) + outputPadding[1] + 1, + ) + } else { + x + } + return listOf( + TensorSpec( + name = "convTranspose2d_output", + shape = outputShape, + dtype = inputs.firstOrNull()?.dtype ?: "unknown" + ) + ) + } +} + private class OnnxPlaceholderOperation( opType: String, type: String = "onnx", diff --git a/skainet-io/skainet-io-onnx/src/jvmTest/kotlin/sk/ainet/io/onnx/OnnxConvTransposeExecuteTest.kt b/skainet-io/skainet-io-onnx/src/jvmTest/kotlin/sk/ainet/io/onnx/OnnxConvTransposeExecuteTest.kt new file mode 100644 index 00000000..f3fe71bc --- /dev/null +++ b/skainet-io/skainet-io-onnx/src/jvmTest/kotlin/sk/ainet/io/onnx/OnnxConvTransposeExecuteTest.kt @@ -0,0 +1,260 @@ +package sk.ainet.io.onnx + +import onnx.AttributeProto +import onnx.GraphProto +import onnx.ModelProto +import onnx.NodeProto +import onnx.TensorProto +import onnx.TensorShapeProto +import onnx.TypeProto +import onnx.ValueInfoProto +import pbandk.ByteArr +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.Phase +import sk.ainet.exec.tensor.ops.DefaultCpuOps +import sk.ainet.lang.graph.exec.ComputeGraphExecutor +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.DenseTensorDataFactory +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * Import-and-execute coverage for ONNX `ConvTranspose` mapped to the real + * `convTranspose1d` / `convTranspose2d` operations (#1406). Each test builds a one-node + * ONNX model, imports it and runs it through [ComputeGraphExecutor] on the CPU backend. + */ +class OnnxConvTransposeExecuteTest { + + private val ctx = DirectCpuExecutionContext(phase = Phase.EVAL) + private val ops = DefaultCpuOps(DenseTensorDataFactory()) + + private fun tensor(shape: Shape, data: FloatArray): Tensor = + ctx.fromFloatArray(shape, FP32::class, data) + + private fun assertValues(expected: FloatArray, actual: Tensor, eps: Float = 1e-5f) { + val dims = actual.shape.dimensions + val idx = IntArray(dims.size) + for (flat in expected.indices) { + var rem = flat + for (d in dims.indices.reversed()) { + idx[d] = rem % dims[d] + rem /= dims[d] + } + val a = actual.data.get(*idx) as Float + assertTrue(abs(expected[flat] - a) <= eps, "Mismatch at $flat: expected ${expected[flat]}, got $a") + } + } + + private fun valueInfo(name: String, vararg dims: Long): ValueInfoProto = ValueInfoProto( + name = name, + type = TypeProto( + value = TypeProto.Value.TensorType( + TypeProto.Tensor( + elemType = TensorProto.DataType.FLOAT.value, + shape = TensorShapeProto( + dim = dims.map { TensorShapeProto.Dimension(value = TensorShapeProto.Dimension.Value.DimValue(it)) } + ) + ) + ) + ) + ) + + private fun floatInitializer(name: String, values: FloatArray, vararg dims: Long): TensorProto = TensorProto( + name = name, + dims = dims.toList(), + dataType = TensorProto.DataType.FLOAT.value, + floatData = values.toList(), + ) + + private fun importSingleNode( + node: NodeProto, + inputs: List, + initializers: List = emptyList(), + ) = OnnxToComputeGraphImporter().import( + ModelProto( + graph = GraphProto( + node = listOf(node), + input = inputs, + initializer = initializers, + ) + ).toGraphView() + ) + + @Test + fun convTranspose2d_defaultAttributes_importsAndExecutes() { + val node = NodeProto( + name = "ct0", + opType = "ConvTranspose", + input = listOf("x", "w"), + output = listOf("y"), + ) + val weight = floatArrayOf(1f, 1f, 1f, 1f) // [1, 1, 2, 2] ones + val graph = importSingleNode( + node, + inputs = listOf(valueInfo("x", 1, 1, 2, 2)), + initializers = listOf(floatInitializer("w", weight, 1, 1, 2, 2)), + ) + + assertEquals("convTranspose2d", graph.nodes.first { it.id == "ct0" }.operation.name) + + val result = ComputeGraphExecutor(graph, ops).execute( + mapOf( + "x" to tensor(Shape(1, 1, 2, 2), floatArrayOf(0f, 1f, 2f, 3f)), + "w" to tensor(Shape(1, 1, 2, 2), weight), + ) + ).getValue("ct0") + + // Full-overlap scatter of a 2x2 ones kernel over [[0,1],[2,3]]. + assertEquals(Shape(1, 1, 3, 3), result.shape) + assertValues(floatArrayOf(0f, 1f, 1f, 2f, 6f, 4f, 2f, 5f, 3f), result) + } + + @Test + fun convTranspose2d_strideAndBias_importsAndExecutes() { + val node = NodeProto( + name = "ct1", + opType = "ConvTranspose", + input = listOf("x", "w", "b"), + output = listOf("y"), + attribute = listOf(AttributeProto(name = "strides", ints = listOf(2L, 2L))), + ) + val weight = floatArrayOf(1f, 1f, 1f, 1f) + val graph = importSingleNode( + node, + inputs = listOf(valueInfo("x", 1, 1, 2, 2)), + initializers = listOf( + floatInitializer("w", weight, 1, 1, 2, 2), + floatInitializer("b", floatArrayOf(0.5f), 1), + ), + ) + + val result = ComputeGraphExecutor(graph, ops).execute( + mapOf( + "x" to tensor(Shape(1, 1, 2, 2), floatArrayOf(1f, 2f, 3f, 4f)), + "w" to tensor(Shape(1, 1, 2, 2), weight), + "b" to tensor(Shape(1), floatArrayOf(0.5f)), + ) + ).getValue("ct1") + + // Stride 2 with a 2x2 kernel tiles each input value into its own 2x2 block, plus bias. + assertEquals(Shape(1, 1, 4, 4), result.shape) + assertValues( + floatArrayOf( + 1.5f, 1.5f, 2.5f, 2.5f, + 1.5f, 1.5f, 2.5f, 2.5f, + 3.5f, 3.5f, 4.5f, 4.5f, + 3.5f, 3.5f, 4.5f, 4.5f, + ), + result, + ) + } + + @Test + fun convTranspose2d_symmetricPads_importsAndExecutes() { + val node = NodeProto( + name = "ct2", + opType = "ConvTranspose", + input = listOf("x", "w"), + output = listOf("y"), + attribute = listOf(AttributeProto(name = "pads", ints = listOf(1L, 1L, 1L, 1L))), + ) + val weight = floatArrayOf(1f, 1f, 1f, 1f) + val graph = importSingleNode( + node, + inputs = listOf(valueInfo("x", 1, 1, 2, 2)), + initializers = listOf(floatInitializer("w", weight, 1, 1, 2, 2)), + ) + + val result = ComputeGraphExecutor(graph, ops).execute( + mapOf( + "x" to tensor(Shape(1, 1, 2, 2), floatArrayOf(0f, 1f, 2f, 3f)), + "w" to tensor(Shape(1, 1, 2, 2), weight), + ) + ).getValue("ct2") + + // padding=1 crops the 3x3 full-overlap result [[0,1,1],[2,6,4],[2,5,3]] to its center. + assertEquals(Shape(1, 1, 1, 1), result.shape) + assertValues(floatArrayOf(6f), result) + } + + @Test + fun convTranspose1d_importsAndExecutes() { + val node = NodeProto( + name = "ct1d", + opType = "ConvTranspose", + input = listOf("x", "w"), + output = listOf("y"), + attribute = listOf(AttributeProto(name = "kernel_shape", ints = listOf(2L))), + ) + val weight = floatArrayOf(1f, 1f) // [1, 1, 2] ones + val graph = importSingleNode( + node, + inputs = listOf(valueInfo("x", 1, 1, 3)), + initializers = listOf(floatInitializer("w", weight, 1, 1, 2)), + ) + + assertEquals("convTranspose1d", graph.nodes.first { it.id == "ct1d" }.operation.name) + + val result = ComputeGraphExecutor(graph, ops).execute( + mapOf( + "x" to tensor(Shape(1, 1, 3), floatArrayOf(1f, 2f, 3f)), + "w" to tensor(Shape(1, 1, 2), weight), + ) + ).getValue("ct1d") + + assertEquals(Shape(1, 1, 4), result.shape) + assertValues(floatArrayOf(1f, 3f, 5f, 3f), result) + } + + @Test + fun convTranspose_autoPad_staysPlaceholder_andFailsLoudly() { + val node = NodeProto( + name = "ct3", + opType = "ConvTranspose", + input = listOf("x", "w"), + output = listOf("y"), + attribute = listOf( + AttributeProto(name = "auto_pad", s = ByteArr("SAME_UPPER".encodeToByteArray())), + ), + ) + val weight = floatArrayOf(1f, 1f, 1f, 1f) + val graph = importSingleNode( + node, + inputs = listOf(valueInfo("x", 1, 1, 2, 2)), + initializers = listOf(floatInitializer("w", weight, 1, 1, 2, 2)), + ) + + assertEquals("ConvTranspose", graph.nodes.first { it.id == "ct3" }.operation.name) + assertFailsWith { + ComputeGraphExecutor(graph, ops).execute( + mapOf( + "x" to tensor(Shape(1, 1, 2, 2), floatArrayOf(0f, 1f, 2f, 3f)), + "w" to tensor(Shape(1, 1, 2, 2), weight), + ) + ) + } + } + + @Test + fun convTranspose_asymmetricPads_staysPlaceholder() { + val node = NodeProto( + name = "ct4", + opType = "ConvTranspose", + input = listOf("x", "w"), + output = listOf("y"), + attribute = listOf(AttributeProto(name = "pads", ints = listOf(1L, 0L, 0L, 0L))), + ) + val graph = importSingleNode( + node, + inputs = listOf(valueInfo("x", 1, 1, 2, 2)), + initializers = listOf(floatInitializer("w", floatArrayOf(1f, 1f, 1f, 1f), 1, 1, 2, 2)), + ) + + assertEquals("ConvTranspose", graph.nodes.first { it.id == "ct4" }.operation.name) + } +}