diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a0c732f..a36f69f7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,58 @@ ## [Unreleased] +### Added + +- **First-class dynamic dimensions (`Dim`).** A new `sk.ainet.lang.tensor.Dim` vocabulary makes "dynamic + extent" explicit instead of an overloaded `-1`: `Dim.DYNAMIC` is a reserved sentinel (`Int.MIN_VALUE`) + **distinct from reshape's `-1` = infer**, with the dynamic-aware shape arithmetic (`concat`, `compatible`, + `render`, `isDynamic`/`isStatic`) centralized in one place rather than scattered `extent < 0` guards. + `Shape` gains `hasDynamic()`, `isDynamic(axis)`, `dynamicAxes`, and its `volume` now throws on a dynamic + shape (an unknown extent has no materializable element count) instead of returning a corrupt product. + The slice/range DSL is dynamic-aware too: `all()` over a dynamic axis is a valid symbolic full-axis, and + reshape passes a dynamic target through unchanged. The `skainet-compile-hlo` emitter now shares this one + sentinel (`TypeMapper.DYNAMIC_DIM` aliases `Dim.DYNAMIC`), so tracer and emitter agree by construction. +- **Dynamic-shape-safe StableHLO emission (streaming KV-cache decode).** The `skainet-compile-hlo` + emitter now renders a `-1` tensor extent as an MLIR `?` (dynamic) dimension and emits op forms that + `iree-compile` accepts under a dynamic dim, so a single compiled vmfb serves every autoregressive + decode step (growing KV cache) instead of one fixed cache length. A new `TypeMapper.DYNAMIC_DIM = -1` + marker plus a `List.hasDynamic()` predicate gate the dynamic paths, so every static graph is + emitted byte-for-byte unchanged. Verified: one dynamic SDPA vmfb runs at key lengths 3 and 17, and + the full FunctionGemma `with_past` decode graph (real weights, dynamic `1x{nKV}x?x256` cache) + self-compiles from the DSL to a CPU vmfb — a graph that could not be compiled before. +- **Dynamic-shape-safe trace finalization.** `TraceToGraphBuilder.extractFloatArray` no longer probes + `Tensor.volume` for non-dense data, so a dynamic-shaped graph input (e.g. a `?` KV-cache tensor) is left + as an input instead of tripping the (correctly) throwing `volume` on an unknown extent. This lets a decode + graph with dynamic caches finalize to a `ComputeGraph` and compile — verified with the real Moonshine v2 + `with_past` decoder (dynamic self *and* cross caches, `1x8x?x40`) self-compiling from the DSL to a vmfb. +- **Allocation-free shape-only tracing (`VoidTensorOps`).** The trace-time op set now propagates shapes + through a `ShapeOnlyTensorData` that carries a `Shape` but allocates no backing buffer, so a dynamic + (`-1`) extent flows through a whole decode trace instead of throwing `NegativeArraySizeException` when + a real buffer of negative size is allocated. This is what lets a real KV-cache seq dim be traced as + dynamic end-to-end (rather than via a sentinel-dimension + post-emit text substitution). + +### Changed + +- **SDPA scale folded into Q.** `AttentionOperationsConverter` now applies the attention scale as a + scalar constant multiplied into Q *before* the QK `dot_general` (`(q·s)@kᵀ ≡ scores·s`, exact), + instead of a dense splat constant sized to the full scores shape. This drops a scores-sized constant + from every attention graph and, crucially, avoids an invalid dynamic-shape splat when the key/cache + dim is `?`. +- **Softmax broadcasts are dynamic-shape-safe.** When the softmax/scores shape carries a dynamic dim, + `AttentionOperationsConverter` and `ActivationOperationsConverter` broadcast the reduced max/sum with + `stablehlo.dynamic_broadcast_in_dim` (runtime shape operand built via `get_dimension_size` + + `concatenate`) instead of a static `broadcast_in_dim`. Static graphs keep the explicit + `broadcast_in_dim` unchanged. +- **Identity reshapes/slices are elided.** `ShapeOperationsConverter` returns the operand SSA value with + no emitted op when a reshape's input and result types are identical, or a `slice`/`narrow` covers the + full extent of every axis (the KV-cache "cache-as-output-sink" and full-cache head-expansion patterns). + Besides being a no-op, this is the only valid lowering on a dynamic axis — a static `stablehlo.slice` + cannot express a full-extent bound on a `?` dim (its limit would be the `-1` extent, e.g. `0:-1:1`). +- **Concatenate propagates dynamic extents.** Both the trace-time shape inference + (`VoidTensorOps.calculateConcatShape`) and the emitter (`ShapeOperationsConverter` concat) now keep the + concatenated axis dynamic when any operand's extent there is dynamic, instead of numerically summing it + (which turned a growing cache `? ++ 1` into a bogus static `0`). + ## [0.37.0] - 2026-07-25 ### Added diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt index 1bca097b..832f254b 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt @@ -368,19 +368,9 @@ public class TraceToGraphBuilder( return buffer.copyOf() } - // Fallback for other data types if possible - if (tensor.volume > 0) { - val result = FloatArray(tensor.volume) - // This is slow but generic. Better if we have a way to get values. - // But usually weights are FloatArrayTensorData in the contexts we use for export. - return try { - // We don't have a good way to iterate over all indices generically without recursion - // for arbitrary rank. Let's stick to FloatArrayTensorData for now as it's the most common. - null - } catch (e: Exception) { - null - } - } + // Nothing else is materializable here: weights are FloatArrayTensorData in export contexts, and a + // dynamic-shaped tensor (e.g. a `?` KV-cache input) has no volume to probe — never call `.volume` + // on it (it throws by design). Such tensors are graph inputs, not constants to embed, so return null. return null } diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/TypeMapper.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/TypeMapper.kt index 84f3b268..1210cd2c 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/TypeMapper.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/TypeMapper.kt @@ -1,5 +1,6 @@ package sk.ainet.compile.hlo +import sk.ainet.lang.tensor.Dim import sk.ainet.lang.tensor.ops.TensorSpec /** @@ -117,24 +118,34 @@ public class TypeMapper { } /** - * Format tensor shape for MLIR + * Format tensor shape for MLIR. A negative extent renders as `?` (a dynamic dimension — see + * [DYNAMIC_DIM]); a `null` shape is a fully-dynamic tensor (`?`). */ - private fun formatShape(shape: List?): String { + public fun formatShape(shape: List?): String { return when { shape == null -> "?" shape.isEmpty() -> "" - else -> shape.joinToString("x") { if (it < 0) "?" else it.toString() } + else -> shape.joinToString("x") { Dim.render(it) } } } - + /** - * Create a tensor type string with explicit shape + * Create a tensor type string with explicit shape. Renders negative extents as `?` (dynamic). */ public fun createTensorType(shape: List, dtype: String): String { val elementType = mapDType(dtype) if (shape.isEmpty()) return "tensor<$elementType>" // rank-0 scalar - val shapeStr = shape.joinToString("x") - return "tensor<${shapeStr}x${elementType}>" + return "tensor<${formatShape(shape)}x${elementType}>" + } + + public companion object { + /** + * Sentinel extent meaning "dynamic dimension" (`?`) in a [TensorSpec] shape. Threaded from the trace + * (e.g. a KV-cache seq dim) so the emitter renders `?` and picks dynamic-shape-safe op forms, instead + * of the legacy post-emit text substitution. Aliases the canonical [Dim.DYNAMIC] (a reserved sentinel + * distinct from reshape's `-1` = infer), so tracer and emitter agree on one value. + */ + public const val DYNAMIC_DIM: Int = Dim.DYNAMIC } /** diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt index 4fd5c221..192657f6 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt @@ -3,6 +3,8 @@ package sk.ainet.compile.hlo.converters import sk.ainet.compile.hlo.ConversionContext import sk.ainet.compile.hlo.ConversionResult import sk.ainet.compile.hlo.StableHloOperationConverter +import sk.ainet.lang.tensor.Dim +import sk.ainet.lang.tensor.hasDynamic import sk.ainet.lang.graph.GraphNode /** @@ -139,8 +141,14 @@ public class ActivationOperationsConverter : StableHloOperationConverter { val reducedType = if (reducedShape.isEmpty()) { "tensor<$elementType>" } else { - "tensor<${reducedShape.joinToString("x")}x$elementType>" + "tensor<${reducedShape.joinToString("x") { Dim.render(it) }}x$elementType>" } + // Dynamic softmax axis / leading dims (`?`): the reduced max/sum must broadcast back to a dynamic + // output shape, which `stablehlo.broadcast_in_dim` cannot target — use `stablehlo.dynamic_broadcast_in_dim` + // with a runtime `output_dimensions` operand built from the input via `get_dimension_size` (IREE's + // stablehlo pipeline rejects CHLO implicit-broadcast ops as illegal). + val dyn = inputShape.hasDynamic() + val shapeType = "tensor<${rank}xi32>" // Dimensions kept for broadcast_in_dim: every input dim except `axis`, // mapped to its position in the reduced tensor. @@ -161,36 +169,54 @@ public class ActivationOperationsConverter : StableHloOperationConverter { // constant is out of range). val maxIdentity = context.getTypeMapper().negInfBits(elementType) - val operations = listOf( + val operations = buildList { // Reduce-max along the softmax axis (for numerical stability). - "$maxInit = stablehlo.constant dense<$maxIdentity> : tensor<$elementType>", - "$maxValue = stablehlo.reduce(${operands[0]} init: $maxInit) " + - "applies stablehlo.maximum across dimensions = [$axis] : " + - "($outputType, tensor<$elementType>) -> $reducedType", - - // Broadcast reduced max back to the input shape. - "$maxBroadcast = stablehlo.broadcast_in_dim $maxValue, " + - "dims = [$broadcastDims] : ($reducedType) -> $outputType", - - // Subtract the max for numerical stability. - "$shiftedValue = stablehlo.subtract ${operands[0]}, $maxBroadcast : $outputType", - - // Elementwise exponential. - "$expValue = stablehlo.exponential $shiftedValue : $outputType", - - // Reduce-sum along the softmax axis. - "$sumInit = stablehlo.constant dense<0.0> : tensor<$elementType>", - "$sumValue = stablehlo.reduce($expValue init: $sumInit) " + - "applies stablehlo.add across dimensions = [$axis] : " + - "($outputType, tensor<$elementType>) -> $reducedType", - - // Broadcast the sum back to the input shape. - "$sumBroadcast = stablehlo.broadcast_in_dim $sumValue, " + - "dims = [$broadcastDims] : ($reducedType) -> $outputType", - - // Normalize. - "$resultValue = stablehlo.divide $expValue, $sumBroadcast : $outputType" - ) + add("$maxInit = stablehlo.constant dense<$maxIdentity> : tensor<$elementType>") + add( + "$maxValue = stablehlo.reduce(${operands[0]} init: $maxInit) " + + "applies stablehlo.maximum across dimensions = [$axis] : " + + "($outputType, tensor<$elementType>) -> $reducedType", + ) + // Build the runtime output-shape operand once (only needed for dynamic broadcasts). + val shapeOperand: String = if (!dyn) "" else run { + val parts = inputShape.indices.map { d -> + if (Dim.isStatic(inputShape[d])) { + val c = context.nextTempValue() + add("$c = stablehlo.constant dense<${inputShape[d]}> : tensor<1xi32>") + c + } else { + val gd = context.nextTempValue(); val gr = context.nextTempValue() + add("$gd = stablehlo.get_dimension_size ${operands[0]}, dim = $d : ($outputType) -> tensor") + add("$gr = stablehlo.reshape $gd : (tensor) -> tensor<1xi32>") + gr + } + } + val sh = context.nextTempValue() + add("$sh = stablehlo.concatenate ${parts.joinToString(", ")}, dim = 0 : (${parts.joinToString(", ") { "tensor<1xi32>" }}) -> $shapeType") + sh + } + // Subtract the max: static → explicit broadcast_in_dim; dynamic → runtime dynamic_broadcast_in_dim. + if (dyn) { + add("$maxBroadcast = stablehlo.dynamic_broadcast_in_dim $maxValue, $shapeOperand, dims = [$broadcastDims] : ($reducedType, $shapeType) -> $outputType") + } else { + add("$maxBroadcast = stablehlo.broadcast_in_dim $maxValue, dims = [$broadcastDims] : ($reducedType) -> $outputType") + } + add("$shiftedValue = stablehlo.subtract ${operands[0]}, $maxBroadcast : $outputType") + add("$expValue = stablehlo.exponential $shiftedValue : $outputType") + add("$sumInit = stablehlo.constant dense<0.0> : tensor<$elementType>") + add( + "$sumValue = stablehlo.reduce($expValue init: $sumInit) " + + "applies stablehlo.add across dimensions = [$axis] : " + + "($outputType, tensor<$elementType>) -> $reducedType", + ) + // Normalize: static → broadcast_in_dim; dynamic → runtime dynamic_broadcast_in_dim. + if (dyn) { + add("$sumBroadcast = stablehlo.dynamic_broadcast_in_dim $sumValue, $shapeOperand, dims = [$broadcastDims] : ($reducedType, $shapeType) -> $outputType") + } else { + add("$sumBroadcast = stablehlo.broadcast_in_dim $sumValue, dims = [$broadcastDims] : ($reducedType) -> $outputType") + } + add("$resultValue = stablehlo.divide $expValue, $sumBroadcast : $outputType") + } operations.forEach { context.emitOperation(it) } diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt index 283b4f44..920992b3 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt @@ -3,6 +3,8 @@ package sk.ainet.compile.hlo.converters import sk.ainet.compile.hlo.ConversionContext import sk.ainet.compile.hlo.ConversionResult import sk.ainet.compile.hlo.StableHloOperationConverter +import sk.ainet.lang.tensor.Dim +import sk.ainet.lang.tensor.hasDynamic import sk.ainet.lang.graph.GraphNode import kotlin.math.sqrt @@ -56,7 +58,9 @@ public class AttentionOperationsConverter : StableHloOperationConverter { val outSpec = node.outputs.firstOrNull() val mapper = context.getTypeMapper() val elem = outSpec?.let { mapper.mapDType(it.dtype) } ?: "f32" - fun typeOf(shape: List): String = "tensor<${shape.joinToString("x")}x$elem>" + // Render a shape's dims, mapping a dynamic extent (DYNAMIC_DIM = -1) to `?`. + fun dims(shape: List): String = shape.joinToString("x") { Dim.render(it) } + fun typeOf(shape: List): String = "tensor<${dims(shape)}x$elem>" val qType = context.getValueType(operands[0]) ?: typeOf(qShape) val kType = context.getValueType(operands[1]) ?: typeOf(kShape) @@ -77,19 +81,20 @@ public class AttentionOperationsConverter : StableHloOperationConverter { val contractQK = rank - 1 // contract head_dim of Q and K val sdAxis = scoresShape.size - 1 // softmax over key length val reducedShape = scoresShape.dropLast(1) - val reducedType = if (reducedShape.isEmpty()) "tensor<$elem>" else "tensor<${reducedShape.joinToString("x")}x$elem>" + val reducedType = if (reducedShape.isEmpty()) "tensor<$elem>" else typeOf(reducedShape) val bcastDims = (scoresShape.indices).filter { it != sdAxis }.joinToString(", ") val contractAttn = scoresShape.size - 1 // attn key-length axis val contractV = rank - 2 // V key-length axis val causal = (node.operation.parameters["causal"] as? Boolean) ?: false val qAxis = rank - 2 // query position in scores [.., Sq, Sk] - val scoresI32Type = "tensor<${scoresShape.joinToString("x")}xi32>" - val scoresI1Type = "tensor<${scoresShape.joinToString("x")}xi1>" + val scoresI32Type = "tensor<${dims(scoresShape)}xi32>" + val scoresI1Type = "tensor<${dims(scoresShape)}xi1>" - val scores = context.nextTempValue() val scaleC = context.nextTempValue() - val scaled = context.nextTempValue() + val scaleB = context.nextTempValue() + val qScaled = context.nextTempValue() + val scores = context.nextTempValue() val maxInit = context.nextTempValue(); val maxV = context.nextTempValue(); val maxB = context.nextTempValue() val shifted = context.nextTempValue(); val expV = context.nextTempValue() val sumInit = context.nextTempValue(); val sumV = context.nextTempValue(); val sumB = context.nextTempValue() @@ -97,9 +102,15 @@ public class AttentionOperationsConverter : StableHloOperationConverter { val out = context.nextTempValue() val ops = mutableListOf( - "$scores = stablehlo.dot_general ${operands[0]}, ${operands[1]}, ${batchClause}contracting_dims = [$contractQK] x [$contractQK] : ($qType, $kType) -> $scoresType", - "$scaleC = stablehlo.constant dense<$scaleVal> : $scoresType", - "$scaled = stablehlo.multiply $scores, $scaleC : $scoresType", + // Scale is applied to Q *before* the QK dot — scores·s == (q·s)@kᵀ exactly. This keeps the scale + // constant at the STATIC query type (`[.., Sq, headDim]`) rather than a splat sized to the scores + // shape `[.., Sq, Sk]`, whose Sk (key/cache length) may be dynamic (`?`) in KV-cache decode — a + // dynamic-shape splat constant is invalid StableHLO. Also drops a full scores-sized dense constant + // from static graphs. + "$scaleC = stablehlo.constant dense<$scaleVal> : tensor<$elem>", + "$scaleB = stablehlo.broadcast_in_dim $scaleC, dims = [] : (tensor<$elem>) -> $qType", + "$qScaled = stablehlo.multiply ${operands[0]}, $scaleB : $qType", + "$scores = stablehlo.dot_general $qScaled, ${operands[1]}, ${batchClause}contracting_dims = [$contractQK] x [$contractQK] : ($qType, $kType) -> $scoresType", ) // Explicit additive mask (operands[3]) — e.g. a sliding-window+causal @@ -108,7 +119,7 @@ public class AttentionOperationsConverter : StableHloOperationConverter { // iota causal path. Broadcast (trailing-aligned) to the scores shape // and add. Without this the masked layers run UNMASKED (attend to // future tokens) — correct only at position 0. - var softmaxIn = scaled + var softmaxIn = scores // scores are already scaled (scale folded into Q above) val maskOperand = operands.getOrNull(3) if (maskOperand != null) { val maskShape = node.inputs.getOrNull(3)?.shape ?: scoresShape @@ -123,7 +134,7 @@ public class AttentionOperationsConverter : StableHloOperationConverter { mb } val masked = context.nextTempValue() - ops += "$masked = stablehlo.add $scaled, $maskBc : $scoresType" + ops += "$masked = stablehlo.add $scores, $maskBc : $scoresType" softmaxIn = masked } else if (causal) { val iotaQ = context.nextTempValue(); val iotaK = context.nextTempValue() @@ -139,19 +150,49 @@ public class AttentionOperationsConverter : StableHloOperationConverter { // masked-fill select, which can trip downstream greedy constant-folding. ops += "$ninf = stablehlo.constant dense<-1.000000e+30> : $scoresType" ops += "$maskAdd = stablehlo.select $keep, $zeros, $ninf : $scoresI1Type, $scoresType" - ops += "$masked = stablehlo.add $scaled, $maskAdd : $scoresType" + ops += "$masked = stablehlo.add $scores, $maskAdd : $scoresType" softmaxIn = masked } - // softmax(softmaxIn) over the key-length axis + // softmax(softmaxIn) over the key-length axis. When the scores shape is dynamic (the `?` key/cache dim + // of KV-cache decode), the reduced max/sum must broadcast back to the dynamic scores shape. A static + // `stablehlo.broadcast_in_dim` cannot target a dynamic shape, so we use `stablehlo.dynamic_broadcast_in_dim` + // with a runtime `output_dimensions` operand (built once from the scores tensor via `get_dimension_size`). + // Static graphs keep the original explicit `broadcast_in_dim` path (byte-for-byte unchanged). + val dyn = scoresShape.hasDynamic() + val shapeType = "tensor<${scoresShape.size}xi32>" + val scoresShapeOperand: String = if (!dyn) "" else run { + val parts = scoresShape.indices.map { d -> + if (Dim.isStatic(scoresShape[d])) { + val c = context.nextTempValue() + ops += "$c = stablehlo.constant dense<${scoresShape[d]}> : tensor<1xi32>" + c + } else { + val gd = context.nextTempValue(); val gr = context.nextTempValue() + ops += "$gd = stablehlo.get_dimension_size $scores, dim = $d : ($scoresType) -> tensor" + ops += "$gr = stablehlo.reshape $gd : (tensor) -> tensor<1xi32>" + gr + } + } + val sh = context.nextTempValue() + ops += "$sh = stablehlo.concatenate ${parts.joinToString(", ")}, dim = 0 : (${parts.joinToString(", ") { "tensor<1xi32>" }}) -> $shapeType" + sh + } + fun broadcastBack(src: String, dst: String) { + if (dyn) { + ops += "$dst = stablehlo.dynamic_broadcast_in_dim $src, $scoresShapeOperand, dims = [$bcastDims] : ($reducedType, $shapeType) -> $scoresType" + } else { + ops += "$dst = stablehlo.broadcast_in_dim $src, dims = [$bcastDims] : ($reducedType) -> $scoresType" + } + } ops += "$maxInit = stablehlo.constant dense<${mapper.negInfBits(elem)}> : tensor<$elem>" ops += "$maxV = stablehlo.reduce($softmaxIn init: $maxInit) applies stablehlo.maximum across dimensions = [$sdAxis] : ($scoresType, tensor<$elem>) -> $reducedType" - ops += "$maxB = stablehlo.broadcast_in_dim $maxV, dims = [$bcastDims] : ($reducedType) -> $scoresType" + broadcastBack(maxV, maxB) ops += "$shifted = stablehlo.subtract $softmaxIn, $maxB : $scoresType" ops += "$expV = stablehlo.exponential $shifted : $scoresType" ops += "$sumInit = stablehlo.constant dense<0.0> : tensor<$elem>" ops += "$sumV = stablehlo.reduce($expV init: $sumInit) applies stablehlo.add across dimensions = [$sdAxis] : ($scoresType, tensor<$elem>) -> $reducedType" - ops += "$sumB = stablehlo.broadcast_in_dim $sumV, dims = [$bcastDims] : ($reducedType) -> $scoresType" + broadcastBack(sumV, sumB) ops += "$attn = stablehlo.divide $expV, $sumB : $scoresType" ops += "$out = stablehlo.dot_general $attn, ${operands[2]}, ${batchClause}contracting_dims = [$contractAttn] x [$contractV] : ($scoresType, $vType) -> $outputType" ops.forEach { context.emitOperation(it) } diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ShapeOperationsConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ShapeOperationsConverter.kt index a86f639f..b9abea14 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ShapeOperationsConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ShapeOperationsConverter.kt @@ -4,6 +4,7 @@ import sk.ainet.compile.hlo.ConversionContext import sk.ainet.compile.hlo.ConversionResult import sk.ainet.compile.hlo.StableHloOperationConverter import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.tensor.Dim import sk.ainet.lang.tensor.ops.TensorSpec /** @@ -108,7 +109,10 @@ public class ShapeOperationsConverter : StableHloOperationConverter { axis in inShapes[0].indices ) { val outShape = inShapes[0].toMutableList() - outShape[axis] = inShapes.sumOf { it[axis] } + // A dynamic extent on the concat axis of ANY operand makes the concatenated extent dynamic + // too — summing it would emit a bogus static dim (e.g. `? + 1` → `0`, an invalid + // `tensor<…x0x…>`). [Dim.concat] matches VoidTensorOps.calculateConcatShape. + outShape[axis] = Dim.concat(inShapes.map { it[axis] }) context.getTypeMapper().mapTensorType( TensorSpec("${node.id}_out", outShape, outputSpec?.dtype ?: node.inputs[0].dtype), ) @@ -166,6 +170,18 @@ public class ShapeOperationsConverter : StableHloOperationConverter { val strides = (node.operation.parameters["strides"] as? List) ?: List(rank) { 1 } + // Elide a full-range (identity) slice: start 0, stride 1, and limit == the full extent on EVERY + // axis. This is a no-op copy, and crucially a `stablehlo.slice` cannot express a full-extent bound + // on a dynamic axis (the limit would be the `-1`/`?` extent, e.g. `0:-1:1`, which is invalid). + // Mirrors the identity-reshape elision (the decode cache is sliced full, then head-expanded). + val isIdentitySlice = rank > 0 && + starts.size == rank && limits.size == rank && strides.size == rank && + starts.all { it == 0 } && strides.all { it == 1 } && + (0 until rank).all { limits[it] == inputShape[it] } + if (isIdentitySlice) { + return ConversionResult.Success(outputValueName = operands[0], emittedOperations = emptyList()) + } + val resultValue = context.nextTempValue() val operation = sliceLine(resultValue, operands[0], starts, limits, strides, resolveOperandType(operands[0], node, context), outputType) @@ -222,6 +238,13 @@ public class ShapeOperationsConverter : StableHloOperationConverter { val limits = List(rank) { if (it == dim) start + length else inputShape[it] } val strides = List(rank) { 1 } + // Full-extent narrow on the target axis (start 0, length == the axis extent) is a no-op copy; + // elide it. On a dynamic axis the limit would be `start + (-1)` — invalid as a static bound — so + // eliding is both an optimization and the only valid lowering. Mirrors the identity slice/reshape. + if (start == 0 && length == inputShape[dim]) { + return ConversionResult.Success(outputValueName = operands[0], emittedOperations = emptyList()) + } + val resultValue = context.nextTempValue() val operation = sliceLine(resultValue, operands[0], starts, limits, strides, resolveOperandType(operands[0], node, context), outputType) @@ -338,6 +361,19 @@ public class ShapeOperationsConverter : StableHloOperationConverter { ) val inputType = resolveOperandType(operands[0], node, context) + + // Identity reshape (input type == output type): elide it. The KV-cache decode graphs route a cache + // tensor through `reshape(x, x.shape)` purely to make it a distinct graph-output sink; emitting + // `stablehlo.reshape` with an unchanged result is a no-op, and is outright INVALID when the result + // carries a dynamic dim (`?`) — `stablehlo.reshape` requires a statically-shaped result. Pass the + // operand through: this node's SSA name resolves to it (and `func.return` may list it more than once). + if (inputType == resultType) { + return ConversionResult.Success( + outputValueName = operands[0], + emittedOperations = emptyList(), + ) + } + val resultValue = context.nextTempValue() val operation = "$resultValue = stablehlo.reshape ${operands[0]} : ($inputType) -> $resultType" context.emitOperation(operation) diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/DynamicShapeHloExportTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/DynamicShapeHloExportTest.kt new file mode 100644 index 00000000..54a8d155 --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/DynamicShapeHloExportTest.kt @@ -0,0 +1,130 @@ +package sk.ainet.compile.hlo + +import sk.ainet.lang.graph.DefaultComputeGraph +import sk.ainet.lang.graph.GraphEdge +import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.tensor.ops.InputOperation +import sk.ainet.lang.tensor.ops.Operation +import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.tensor.ops.ValidationResult +import sk.ainet.lang.types.DType +import sk.ainet.lang.tensor.Tensor +import kotlin.test.Test +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +/** + * Verifies the emitter is dynamic-shape-safe for KV-cache decode: when a tensor dim is dynamic + * ([TypeMapper.DYNAMIC_DIM] = -1), the converters must emit forms `iree-compile` accepts under `?` — + * no splat constant sized to a dynamic dim, no static `broadcast_in_dim`/`reshape` to a dynamic result. + * See the "make the emitter dynamic-shape-safe" change (scale-on-Q, dynamic_broadcast_in_dim softmax broadcast, + * identity-reshape elision). The `-1` renders as `?` via `TypeMapper.formatShape`. + */ +class DynamicShapeHloExportTest { + + private fun op(opName: String, params: Map = emptyMap()): Operation = object : Operation { + override val name = opName + override val type = "test" + override val parameters = params + override fun execute(inputs: List>) = + throw UnsupportedOperationException("conversion-only test op") + override fun validateInputs(inputs: List) = ValidationResult.Valid + override fun inferOutputs(inputs: List) = inputs + override fun clone(newParameters: Map) = op(opName, newParameters) + override fun serialize() = mapOf("name" to name, "type" to type) + } + + @Test + fun sdpa_with_dynamic_key_dim_emits_dynamic_safe_hlo() { + // Decode-shaped SDPA: query seq = 1 (static), key/value seq = DYNAMIC (-1 → `?`), the growing KV cache. + val q = TensorSpec("q", listOf(1, 8, 1, 40), "FP32") + val k = TensorSpec("k", listOf(1, 8, TypeMapper.DYNAMIC_DIM, 40), "FP32") + val v = TensorSpec("v", listOf(1, 8, TypeMapper.DYNAMIC_DIM, 40), "FP32") + val out = TensorSpec("out", listOf(1, 8, 1, 40), "FP32") + + val g = DefaultComputeGraph() + val nq = GraphNode("q", InputOperation(), emptyList(), listOf(q)) + val nk = GraphNode("k", InputOperation(), emptyList(), listOf(k)) + val nv = GraphNode("v", InputOperation(), emptyList(), listOf(v)) + val sdpa = GraphNode("sdpa", op("scaledDotProductAttention", mapOf("causal" to false)), + listOf(q, k, v), listOf(out)) + listOf(nq, nk, nv, sdpa).forEach { g.addNode(it) } + g.addEdge(GraphEdge("e0", nq, sdpa, 0, 0, q)) + g.addEdge(GraphEdge("e1", nk, sdpa, 0, 1, k)) + g.addEdge(GraphEdge("e2", nv, sdpa, 0, 2, v)) + + val mlir = StableHloConverterFactory.createBasic().convert(g, "dyn_sdpa").content + + assertTrue(mlir.contains("x?x"), "dynamic key dim must render as `?`:\n$mlir") + // No splat constant sized to a dynamic tensor (the scale bug): a `dense : tensor<…?…>`. + assertFalse( + mlir.contains(Regex("""stablehlo\.constant dense<[^>\[]*> : tensor<[^>]*\?[^>]*>""")), + "must not emit a dynamic-shape splat constant:\n$mlir", + ) + // Softmax broadcasts over the dynamic dim go through a runtime dynamic_broadcast_in_dim + // (IREE's stablehlo pipeline rejects CHLO implicit-broadcast ops), and no static broadcast_in_dim + // may target the dynamic scores shape. + assertTrue(mlir.contains("stablehlo.dynamic_broadcast_in_dim"), "dynamic softmax must use dynamic_broadcast_in_dim:\n$mlir") + assertTrue(mlir.contains("stablehlo.get_dimension_size"), "dynamic broadcast needs a runtime shape operand:\n$mlir") + assertFalse(mlir.contains("chlo."), "must not emit CHLO ops (illegal in IREE's stablehlo pipeline):\n$mlir") + // Scale folded into Q: a scalar scale const + a multiply feeding the first dot_general. + assertTrue( + mlir.contains(Regex("""stablehlo\.constant dense<[^>\[]*> : tensor""")) && + mlir.contains("stablehlo.multiply"), + "scale must be a scalar folded into Q (multiply), not a scores-sized splat:\n$mlir", + ) + } + + @Test + fun identity_reshape_on_dynamic_tensor_is_elided() { + // The decode cache-as-output-sink trick: reshape(x, x.shape) with x carrying a dynamic dim. + val t = TensorSpec("x", listOf(1, 8, TypeMapper.DYNAMIC_DIM, 40), "FP32") + val g = DefaultComputeGraph() + val nx = GraphNode("x", InputOperation(), emptyList(), listOf(t)) + val nr = GraphNode("rs", op("reshape", mapOf("outputShape" to listOf(1, 8, TypeMapper.DYNAMIC_DIM, 40))), + listOf(t), listOf(TensorSpec("rs_out", listOf(1, 8, TypeMapper.DYNAMIC_DIM, 40), "FP32"))) + g.addNode(nx); g.addNode(nr) + g.addEdge(GraphEdge("e0", nx, nr, 0, 0, t)) + + val mlir = StableHloConverterFactory.createBasic().convert(g, "id_reshape").content + assertFalse(mlir.contains("stablehlo.reshape"), "identity reshape must be elided, not emitted:\n$mlir") + } + + @Test + fun concat_of_dynamic_cache_stays_dynamic_not_zero() { + // The growing KV cache: concat(past[..,?,..], step[..,1,..]) along the seq axis must stay `?`, + // never `? + 1 = 0` (a bogus static `tensor<…x0x…>` that iree-compile rejects). + val past = TensorSpec("past", listOf(1, 4, TypeMapper.DYNAMIC_DIM, 256), "FP32") + val step = TensorSpec("step", listOf(1, 4, 1, 256), "FP32") + val outc = TensorSpec("full", listOf(1, 4, TypeMapper.DYNAMIC_DIM, 256), "FP32") + + val g = DefaultComputeGraph() + val np = GraphNode("past", InputOperation(), emptyList(), listOf(past)) + val ns = GraphNode("step", InputOperation(), emptyList(), listOf(step)) + val nc = GraphNode("cat", op("concat", mapOf("dim" to 2)), listOf(past, step), listOf(outc)) + listOf(np, ns, nc).forEach { g.addNode(it) } + g.addEdge(GraphEdge("e0", np, nc, 0, 0, past)) + g.addEdge(GraphEdge("e1", ns, nc, 0, 1, step)) + + val mlir = StableHloConverterFactory.createBasic().convert(g, "dyn_concat").content + assertTrue(mlir.contains("stablehlo.concatenate"), "concat must emit:\n$mlir") + assertTrue(mlir.contains("x?x256"), "concatenated seq axis must stay dynamic `?`:\n$mlir") + assertFalse(mlir.contains("x0x256"), "dynamic `? ++ 1` must NOT collapse to a static 0 dim:\n$mlir") + } + + @Test + fun full_extent_narrow_on_dynamic_axis_is_elided() { + // Full-cache head-expansion slices the whole (dynamic) seq axis; a static stablehlo.slice cannot + // express a full-extent bound on `?` (limit would be the `-1`-extent), so it must be elided. + val t = TensorSpec("x", listOf(1, 4, TypeMapper.DYNAMIC_DIM, 256), "FP32") + val g = DefaultComputeGraph() + val nx = GraphNode("x", InputOperation(), emptyList(), listOf(t)) + val nn = GraphNode("nw", op("narrow", mapOf("dim" to 2, "start" to 0, "length" to TypeMapper.DYNAMIC_DIM)), + listOf(t), listOf(TensorSpec("nw_out", listOf(1, 4, TypeMapper.DYNAMIC_DIM, 256), "FP32"))) + g.addNode(nx); g.addNode(nn) + g.addEdge(GraphEdge("e0", nx, nn, 0, 0, t)) + + val mlir = StableHloConverterFactory.createBasic().convert(g, "dyn_narrow").content + assertFalse(mlir.contains("stablehlo.slice"), "full-extent narrow on a dynamic axis must be elided:\n$mlir") + } +} diff --git a/skainet-compile/skainet-compile-hlo/src/jvmTest/kotlin/sk/ainet/compile/hlo/DynamicShapeHloDumpTest.kt b/skainet-compile/skainet-compile-hlo/src/jvmTest/kotlin/sk/ainet/compile/hlo/DynamicShapeHloDumpTest.kt new file mode 100644 index 00000000..a5bd841f --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/jvmTest/kotlin/sk/ainet/compile/hlo/DynamicShapeHloDumpTest.kt @@ -0,0 +1,49 @@ +package sk.ainet.compile.hlo + +import sk.ainet.lang.graph.DefaultComputeGraph +import sk.ainet.lang.graph.GraphEdge +import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.ops.InputOperation +import sk.ainet.lang.tensor.ops.Operation +import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.tensor.ops.ValidationResult +import sk.ainet.lang.types.DType +import java.io.File +import kotlin.test.Test + +/** Dumps the emitted dynamic-key SDPA StableHLO to `DYN_MLIR_OUT` (if set) so it can be fed to `iree-compile` + * to confirm the dynamic-shape-safe emission actually compiles. No-op assertion when the env var is unset. */ +class DynamicShapeHloDumpTest { + private fun op(opName: String, params: Map = emptyMap()): Operation = object : Operation { + override val name = opName + override val type = "test" + override val parameters = params + override fun execute(inputs: List>) = throw UnsupportedOperationException() + override fun validateInputs(inputs: List) = ValidationResult.Valid + override fun inferOutputs(inputs: List) = inputs + override fun clone(newParameters: Map) = op(opName, newParameters) + override fun serialize() = mapOf("name" to name) + } + + @Test + fun dumpDynamicSdpa() { + val out = System.getenv("DYN_MLIR_OUT") ?: return + val q = TensorSpec("q", listOf(1, 8, 1, 40), "FP32") + val k = TensorSpec("k", listOf(1, 8, TypeMapper.DYNAMIC_DIM, 40), "FP32") + val v = TensorSpec("v", listOf(1, 8, TypeMapper.DYNAMIC_DIM, 40), "FP32") + val o = TensorSpec("out", listOf(1, 8, 1, 40), "FP32") + val g = DefaultComputeGraph() + val nq = GraphNode("q", InputOperation(), emptyList(), listOf(q)) + val nk = GraphNode("k", InputOperation(), emptyList(), listOf(k)) + val nv = GraphNode("v", InputOperation(), emptyList(), listOf(v)) + val sdpa = GraphNode("sdpa", op("scaledDotProductAttention", mapOf("causal" to false)), + listOf(q, k, v), listOf(o)) + listOf(nq, nk, nv, sdpa).forEach { g.addNode(it) } + g.addEdge(GraphEdge("e0", nq, sdpa, 0, 0, q)) + g.addEdge(GraphEdge("e1", nk, sdpa, 0, 1, k)) + g.addEdge(GraphEdge("e2", nv, sdpa, 0, 2, v)) + File(out).writeText(StableHloConverterFactory.createBasic().convert(g, "dyn_sdpa").content) + println("WROTE_MLIR $out") + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/Dim.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/Dim.kt new file mode 100644 index 00000000..bd8d4ad6 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/Dim.kt @@ -0,0 +1,50 @@ +package sk.ainet.lang.tensor + +/** + * Vocabulary for a single tensor-dimension extent, and the canonical home of the [DYNAMIC] marker. + * + * A [Shape] stores extents as plain `Int`s (kept for storage/backend compatibility). An extent is either: + * - a concrete, materializable size (`>= 0`), or + * - [DYNAMIC] — a size unknown at compile time (e.g. a growing KV-cache sequence length). It renders as + * `?` in MLIR, and a single compiled program then serves every concrete size at that axis. + * + * [DYNAMIC] is a RESERVED sentinel ([Int.MIN_VALUE]), deliberately distinct from `-1`, which the reshape + * DSL uses for "infer this dimension from the total volume". Keeping the two values distinct removes the + * historical overloading where a dynamic axis and a reshape-infer slot were both `-1` and could collide. + * + * All dynamic-aware shape arithmetic lives here so ops don't scatter ad-hoc `extent < 0` guards. + */ +public object Dim { + /** Reserved sentinel extent meaning "dynamic / unknown size". Distinct from reshape's `-1` = infer. */ + public const val DYNAMIC: Int = Int.MIN_VALUE + + /** True iff [extent] is the [DYNAMIC] sentinel. */ + public fun isDynamic(extent: Int): Boolean = extent == DYNAMIC + + /** True iff [extent] is a concrete, materializable size (`>= 0`). */ + public fun isStatic(extent: Int): Boolean = extent >= 0 + + /** + * Extent of concatenating many tensors along one axis: [DYNAMIC] if ANY input is dynamic there, else + * the sum. (A growing cache `? ++ n` must stay `?`; numerically summing it would corrupt the shape, + * e.g. `? + 1` collapsing to a bogus static `0`.) + */ + public fun concat(extents: List): Int = + if (extents.any { isDynamic(it) }) DYNAMIC else extents.sum() + + /** + * Are two extents compatible for an elementwise / broadcast op — equal, or either side dynamic? + * A dynamic axis is compatible with any concrete size (the concrete one wins as the known extent). + */ + public fun compatible(a: Int, b: Int): Boolean = isDynamic(a) || isDynamic(b) || a == b + + /** Render an extent for MLIR: [DYNAMIC] (or any non-concrete value) as `?`, else the decimal size. */ + public fun render(extent: Int): String = if (isStatic(extent)) extent.toString() else "?" +} + +/** True if any extent in this dimension list is [Dim.DYNAMIC]. Lets both the tracer and the emitter pick + * dynamic-shape-safe forms only when needed, leaving fully-static shapes on their original code paths. */ +public fun List.hasDynamic(): Boolean = any { Dim.isDynamic(it) } + +/** True if any extent is [Dim.DYNAMIC]. */ +public fun IntArray.hasDynamic(): Boolean = any { Dim.isDynamic(it) } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/Shape.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/Shape.kt index b1a93894..6b0b807a 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/Shape.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/Shape.kt @@ -8,11 +8,27 @@ public data class Shape(val dimensions: IntArray) { } val volume: Int - get() = dimensions.fold(1) { a, x -> a * x } + get() { + require(!dimensions.hasDynamic()) { + "volume is undefined for a dynamic shape (${dimensions.joinToString(" x ", "[", "]") { Dim.render(it) }}); " + + "a dynamic extent has no materializable element count" + } + return dimensions.fold(1) { a, x -> a * x } + } val rank: Int get() = dimensions.size + /** True if any extent is [Dim.DYNAMIC] (unknown at compile time). */ + public fun hasDynamic(): Boolean = dimensions.hasDynamic() + + /** True if the extent on [axis] is [Dim.DYNAMIC]. */ + public fun isDynamic(axis: Int): Boolean = Dim.isDynamic(dimensions[axis]) + + /** Indices of every dynamic axis (empty for a fully-static shape). */ + public val dynamicAxes: List + get() = dimensions.indices.filter { Dim.isDynamic(dimensions[it]) } + public fun index(indices: IntArray): Int { assert( { indices.size == dimensions.size }, @@ -41,10 +57,11 @@ public data class Shape(val dimensions: IntArray) { } override fun toString(): String { - // Create a string representation of the dimensions array - val dimensionsString = dimensions.joinToString(separator = " x ", prefix = "[", postfix = "]") - // Return the formatted string including dimensions and volume - return "Shape: Dimensions = $dimensionsString, Size (Volume) = $volume" + // Render each extent via Dim (a dynamic extent shows as `?`), and omit the volume when it is + // undefined (any dynamic extent) rather than computing a corrupt product. + val dimensionsString = dimensions.joinToString(separator = " x ", prefix = "[", postfix = "]") { Dim.render(it) } + val volumeString = if (dimensions.hasDynamic()) "dynamic" else volume.toString() + return "Shape: Dimensions = $dimensionsString, Size (Volume) = $volumeString" } } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/Slice.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/Slice.kt index 38a38317..5b0ef65e 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/Slice.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/Slice.kt @@ -130,11 +130,25 @@ public sealed class Slice { * @return true if the slice is valid for the given dimension size */ public fun isValid(dimensionSize: Int): Boolean { - require(dimensionSize >= 0) { "Dimension size must be non-negative: $dimensionSize" } - + require(dimensionSize >= 0 || Dim.isDynamic(dimensionSize)) { + "Dimension size must be non-negative or dynamic: $dimensionSize" + } + + // Dynamic (unknown-extent) axis: `all()` is the symbolic full axis and is always valid; a partial + // Range/At/Step can't be bounds-checked against an unknown size, so only its own self-consistent, + // non-negative (never from-end) bounds are validated. + if (Dim.isDynamic(dimensionSize)) { + return when (this) { + is Range -> start >= 0 && end >= start + is At -> index >= 0 + is All -> true + is Step -> start >= 0 && end >= start && step > 0 + } + } + return when (this) { is Range -> start < dimensionSize && end <= dimensionSize - is At -> index < dimensionSize + is At -> index < dimensionSize is All -> true // Always valid is Step -> start < dimensionSize && end <= dimensionSize } @@ -218,8 +232,21 @@ public sealed class Slice { * @return a normalized slice equivalent to this slice */ public fun normalize(dimensionSize: Int): Slice { - require(dimensionSize >= 0) { "Dimension size must be non-negative: $dimensionSize" } - + require(dimensionSize >= 0 || Dim.isDynamic(dimensionSize)) { + "Dimension size must be non-negative or dynamic: $dimensionSize" + } + + // A dynamic axis has no known size to resolve from-end (negative) indices against, so those are + // rejected; the full-axis `all()` stays symbolic and explicit non-negative bounds pass through. + if (Dim.isDynamic(dimensionSize)) { + return when (this) { + is All -> this + is Range -> { require(start >= 0 && end >= 0) { "dynamic-axis Range needs non-negative bounds: $this" }; this } + is At -> { require(index >= 0) { "dynamic-axis At needs a non-negative index: $this" }; this } + is Step -> { require(start >= 0 && end >= 0) { "dynamic-axis Step needs non-negative bounds: $this" }; this } + } + } + return when (this) { is Range -> { val normStart = if (start < 0) maxOf(0, dimensionSize + start) else minOf(start, dimensionSize) diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt index 44f4467e..caa5a718 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt @@ -2,17 +2,26 @@ package sk.ainet.lang.tensor.ops import sk.ainet.lang.ops.Backend import sk.ainet.lang.ops.InProgress +import sk.ainet.lang.tensor.Dim import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.VoidOpsTensor +import sk.ainet.lang.tensor.hasDynamic import sk.ainet.lang.tensor.data.DenseTensorDataFactory +import sk.ainet.lang.tensor.data.TensorData import sk.ainet.lang.types.DType import sk.ainet.lang.tensor.data.views.UnsqueezedTensorData +import kotlin.reflect.KClass @Backend(id = "void", displayName = "Shape-only", internal = true) public class VoidTensorOps : TensorOps { - - private val dataFactory = DenseTensorDataFactory() + + // Shape-only tracing. A STATIC shape still gets a real (readable) zeros buffer — existing code + // creates void tensors and reads their zeros, so that behavior must be preserved. A DYNAMIC shape + // (a `Dim.DYNAMIC` extent) cannot be allocated at all (it would throw NegativeArraySizeException), + // so it gets an allocation-free ShapeOnlyTensorData that carries only the Shape — which is exactly + // what lets a dynamic KV-cache seq dim thread through a decode trace. See ShapeOnlyDataFactory. + private val dataFactory = ShapeOnlyDataFactory /** * Validates that two shapes are compatible for element-wise operations. @@ -694,8 +703,17 @@ public class VoidTensorOps : TensorOps { * Validates that total volume remains the same and no illegal dimensions are provided. */ private fun calculateReshapeTargetShape(originalShape: Shape, target: Shape): Shape { - val total = originalShape.volume val dims = target.dimensions + // A dynamic extent (either in the target or the input) passes through unchanged: volume-based + // inference/validation is undefined when an extent is unknown. This is distinct from `-1` = infer, + // which the distinct [Dim.DYNAMIC] sentinel keeps separate. + if (dims.hasDynamic() || originalShape.dimensions.hasDynamic()) { + require(dims.none { it == -1 }) { + "reshape cannot infer a `-1` dimension while the shape is dynamic: ${dims.toList()}" + } + return Shape(dims.copyOf()) + } + val total = originalShape.volume var inferIndex = -1 var product = 1 for ((i, d) in dims.withIndex()) { @@ -933,7 +951,8 @@ public class VoidTensorOps : TensorOps { throw IllegalArgumentException("All tensors must have the same number of dimensions for concatenation") } for (i in shape.dimensions.indices) { - if (i != actualDim && shape.dimensions[i] != firstShape.dimensions[i]) { + // Off-axis extents must match; a dynamic extent is compatible with any concrete size. + if (i != actualDim && !Dim.compatible(shape.dimensions[i], firstShape.dimensions[i])) { throw IllegalArgumentException( "All tensors must have the same shape except in the concatenation dimension. " + "Dimension $i: ${firstShape.dimensions[i]} vs ${shape.dimensions[i]}" @@ -942,10 +961,13 @@ public class VoidTensorOps : TensorOps { } } - // Calculate result shape + // Calculate result shape. [Dim.concat] keeps the concatenated axis dynamic when any input is + // dynamic there (a growing KV cache `? ++ 1` stays `?`) instead of numerically summing it, which + // would corrupt the shape — exactly what blocked threading a real dynamic seq dim through a decode + // trace. val resultDims = firstShape.dimensions.copyOf() - resultDims[actualDim] = shapes.sumOf { it.dimensions[actualDim] } - + resultDims[actualDim] = Dim.concat(shapes.map { it.dimensions[actualDim] }) + return Shape(resultDims) } @@ -1066,3 +1088,25 @@ public class VoidTensorOps : TensorOps { return actualDim } } + +/** + * A [TensorData] that carries only a [Shape] and allocates NO backing buffer. Used by + * [VoidTensorOps] for shape-only tracing: element access is never valid (nothing to read/write), + * but crucially the shape may contain a dynamic extent (`-1`) that a real allocation would reject. + */ +private class ShapeOnlyTensorData(override val shape: Shape) : TensorData { + private fun noData(): Nothing = + error("shape-only (void) tensor carries no data — tracing propagates shapes only") + override fun get(vararg indices: Int): V = noData() + override fun set(vararg indices: Int, value: V): Unit = noData() + override fun copyToFloatArray(): FloatArray = noData() +} + +/** Drop-in for the one `dataFactory.zeros` call VoidTensorOps makes. A static shape delegates to + * [DenseTensorDataFactory] (real, readable zeros — preserves existing behavior); only a DYNAMIC shape, + * which cannot be allocated, gets the allocation-free [ShapeOnlyTensorData] so its `-1` extent survives. */ +private object ShapeOnlyDataFactory { + private val dense = DenseTensorDataFactory() + fun zeros(shape: Shape, dtype: KClass): TensorData = + if (shape.hasDynamic()) ShapeOnlyTensorData(shape) else dense.zeros(shape, dtype) +} diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/DimTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/DimTest.kt new file mode 100644 index 00000000..952dde5a --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/tensor/DimTest.kt @@ -0,0 +1,96 @@ +package sk.ainet.lang.tensor + +import sk.ainet.lang.tensor.ops.VoidTensorOps +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +/** + * The first-class dynamic-dimension vocabulary: [Dim.DYNAMIC] is a reserved sentinel distinct from + * reshape's `-1` = infer, dynamic-aware shape arithmetic lives in [Dim], and [Shape] / the slice DSL / + * the shape-only tracer all treat a dynamic extent explicitly instead of doing integer math on `-1`. + */ +class DimTest { + + @Test + fun dynamic_sentinel_is_distinct_from_reshape_infer() { + assertTrue(Dim.isDynamic(Dim.DYNAMIC)) + assertFalse(Dim.isDynamic(-1), "reshape's -1 (infer) must NOT read as dynamic") + assertFalse(Dim.isDynamic(0)) + assertTrue(Dim.isStatic(0)) + assertFalse(Dim.isStatic(Dim.DYNAMIC)) + } + + @Test + fun dim_arithmetic_keeps_dynamic() { + assertEquals(6, Dim.concat(listOf(2, 4))) + assertEquals(Dim.DYNAMIC, Dim.concat(listOf(Dim.DYNAMIC, 1)), "? ++ 1 must stay ? (not 0)") + assertEquals(Dim.DYNAMIC, Dim.concat(listOf(3, Dim.DYNAMIC, 5))) + assertTrue(Dim.compatible(Dim.DYNAMIC, 7), "dynamic is compatible with any concrete size") + assertTrue(Dim.compatible(7, 7)) + assertFalse(Dim.compatible(7, 8)) + assertEquals("?", Dim.render(Dim.DYNAMIC)) + assertEquals("40", Dim.render(40)) + } + + @Test + fun shape_reports_dynamic_axes_and_guards_volume() { + val s = Shape(1, 8, Dim.DYNAMIC, 40) + assertTrue(s.hasDynamic()) + assertTrue(s.isDynamic(2)) + assertFalse(s.isDynamic(1)) + assertEquals(listOf(2), s.dynamicAxes) + assertTrue(s.toString().contains("?")) + // volume is undefined for a dynamic shape — must throw rather than return a corrupt product. + assertFailsWith { s.volume } + // static shapes are unaffected. + assertEquals(1 * 8 * 5 * 40, Shape(1, 8, 5, 40).volume) + } + + @Test + fun void_concat_keeps_growing_cache_dynamic() { + val ops = VoidTensorOps() + val past = VoidOpsTensor( + ShapeOnly(Shape(1, 4, Dim.DYNAMIC, 256)), sk.ainet.lang.types.FP32::class, + ) + val step = VoidOpsTensor( + ShapeOnly(Shape(1, 4, 1, 256)), sk.ainet.lang.types.FP32::class, + ) + val cat = ops.concat(listOf(past, step), dim = 2) + assertEquals(listOf(1, 4, Dim.DYNAMIC, 256), cat.shape.dimensions.toList(), "past ++ step keeps the seq axis dynamic") + } + + @Test + fun void_reshape_passes_dynamic_target_through() { + // A reshape whose target carries a dynamic extent (the cache-as-output-sink `reshape(x, x.shape)`) + // passes through unchanged — it must NOT be mistaken for a `-1` = infer slot. + val ops = VoidTensorOps() + val x = VoidOpsTensor( + ShapeOnly(Shape(1, 4, Dim.DYNAMIC, 256)), sk.ainet.lang.types.FP32::class, + ) + val r = ops.reshape(x, Shape(1, 4, Dim.DYNAMIC, 256)) + assertEquals(listOf(1, 4, Dim.DYNAMIC, 256), r.shape.dimensions.toList()) + } + + @Test + fun slice_all_is_symbolic_full_axis_over_dynamic() { + val all = Slice.All() + assertTrue(all.isValid(Dim.DYNAMIC)) + assertEquals(Dim.DYNAMIC, all.getResultSize(Dim.DYNAMIC), "all() over a dynamic axis stays dynamic") + // A partial range with concrete non-negative bounds over a dynamic axis yields a concrete size. + val r = Slice.Range(2, 5) + assertTrue(r.isValid(Dim.DYNAMIC)) + assertEquals(3, r.getResultSize(Dim.DYNAMIC)) + } +} + +/** Minimal shape-only TensorData for the concat test (mirrors VoidTensorOps' internal one). */ +private class ShapeOnly( + override val shape: Shape, +) : sk.ainet.lang.tensor.data.TensorData { + override fun get(vararg indices: Int): V = error("shape-only") + override fun set(vararg indices: Int, value: V) = error("shape-only") + override fun copyToFloatArray(): FloatArray = error("shape-only") +}