Fix StableHLO MLIR export to produce valid IREE-compilable output

The HLO generator was producing broken MLIR because the graph builder
did not synthesize nodes for tensor inputs without known producers
(model inputs and weights). This adds a finalize() step to
TraceToGraphBuilder that creates synthetic "input" and "weight" nodes
for unresolved references, fixes dense literal formatting to produce
properly nested arrays matching tensor rank, and corrects the
stablehlo.convolution assembly syntax (attribute dict, functional type
annotation, batch_group_count).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: michalharakal

rgb2grayscale.mlir

@@ -0,0 +1,9 @@
+module {
+  func.func @rgb2grayscale(%arg0: tensor<1x3x4x4xf32>) -> (tensor<1x1x4x4xf32>) {
+    // input n1_input: t0 : tensor<1x3x4x4xf32>
+    // weight n2_weight: frozen parameter
+    %v0 = stablehlo.constant dense<[[[[0.2989]], [[0.587]], [[0.114]]]]> : tensor<1x3x1x1xf32>
+    %v1 = stablehlo.convolution(%arg0, %v0) dim_numbers = [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1], window = {stride = [1, 1], pad = [[0, 0], [0, 0]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x4xf32>, tensor<1x3x1x1xf32>) -> tensor<1x1x4x4xf32>
+    return %v1 : tensor<1x1x4x4xf32>
+  }
+}

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt

@@ -240,13 +240,19 @@ public open class DefaultExecutionTape(
         return newTape
     }
 
-    public fun toComputeGraph(): ComputeGraph {
+    public fun toComputeGraph(
+        synthesizeExternalInputs: Boolean = false,
+        inputTensorIds: Set<String> = emptySet()
+    ): ComputeGraph {
         // Prefer trace-based offline build when traces are available to ensure
         // consistency with online GraphSink wiring rules (PRD FR6).
         if (_traces.isNotEmpty()) {
             val graph = DefaultComputeGraph()
             val builder = TraceToGraphBuilder(graph, session)
             builder.addAll(_traces)
+            if (synthesizeExternalInputs) {
+                builder.finalize(inputTensorIds)
+            }
             return graph
         }
 

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/tape/extensions.kt

@@ -5,11 +5,20 @@ import sk.ainet.lang.graph.DefaultComputeGraph
 import sk.ainet.tape.ExecutionTape
 
 /**
- * Convert the tape to a compute graph
+ * Convert the tape to a compute graph.
+ *
+ * @param synthesizeExternalInputs When true, placeholder "input" and "weight" constant nodes are
+ *   created for tensor inputs that have no known producer in the trace. Required for StableHLO
+ *   compilation where every operand must be wired through graph edges.
+ * @param inputTensorIds Tensor IDs that should always become function arguments (model inputs)
+ *   rather than constants, even if their data is resolvable.
  */
-public fun ExecutionTape.toComputeGraph(): ComputeGraph {
+public fun ExecutionTape.toComputeGraph(
+    synthesizeExternalInputs: Boolean = false,
+    inputTensorIds: Set<String> = emptySet()
+): ComputeGraph {
     return when (this) {
-        is sk.ainet.lang.graph.DefaultExecutionTape -> this.toComputeGraph()
+        is sk.ainet.lang.graph.DefaultExecutionTape -> this.toComputeGraph(synthesizeExternalInputs, inputTensorIds)
         else -> DefaultComputeGraph()
     }
 }

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt

@@ -18,9 +18,10 @@ import sk.ainet.lang.tensor.ops.ValidationResult
  *   "e_<srcNodeId>_<srcOut>__<dstNodeId>_<dstIn>"
  *   Example: e_n0_Add_0__n1_Relu_0. This is deterministic given the node IDs and wiring.
  *
- * Note: This builder does NOT synthesize explicit "input" placeholder nodes for tensors without a
- * known producer. Only real operation nodes are created, and edges are added solely between such
- * nodes when a producer is known. This matches the expectations of TracingAcceptanceTest.
+ * By default this builder does NOT synthesize explicit "input" placeholder nodes for tensors
+ * without a known producer. Call [finalize] after adding all traces to synthesize "input" and
+ * "weight" constant nodes for unresolved external inputs. This is required for StableHLO
+ * compilation where every operand must be wired through graph edges.
  */
 public class TraceToGraphBuilder(
     private val graph: ComputeGraph,
@@ -32,6 +33,14 @@ public class TraceToGraphBuilder(
     private data class Producer(val node: GraphNode, val outIndex: Int, val spec: TensorSpec)
     private val producersByTensorId = mutableMapOf<String, Producer>()
 
+    private data class UnresolvedRef(
+        val tensorRef: TensorRef,
+        val consumerNode: GraphNode,
+        val inputIndex: Int,
+        val spec: TensorSpec
+    )
+    private val unresolvedByTensorId = mutableMapOf<String, MutableList<UnresolvedRef>>()
+
     /**
      * Add a single OpTrace into the graph, wiring known producers to inputs
      * and registering the outputs as new producers.
@@ -106,6 +115,35 @@ public class TraceToGraphBuilder(
                         }
                     }
                 }
+                "conv2d" -> {
+                    // For Conv2d layer: conv2d(input, weight, bias?)
+                    // Resolve weight tensor (second input) from session
+                    if (trace.inputs.size >= 2) {
+                        val weightRef = trace.inputs[1]
+                        if (!producersByTensorId.containsKey(weightRef.id)) {
+                            val tensor = session?.resolve(weightRef)
+                            if (tensor != null) {
+                                val values = extractFloatArray(tensor)
+                                if (values != null) {
+                                    parameters["weights"] = values
+                                }
+                            }
+                        }
+                    }
+                    // Resolve optional bias tensor (third input) from session
+                    if (trace.inputs.size >= 3) {
+                        val biasRef = trace.inputs[2]
+                        if (!producersByTensorId.containsKey(biasRef.id)) {
+                            val tensor = session?.resolve(biasRef)
+                            if (tensor != null) {
+                                val values = extractFloatArray(tensor)
+                                if (values != null) {
+                                    parameters["bias_values"] = values
+                                }
+                            }
+                        }
+                    }
+                }
             }
         }
 
@@ -114,8 +152,6 @@ public class TraceToGraphBuilder(
         val inputSpecs = buildInputSpecs(trace)
         val outputSpecs = buildOutputSpecs(trace)
 
-        // Do not synthesize placeholder input nodes; leave unknown producers unresolved.
-
         val nodeId = "n${nextNodeId++}_${trace.opType}"
         val node = GraphNode(
             id = nodeId,
@@ -125,7 +161,7 @@ public class TraceToGraphBuilder(
         )
         graph.addNode(node)
 
-        // Wire edges from producers
+        // Wire edges from producers; track unresolved inputs for later finalization
         trace.inputs.forEachIndexed { idx, tRef ->
             val prod = producersByTensorId[tRef.id]
             if (prod != null) {
@@ -141,6 +177,15 @@ public class TraceToGraphBuilder(
                         tensorSpec = tensorSpec
                     )
                 )
+            } else {
+                // Track for finalize() — no placeholder synthesized here by default
+                val spec = inputSpecs.getOrNull(idx) ?: TensorSpec(
+                    name = tRef.id,
+                    shape = tRef.shape.dimensions.toList(),
+                    dtype = tRef.dtype::class.simpleName ?: "FP32"
+                )
+                unresolvedByTensorId.getOrPut(tRef.id) { mutableListOf() }
+                    .add(UnresolvedRef(tRef, node, idx, spec))
             }
         }
 
@@ -159,6 +204,95 @@ public class TraceToGraphBuilder(
         traces.forEach { addTrace(it) }
     }
 
+    /**
+     * Synthesize placeholder nodes for tensor inputs that had no known producer.
+     *
+     * For each unresolved tensor:
+     * - If the tensor ID is in [inputTensorIds], an "input" placeholder node is created
+     *   (representing a function argument).
+     * - Else if the original tensor can be resolved from the session and contains constant data
+     *   (e.g. model weights), a "weight" constant node is created.
+     * - Otherwise an "input" placeholder node is created as a fallback.
+     *
+     * Edges are wired from the new nodes to every consumer that referenced the tensor.
+     * Call this after [addAll] when building graphs for compilation.
+     *
+     * @param inputTensorIds Tensor IDs that should always become function arguments (model inputs).
+     */
+    public fun finalize(inputTensorIds: Set<String> = emptySet()) {
+        for ((tensorId, refs) in unresolvedByTensorId) {
+            val firstRef = refs.first()
+            val spec = firstRef.spec
+
+            // If explicitly marked as model input, create an input node unconditionally
+            val forceInput = inputTensorIds.contains(tensorId)
+
+            // Try to resolve as a constant from the session
+            val tensor = if (!forceInput) session?.resolve(firstRef.tensorRef) else null
+            val constantValues = tensor?.let { extractFloatArray(it) }
+
+            val syntheticNode: GraphNode
+            if (constantValues != null) {
+                // Create a constant/weight node with embedded values
+                val weightShape = tensor!!.shape.dimensions.toList()
+                val weightDtype = tensor.dtype.simpleName ?: "FP32"
+                val nodeId = "n${nextNodeId++}_weight"
+                val op = TraceBackedOperation(
+                    name = "weight",
+                    type = "constant",
+                    parameters = mapOf(
+                        "initial_value" to constantValues.toList(),
+                        "trainable" to false
+                    )
+                )
+                syntheticNode = GraphNode(
+                    id = nodeId,
+                    operation = op,
+                    inputs = emptyList(),
+                    outputs = listOf(TensorSpec(
+                        name = tensorId,
+                        shape = weightShape,
+                        dtype = weightDtype
+                    ))
+                )
+            } else {
+                // Create an input placeholder node
+                val nodeId = "n${nextNodeId++}_input"
+                val op = TraceBackedOperation(
+                    name = "input",
+                    type = "input",
+                    parameters = emptyMap()
+                )
+                syntheticNode = GraphNode(
+                    id = nodeId,
+                    operation = op,
+                    inputs = emptyList(),
+                    outputs = listOf(spec)
+                )
+            }
+
+            graph.addNode(syntheticNode)
+
+            // Wire edges to all consumers
+            for (ref in refs) {
+                graph.addEdge(
+                    GraphEdge(
+                        id = "e_${syntheticNode.id}_0__${ref.consumerNode.id}_${ref.inputIndex}",
+                        source = syntheticNode,
+                        destination = ref.consumerNode,
+                        sourceOutputIndex = 0,
+                        destinationInputIndex = ref.inputIndex,
+                        tensorSpec = ref.spec
+                    )
+                )
+            }
+
+            // Register as producer
+            producersByTensorId[tensorId] = Producer(syntheticNode, 0, spec)
+        }
+        unresolvedByTensorId.clear()
+    }
+
     private fun buildInputSpecs(trace: OpTrace): List<TensorSpec> {
         val shapes = (trace.attributes["inputShapes"] as? List<*>)?.map { it as? List<Int> }
         val dtypes = (trace.attributes["inputDTypes"] as? List<*>)?.map { it?.toString() }

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/TypeMapper.kt

@@ -25,7 +25,7 @@ public class TypeMapper {
         "UI32", "UINT32" -> "ui32"
         "UI64", "UINT64" -> "ui64"
         "BF16", "BFLOAT16" -> "bf16"
-        "F16", "FLOAT16" -> "f16"
+        "FP16", "F16", "FLOAT16" -> "f16"
         "BOOL", "BOOLEAN" -> "i1"
         else -> {
             // Default fallback with warning

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ConstantOperationsConverter.kt

@@ -342,55 +342,54 @@ public class ConstantOperationsConverter : StableHloOperationConverter {
     }
 
     /**
-     * Format tensor values for MLIR dense constant
+     * Format tensor values for MLIR dense constant.
+     * MLIR dense<> syntax requires nested brackets matching the tensor rank:
+     *   scalar:  dense<42.0>
+     *   1D [3]:  dense<[v0, v1, v2]>
+     *   2D [2,3]: dense<[[v0,v1,v2],[v3,v4,v5]]>
+     *   4D [1,3,1,1]: dense<[[[[v0],[v1],[v2]]]]>
      */
     private fun formatTensorValues(values: List<*>, outputSpec: TensorSpec?): String {
         val shape = outputSpec?.shape ?: emptyList()
-        
+
         return when {
             values.isEmpty() -> "0.0"
-            shape.isEmpty() || shape.size == 1 -> {
-                // 1D tensor or scalar
-                values.joinToString(", ") { formatConstantValue(it as Number) }
-            }
-            shape.size == 2 -> {
-                // 2D tensor - format as nested arrays
-                formatAs2DTensor(values, shape)
-            }
-            else -> {
-                // Multi-dimensional tensor - flatten for now
-                values.joinToString(", ") { formatConstantValue(it as Number) }
-            }
+            values.size == 1 -> formatConstantValue(values[0] as Number)
+            shape.isEmpty() -> "[" + values.joinToString(", ") { formatConstantValue(it as Number) } + "]"
+            else -> formatNestedTensor(values, shape, 0, IntArray(1))
         }
     }
-    
+
     /**
-     * Format values as a 2D tensor for MLIR
+     * Recursively format a flat list of values into nested MLIR dense literal
+     * matching the given shape. [offset] tracks the current position in the flat values list.
      */
-    private fun formatAs2DTensor(values: List<*>, shape: List<Int>): String {
-        if (shape.size != 2) return values.joinToString(", ") { formatConstantValue(it as Number) }
-        
-        val rows = shape[0]
-        val cols = shape[1]
-        val result = StringBuilder()
-        
-        result.append("[")
-        for (i in 0 until rows) {
-            if (i > 0) result.append(", ")
-            result.append("[")
-            for (j in 0 until cols) {
-                if (j > 0) result.append(", ")
-                val index = i * cols + j
-                if (index < values.size) {
-                    result.append(formatConstantValue(values[index] as Number))
+    private fun formatNestedTensor(values: List<*>, shape: List<Int>, dim: Int, offset: IntArray): String {
+        if (dim == shape.size - 1) {
+            // Innermost dimension: emit a flat array of values
+            val size = shape[dim]
+            val sb = StringBuilder("[")
+            for (i in 0 until size) {
+                if (i > 0) sb.append(", ")
+                val idx = offset[0]++
+                if (idx < values.size) {
+                    sb.append(formatConstantValue(values[idx] as Number))
                 } else {
-                    result.append("0.0")
+                    sb.append("0.0")
                 }
             }
-            result.append("]")
+            sb.append("]")
+            return sb.toString()
         }
-        result.append("]")
-        
-        return result.toString()
+
+        // Non-innermost dimension: recurse
+        val size = shape[dim]
+        val sb = StringBuilder("[")
+        for (i in 0 until size) {
+            if (i > 0) sb.append(", ")
+            sb.append(formatNestedTensor(values, shape, dim + 1, offset))
+        }
+        sb.append("]")
+        return sb.toString()
     }
 }