diff --git a/CHANGELOG.md b/CHANGELOG.md index a3cccb6d7..b8867d628 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,8 +2,28 @@ ## [Unreleased] +### Added + +- **`BufferHandle.Floats` — array-free path for ≥2 GiB constants** + ([#1247](https://github.com/SKaiNET-developers/SKaiNET/issues/1247)): external FP32 constants + now ride the aliased `FloatArray` end-to-end (graph → `ExternalParameterRef` → `.irpa`), never + serializing to a single `ByteArray` — the gemma3n tied embedding (262144x2048 FP32 = + `Int.MAX_VALUE` + 1 bytes) structurally cannot exist as one byte buffer. `IrpaWriter` streams + the values little-endian in 64 MiB chunks; `DefaultBufferResolver` reads through a chunked byte + view. Constant element counts fold in `Long`, and an oversized single-buffer serialization now + throws `ConstantTooLargeException` with the remediation instead of the + `NegativeArraySizeException` that was previously mistaken for a registry miss. + ### Changed +- **Graph constants alias live weights; packed params fail loudly** + ([#1247](https://github.com/SKaiNET-developers/SKaiNET/issues/1247)): `TraceToGraphBuilder` no + longer copies every frozen float weight into the graph — the constant's `initial_value` aliases + the live buffer (read-only contract), halving weight residency during export. BF16/FP16 dense + weights widen to one FP32 copy. A frozen parameter with packed storage (Q4_K, Q8_0, ternary, …) + now throws `PackedConstantException` instead of silently becoming a function argument and + producing an unservable module; `PackedConstantHandling.DEQUANTIZE` (threaded through + `toComputeGraph`) opts into dense FP32 extraction instead. - **StableHLO conversion fails loudly by default** ([#1247](https://github.com/SKaiNET-developers/SKaiNET/issues/1247)): converter failures were MLIR comments — a graph whose first node failed to lower could cascade through every downstream diff --git a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api index 08c8017af..bfee28fd0 100644 --- a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api +++ b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api @@ -85,8 +85,8 @@ public class sk/ainet/lang/graph/DefaultExecutionTape : sk/ainet/tape/ExecutionT protected final fun set_recordingStrategy (Lsk/ainet/tape/TapeRecordingStrategy;)V public fun startRecording ()V public fun stopRecording ()V - public final fun toComputeGraph (ZLjava/util/Set;Z)Lsk/ainet/lang/graph/ComputeGraph; - public static synthetic fun toComputeGraph$default (Lsk/ainet/lang/graph/DefaultExecutionTape;ZLjava/util/Set;ZILjava/lang/Object;)Lsk/ainet/lang/graph/ComputeGraph; + public final fun toComputeGraph (ZLjava/util/Set;ZLsk/ainet/lang/trace/PackedConstantHandling;)Lsk/ainet/lang/graph/ComputeGraph; + public static synthetic fun toComputeGraph$default (Lsk/ainet/lang/graph/DefaultExecutionTape;ZLjava/util/Set;ZLsk/ainet/lang/trace/PackedConstantHandling;ILjava/lang/Object;)Lsk/ainet/lang/graph/ComputeGraph; } public final class sk/ainet/lang/graph/DefaultGradientTape : sk/ainet/lang/graph/DefaultExecutionTape, sk/ainet/lang/tensor/ops/DifferentiableTensorOps, sk/ainet/tape/GradientTape { @@ -508,8 +508,8 @@ public final class sk/ainet/lang/graph/utils/GraphvizKt { } public final class sk/ainet/lang/tape/ExtensionsKt { - public static final fun toComputeGraph (Lsk/ainet/tape/ExecutionTape;ZLjava/util/Set;)Lsk/ainet/lang/graph/ComputeGraph; - public static synthetic fun toComputeGraph$default (Lsk/ainet/tape/ExecutionTape;ZLjava/util/Set;ILjava/lang/Object;)Lsk/ainet/lang/graph/ComputeGraph; + public static final fun toComputeGraph (Lsk/ainet/tape/ExecutionTape;ZLjava/util/Set;Lsk/ainet/lang/trace/PackedConstantHandling;)Lsk/ainet/lang/graph/ComputeGraph; + public static synthetic fun toComputeGraph$default (Lsk/ainet/tape/ExecutionTape;ZLjava/util/Set;Lsk/ainet/lang/trace/PackedConstantHandling;ILjava/lang/Object;)Lsk/ainet/lang/graph/ComputeGraph; } public final class sk/ainet/lang/trace/GraphSink : sk/ainet/lang/trace/OpSink { @@ -517,14 +517,28 @@ public final class sk/ainet/lang/trace/GraphSink : sk/ainet/lang/trace/OpSink { public fun onOpExecuted (Lsk/ainet/lang/trace/OpTrace;)V } +public final class sk/ainet/lang/trace/PackedConstantException : java/lang/IllegalStateException { + public fun (Ljava/lang/String;Ljava/lang/String;)V + public final fun getEncodingName ()Ljava/lang/String; + public final fun getTensorId ()Ljava/lang/String; +} + +public final class sk/ainet/lang/trace/PackedConstantHandling : java/lang/Enum { + public static final field DEQUANTIZE Lsk/ainet/lang/trace/PackedConstantHandling; + public static final field FAIL Lsk/ainet/lang/trace/PackedConstantHandling; + public static fun getEntries ()Lkotlin/enums/EnumEntries; + public static fun valueOf (Ljava/lang/String;)Lsk/ainet/lang/trace/PackedConstantHandling; + public static fun values ()[Lsk/ainet/lang/trace/PackedConstantHandling; +} + public final class sk/ainet/lang/trace/TapeSink : sk/ainet/lang/trace/OpSink { public fun (Lsk/ainet/tape/ExecutionTape;)V public fun onOpExecuted (Lsk/ainet/lang/trace/OpTrace;)V } public final class sk/ainet/lang/trace/TraceToGraphBuilder { - public fun (Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/lang/trace/TraceSession;Z)V - public synthetic fun (Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/lang/trace/TraceSession;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun (Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/lang/trace/TraceSession;ZLsk/ainet/lang/trace/PackedConstantHandling;)V + public synthetic fun (Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/lang/trace/TraceSession;ZLsk/ainet/lang/trace/PackedConstantHandling;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun addAll (Ljava/lang/Iterable;)V public final fun addTrace (Lsk/ainet/lang/trace/OpTrace;)V public final fun finalize (Ljava/util/Set;Z)V diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt index ff69a5eea..be52912e2 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt @@ -250,13 +250,19 @@ public open class DefaultExecutionTape( public fun toComputeGraph( synthesizeExternalInputs: Boolean = false, inputTensorIds: Set = emptySet(), - embedConstants: Boolean = true + embedConstants: Boolean = true, + packedConstants: sk.ainet.lang.trace.PackedConstantHandling = + sk.ainet.lang.trace.PackedConstantHandling.FAIL ): 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, embedWeightData = embedConstants) + val builder = TraceToGraphBuilder( + graph, session, + embedWeightData = embedConstants, + packedConstants = packedConstants + ) builder.addAll(_traces) if (synthesizeExternalInputs) { builder.finalize(inputTensorIds, embedConstants = embedConstants) diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/tape/extensions.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/tape/extensions.kt index 9e34542bb..3d0d39b48 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/tape/extensions.kt +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/tape/extensions.kt @@ -12,13 +12,18 @@ import sk.ainet.tape.ExecutionTape * 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. + * @param packedConstants How frozen parameters with packed/quantized storage are treated during + * constant synthesis — fail loudly (default) or dequantize to dense FP32 (issue #1247). */ public fun ExecutionTape.toComputeGraph( synthesizeExternalInputs: Boolean = false, - inputTensorIds: Set = emptySet() + inputTensorIds: Set = emptySet(), + packedConstants: sk.ainet.lang.trace.PackedConstantHandling = + sk.ainet.lang.trace.PackedConstantHandling.FAIL ): ComputeGraph { return when (this) { - is sk.ainet.lang.graph.DefaultExecutionTape -> this.toComputeGraph(synthesizeExternalInputs, inputTensorIds) + is sk.ainet.lang.graph.DefaultExecutionTape -> + this.toComputeGraph(synthesizeExternalInputs, inputTensorIds, packedConstants = packedConstants) else -> DefaultComputeGraph() } } diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/PackedConstantHandling.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/PackedConstantHandling.kt new file mode 100644 index 000000000..dfcf345e0 --- /dev/null +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/PackedConstantHandling.kt @@ -0,0 +1,45 @@ +package sk.ainet.lang.trace + +/** + * How [TraceToGraphBuilder.finalize] treats a frozen parameter whose data is + * packed/quantized ([sk.ainet.lang.tensor.storage.PackedBlockStorage]) and + * therefore cannot be embedded as a float constant directly. + * + * Historically such tensors fell through to the "input" placeholder branch, + * silently turning model weights into function arguments — issue #1247 + * measured a `func @gemma3n` with 190+ weight args and zero dot ops that + * still exported with exit 0. + */ +public enum class PackedConstantHandling { + /** + * Throw [PackedConstantException] naming the tensor and its encoding. + * Default: an unservable module must never be produced silently. + */ + FAIL, + + /** + * Dequantize the packed data to a dense FP32 constant at extraction time + * via [sk.ainet.lang.tensor.storage.PackedBlockStorage.toFloatArray]. + * Opt-in: costs one dense copy of each packed weight (a Q4_K matrix + * grows ~8x), which is exactly the memory class #1247 is fighting — + * use only when the export target genuinely needs dense constants. + */ + DEQUANTIZE, +} + +/** + * A frozen parameter with packed/quantized storage reached constant + * extraction under [PackedConstantHandling.FAIL]. + */ +public class PackedConstantException( + public val tensorId: String, + public val encodingName: String?, +) : IllegalStateException( + "Frozen parameter '$tensorId' has packed storage" + + (encodingName?.let { " (encoding $it)" } ?: "") + + " and cannot be embedded as a float graph constant. Refusing to fall " + + "back to a function-argument placeholder (that silently produces an " + + "unservable module — issue #1247). Either load this weight dense, or " + + "pass PackedConstantHandling.DEQUANTIZE to dequantize it to FP32 at " + + "extraction (costs one dense copy per packed weight)." +) 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 53b920daa..fad7331da 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 @@ -32,7 +32,14 @@ import sk.ainet.lang.tensor.ops.withTensorId public class TraceToGraphBuilder( private val graph: ComputeGraph, private val session: TraceSession? = null, - private val embedWeightData: Boolean = true + private val embedWeightData: Boolean = true, + /** + * How [finalize] treats frozen parameters whose data is packed/quantized + * and cannot become a float constant. Default [PackedConstantHandling.FAIL] + * throws instead of silently synthesizing a function-argument placeholder + * for a model weight (issue #1247). + */ + private val packedConstants: PackedConstantHandling = PackedConstantHandling.FAIL ) { private var nextNodeId = 0L @@ -263,12 +270,30 @@ public class TraceToGraphBuilder( // Try to resolve as a constant from the session val tensor = if (!forceInput && embedConstants) session?.resolve(firstRef.tensorRef) else null - val constantValues = tensor?.let { extractFloatArray(it) } + var constantValues = tensor?.let { extractFloatArray(it) } // Resolved tensors that carry a concrete storage encoding (Q4_K, // Q8_0, TernaryPacked, TurboQuant, …) propagate it onto the // produced spec so later compile stages can preserve the // quantization instead of silently re-materializing FP32. - val encoding = tensor?.data?.inferTensorEncoding() + var encoding = tensor?.data?.inferTensorEncoding() + + // A frozen parameter with packed storage must never fall through to + // the "input" placeholder branch: that silently turns model weights + // into function arguments and produces an unservable module with + // exit 0 (issue #1247, the 190+-arg gemma3n export). + val packedData = tensor?.data as? sk.ainet.lang.tensor.storage.PackedBlockStorage + if (constantValues == null && packedData != null) { + when (packedConstants) { + PackedConstantHandling.FAIL -> + throw PackedConstantException(tensorId, encoding?.name) + PackedConstantHandling.DEQUANTIZE -> { + constantValues = packedData.toFloatArray() + // The embedded constant is now dense FP32 — carrying the + // packed encoding forward would misdescribe it. + encoding = null + } + } + } val syntheticNode: GraphNode val producedSpec: TensorSpec @@ -376,12 +401,27 @@ public class TraceToGraphBuilder( private fun extractFloatArray(tensor: sk.ainet.lang.tensor.Tensor<*, *>): FloatArray? { val data = tensor.data if (data is sk.ainet.lang.tensor.data.FloatArrayTensorData) { - val buffer = data.buffer - return buffer.copyOf() + // ALIASED, not copied (#1247): the graph constant shares the live + // weight buffer. Copying doubled residency of every model weight — + // the module tree and TraceSession keep the original alive for the + // whole build, so a full E2B extraction OOMed a 46 GB heap. + // Contract: consumers of "initial_value"/"weights" parameters are + // read-only (the HLO ConstantByteSerializer and inline emitters); + // never mutate the array behind these parameters. + return data.buffer } - - // 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` + + // Narrow floats (BF16/FP16 dense) widen to one FP32 copy — this is a + // real conversion, not an alias, and only valid for static shapes. + if (data is sk.ainet.lang.tensor.data.NarrowFloatTensorData && + !tensor.shape.hasDynamic() + ) { + return data.copyToFloatArray() + } + + // Nothing else is materializable here: packed storage is handled by the + // caller via PackedConstantHandling, 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-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ConstantExtractionTest.kt b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ConstantExtractionTest.kt new file mode 100644 index 000000000..7a119227f --- /dev/null +++ b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ConstantExtractionTest.kt @@ -0,0 +1,225 @@ +package sk.ainet.compile.graph + +import sk.ainet.context.Phase +import sk.ainet.exec.tensor.ops.DefaultCpuOps +import sk.ainet.lang.graph.DefaultExecutionTape +import sk.ainet.lang.graph.DefaultGradientTape +import sk.ainet.lang.graph.DefaultGraphExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.data.Bf16DenseTensorData +import sk.ainet.lang.tensor.data.DenseTensorDataFactory +import sk.ainet.lang.tensor.data.FloatArrayTensorData +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.trace.OpTrace +import sk.ainet.lang.trace.PackedConstantException +import sk.ainet.lang.trace.PackedConstantHandling +import sk.ainet.lang.tensor.ops.tensorEncoding +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** + * #1247 constant-extraction contract: + * - float weights are ALIASED into the graph, never copied (double residency + * of every model weight OOMed the E2B export at a 46 GB heap); + * - packed frozen params fail loudly by default instead of silently becoming + * function arguments (the 190+-arg unservable gemma3n module), with + * dequantize-to-FP32 as the opt-in; + * - narrow-float (BF16/FP16) dense weights widen to one FP32 constant. + */ +class ConstantExtractionTest { + + private fun ctx(): DefaultGraphExecutionContext { + val dataFactory = DenseTensorDataFactory() + return DefaultGraphExecutionContext( + baseOps = DefaultCpuOps(dataFactory), + phase = Phase.TRAIN, + tensorDataFactory = dataFactory, + createTapeFactory = { _ -> DefaultGradientTape(true) }, + ) + } + + @Test + fun float_weight_constant_aliases_the_live_buffer() { + val trainCtx = ctx() + val input = trainCtx.fromFloatArray(Shape(2, 4), FP32::class, FloatArray(8) { it.toFloat() }) + val weight = trainCtx.fromFloatArray(Shape(4, 3), FP32::class, FloatArray(12) { it * 0.5f }) + val output = trainCtx.fromFloatArray(Shape(2, 3), FP32::class, FloatArray(6)) + + val tape = DefaultExecutionTape(trainCtx.session) + tape.startRecording() + val inputRef = tape.session.refOf(input) + val weightRef = tape.session.refOf(weight) + val outputRef = tape.session.refOf(output) + tape.recordTrace( + OpTrace( + opType = "matmul", + inputs = listOf(inputRef, weightRef), + outputs = listOf(outputRef), + attributes = emptyMap() + ) + ) + tape.stopRecording() + + val graph = tape.toComputeGraph( + synthesizeExternalInputs = true, + inputTensorIds = setOf(inputRef.id) + ) + + val weightNode = graph.nodes.single { it.operation.type == "constant" } + val embedded = weightNode.operation.parameters["initial_value"] as FloatArray + val liveBuffer = (weight.data as FloatArrayTensorData).buffer + assertSame( + liveBuffer, embedded, + "the graph constant must alias the live weight buffer, not copy it (#1247 double residency)" + ) + } + + @Test + fun packed_frozen_param_fails_loudly_by_default() { + val trainCtx = ctx() + val packed = Q8_0BlockTensorData.fromRawBytes(Shape(32), ByteArray(34)) + @Suppress("UNCHECKED_CAST") + val weight = trainCtx.fromData(packed as sk.ainet.lang.tensor.data.TensorData, FP32::class) + val input = trainCtx.fromFloatArray(Shape(32), FP32::class, FloatArray(32)) + val output = trainCtx.fromFloatArray(Shape(32), FP32::class, FloatArray(32)) + + val tape = DefaultExecutionTape(trainCtx.session) + tape.startRecording() + val inputRef = tape.session.refOf(input) + val weightRef = tape.session.refOf(weight) + val outputRef = tape.session.refOf(output) + tape.recordTrace( + OpTrace( + opType = "add", + inputs = listOf(inputRef, weightRef), + outputs = listOf(outputRef), + attributes = emptyMap() + ) + ) + tape.stopRecording() + + val e = assertFailsWith { + tape.toComputeGraph( + synthesizeExternalInputs = true, + inputTensorIds = setOf(inputRef.id) + ) + } + val message = e.message ?: "" + assertTrue("Q8_0" in message, "exception must name the packed encoding: $message") + assertTrue(weightRef.id in message, "exception must name the tensor: $message") + assertTrue("DEQUANTIZE" in message, "exception must point at the remediation: $message") + } + + @Test + fun packed_frozen_param_dequantizes_when_opted_in() { + val trainCtx = ctx() + val packed = Q8_0BlockTensorData.fromRawBytes(Shape(32), ByteArray(34)) + @Suppress("UNCHECKED_CAST") + val weight = trainCtx.fromData(packed as sk.ainet.lang.tensor.data.TensorData, FP32::class) + val input = trainCtx.fromFloatArray(Shape(32), FP32::class, FloatArray(32)) + val output = trainCtx.fromFloatArray(Shape(32), FP32::class, FloatArray(32)) + + val tape = DefaultExecutionTape(trainCtx.session) + tape.startRecording() + val inputRef = tape.session.refOf(input) + val weightRef = tape.session.refOf(weight) + val outputRef = tape.session.refOf(output) + tape.recordTrace( + OpTrace( + opType = "add", + inputs = listOf(inputRef, weightRef), + outputs = listOf(outputRef), + attributes = emptyMap() + ) + ) + tape.stopRecording() + + val graph = tape.toComputeGraph( + synthesizeExternalInputs = true, + inputTensorIds = setOf(inputRef.id), + packedConstants = PackedConstantHandling.DEQUANTIZE + ) + + val weightNode = graph.nodes.single { it.operation.type == "constant" } + val embedded = weightNode.operation.parameters["initial_value"] as FloatArray + assertEquals(32, embedded.size, "dequantized constant must carry the full logical extent") + // The embedded constant is dense FP32 now — the packed encoding must + // not ride along and misdescribe it. + val spec = weightNode.outputs.single() + assertEquals( + null, + spec.tensorEncoding, + "a dequantized constant must not carry the packed encoding" + ) + } + + @Test + fun narrow_float_weight_widens_to_fp32_constant() { + val trainCtx = ctx() + // 4 bf16 zeros: widening must produce 4 FP32 zeros. + val bf16 = Bf16DenseTensorData(Shape(4), ByteArray(8)) + @Suppress("UNCHECKED_CAST") + val weight = trainCtx.fromData(bf16 as sk.ainet.lang.tensor.data.TensorData, FP32::class) + val input = trainCtx.fromFloatArray(Shape(4), FP32::class, FloatArray(4)) + val output = trainCtx.fromFloatArray(Shape(4), FP32::class, FloatArray(4)) + + val tape = DefaultExecutionTape(trainCtx.session) + tape.startRecording() + val inputRef = tape.session.refOf(input) + val weightRef = tape.session.refOf(weight) + val outputRef = tape.session.refOf(output) + tape.recordTrace( + OpTrace( + opType = "add", + inputs = listOf(inputRef, weightRef), + outputs = listOf(outputRef), + attributes = emptyMap() + ) + ) + tape.stopRecording() + + val graph = tape.toComputeGraph( + synthesizeExternalInputs = true, + inputTensorIds = setOf(inputRef.id) + ) + + val weightNode = graph.nodes.single { it.operation.type == "constant" } + val embedded = weightNode.operation.parameters["initial_value"] as FloatArray + assertEquals(4, embedded.size) + assertTrue(embedded.all { it == 0.0f }, "bf16 zeros must widen to FP32 zeros") + } + + @Test + fun forced_input_tensor_still_becomes_input_node() { + val trainCtx = ctx() + val input = trainCtx.fromFloatArray(Shape(2), FP32::class, FloatArray(2)) + val output = trainCtx.fromFloatArray(Shape(2), FP32::class, FloatArray(2)) + + val tape = DefaultExecutionTape(trainCtx.session) + tape.startRecording() + val inputRef = tape.session.refOf(input) + val outputRef = tape.session.refOf(output) + tape.recordTrace( + OpTrace( + opType = "relu", + inputs = listOf(inputRef), + outputs = listOf(outputRef), + attributes = emptyMap() + ) + ) + tape.stopRecording() + + val graph = tape.toComputeGraph( + synthesizeExternalInputs = true, + inputTensorIds = setOf(inputRef.id) + ) + assertTrue( + graph.nodes.any { it.operation.type == "input" }, + "a tensor in inputTensorIds must remain a function argument even though it is resolvable" + ) + } +} diff --git a/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api b/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api index 386cbf230..f6f8f4e3e 100644 --- a/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api +++ b/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api @@ -41,6 +41,10 @@ public final class sk/ainet/compile/hlo/ConstantMaterializationPolicy$SizeThresh public fun toString ()Ljava/lang/String; } +public final class sk/ainet/compile/hlo/ConstantTooLargeException : java/lang/IllegalStateException { + public fun (Ljava/lang/String;)V +} + public final class sk/ainet/compile/hlo/ConversionContext { public fun (Lsk/ainet/compile/hlo/TypeMapper;)V public fun (Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/lang/graph/ComputeGraph;)V diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/ConstantByteSerializer.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/ConstantByteSerializer.kt index 9021ba9b2..1d514fc3e 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/ConstantByteSerializer.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/ConstantByteSerializer.kt @@ -23,6 +23,7 @@ internal fun numberListToLittleEndianBytes( expectedElements: Int ): ByteArray { val count = expectedElements.coerceAtLeast(values.size) + requireSerializableByteCount(count, bytesPerSerializedElement(dtype), dtype) val normalized = dtype.uppercase() return when (normalized) { @@ -91,6 +92,7 @@ internal fun floatArrayToLittleEndianBytes( expectedElements: Int ): ByteArray { val count = expectedElements.coerceAtLeast(values.size) + requireSerializableByteCount(count, bytesPerSerializedElement(dtype), dtype) val n = minOf(count, values.size) return when (dtype.uppercase()) { "FP32", "F32", "FLOAT32" -> { @@ -125,9 +127,60 @@ internal fun floatArrayToLittleEndianBytes( * Expected element count for a (possibly empty) shape. Empty shape * (scalar) means one element; `null` / absent dims degrade to 0 so the * caller can detect "no declared shape". + * + * [Long] arithmetic (#1247): the gemma3n token embedding is + * 262144 x 2048 = 536,870,912 elements — an [Int] fold of its *byte* + * count goes negative, which previously surfaced as a + * `NegativeArraySizeException` inside the serializer and was mistaken + * for a registry miss ("Unsupported op 'weight' … Known names: […]"). + */ +internal fun elementCountFromShape(shape: List?): Long { + if (shape == null) return 0L + if (shape.isEmpty()) return 1L + return shape.fold(1L) { acc, d -> acc * d } +} + +/** + * A constant's serialized form would exceed the JVM single-array ceiling. + * Deliberately NOT an [IllegalArgumentException]: the converter's + * "unsupported dtype" fallback catches that type to retry inline emission, + * and inlining a multi-GiB tensor as text is exactly the wrong recovery. + */ +public class ConstantTooLargeException(message: String) : IllegalStateException(message) + +private fun bytesPerSerializedElement(dtype: String): Long = when (dtype.uppercase()) { + "FP64", "F64", "FLOAT64", "I64", "INT64" -> 8L + else -> 4L +} + +/** + * Narrow a [Long] element count for the byte-serialization paths, which + * address a single array. Throws [ConstantTooLargeException] instead of + * truncating — truncation is how the #1247 embedding turned into a + * `NegativeArraySizeException`. */ -internal fun elementCountFromShape(shape: List?): Int { - if (shape == null) return 0 - if (shape.isEmpty()) return 1 - return shape.fold(1) { acc, d -> acc * d } +internal fun checkedIntElements(elementCount: Long): Int { + if (elementCount > Int.MAX_VALUE - 8L) { + throw ConstantTooLargeException( + "Constant of $elementCount elements exceeds the single-array serialization " + + "ceiling. Use ConstantMaterializationPolicy.ExternalAlways with FP32 values — " + + "the external path carries a FloatArray (BufferHandle.Floats) without byte " + + "serialization (issue #1247)." + ) + } + return elementCount.toInt() +} + +private fun requireSerializableByteCount(count: Int, bytesPerElement: Long, dtype: String) { + val byteCount = count.toLong() * bytesPerElement + // A JVM array tops out just under Int.MAX_VALUE entries; keep a small + // margin for VM-specific header overhead. + if (byteCount > Int.MAX_VALUE - 8L) { + throw ConstantTooLargeException( + "Constant of $count $dtype elements needs $byteCount bytes; single-buffer " + + "serialization caps at 2 GiB - 1. Use ConstantMaterializationPolicy.ExternalAlways " + + "with FP32 values — the external path carries a FloatArray (BufferHandle.Floats) " + + "without any byte serialization (issue #1247)." + ) + } } diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ConstantOperationsConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ConstantOperationsConverter.kt index 3dac1ce8a..01c30d6b4 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ConstantOperationsConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ConstantOperationsConverter.kt @@ -5,6 +5,7 @@ import sk.ainet.compile.hlo.ConversionContext import sk.ainet.compile.hlo.ConversionResult import sk.ainet.compile.hlo.ExternalParameterRef import sk.ainet.compile.hlo.StableHloOperationConverter +import sk.ainet.compile.hlo.checkedIntElements import sk.ainet.compile.hlo.elementCountFromShape import sk.ainet.compile.hlo.floatArrayToLittleEndianBytes import sk.ainet.compile.hlo.numberListToLittleEndianBytes @@ -410,9 +411,9 @@ public class ConstantOperationsConverter : StableHloOperationConverter { val encoding = outputSpec.tensorEncoding ?: TensorEncoding.Dense(bytesPerElement = bytesPerElement(outputSpec.dtype)) val elementCount = elementCountFromShape(outputSpec.shape) - if (elementCount <= 0) return null + if (elementCount <= 0L) return null - val logicalBytes = encoding.physicalBytes(elementCount.toLong()) ?: return null + val logicalBytes = encoding.physicalBytes(elementCount) ?: return null val scope = when (policy) { is ConstantMaterializationPolicy.InlineAlways -> return null is ConstantMaterializationPolicy.ExternalAlways -> policy.scope @@ -424,9 +425,11 @@ public class ConstantOperationsConverter : StableHloOperationConverter { // Serialize now. Fall back to inline on unsupported dtype — // a loud exception here would defeat the "default path is - // safe" invariant of the seam. + // safe" invariant of the seam. A ConstantTooLargeException is NOT + // caught: inlining a multi-GiB tensor as text is the wrong recovery + // (#1247) — it propagates with its actionable message. val bytes = try { - numberListToLittleEndianBytes(values, outputSpec.dtype, elementCount) + numberListToLittleEndianBytes(values, outputSpec.dtype, checkedIntElements(elementCount)) } catch (e: IllegalArgumentException) { context.emitComment( "external materialization fell back to inline for ${node.id}: ${e.message}" @@ -489,9 +492,9 @@ public class ConstantOperationsConverter : StableHloOperationConverter { val encoding = outputSpec.tensorEncoding ?: TensorEncoding.Dense(bytesPerElement = bytesPerElement(outputSpec.dtype)) val elementCount = elementCountFromShape(outputSpec.shape) - if (elementCount <= 0) return null + if (elementCount <= 0L) return null - val logicalBytes = encoding.physicalBytes(elementCount.toLong()) ?: return null + val logicalBytes = encoding.physicalBytes(elementCount) ?: return null val scope = when (policy) { is ConstantMaterializationPolicy.InlineAlways -> return null is ConstantMaterializationPolicy.ExternalAlways -> policy.scope @@ -501,13 +504,31 @@ public class ConstantOperationsConverter : StableHloOperationConverter { } } - val bytes = try { - floatArrayToLittleEndianBytes(values, outputSpec.dtype, elementCount) - } catch (e: IllegalArgumentException) { - context.emitComment( - "external materialization fell back to inline for ${node.id}: ${e.message}" - ) - return null + val normalizedDtype = outputSpec.dtype.uppercase() + val source: BufferHandle = if ( + (normalizedDtype == "FP32" || normalizedDtype == "F32" || normalizedDtype == "FLOAT32") && + values.size.toLong() == elementCount + ) { + // Array-free path (#1247): the FloatArray aliases the live module + // weight all the way from TraceToGraphBuilder — hand it to the + // packager as-is. No byte serialization means no extra copy and + // no 2 GiB ByteArray ceiling: the gemma3n tied embedding + // (262144x2048 = Int.MAX_VALUE + 1 bytes) only exports this way. + BufferHandle.Floats(values) + } else { + // Under-filled initializations and non-FP32 dtypes keep the + // padded byte serialization. ConstantTooLargeException is NOT + // caught — inlining a multi-GiB tensor as text is the wrong + // recovery; it propagates with its actionable message. + val bytes = try { + floatArrayToLittleEndianBytes(values, outputSpec.dtype, checkedIntElements(elementCount)) + } catch (e: IllegalArgumentException) { + context.emitComment( + "external materialization fell back to inline for ${node.id}: ${e.message}" + ) + return null + } + BufferHandle.Owned(bytes) } val key = outputSpec.name.ifEmpty { node.id } @@ -516,7 +537,7 @@ public class ConstantOperationsConverter : StableHloOperationConverter { scope = scope, key = key, encoding = encoding, - source = BufferHandle.Owned(bytes), + source = source, blockOrder = outputSpec.blockOrder, ) ) diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/BigConstantMaterializationTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/BigConstantMaterializationTest.kt new file mode 100644 index 000000000..e4284b4da --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/BigConstantMaterializationTest.kt @@ -0,0 +1,148 @@ +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.Operation +import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.tensor.ops.ValidationResult +import sk.ainet.lang.tensor.storage.BufferHandle +import sk.ainet.lang.types.DType +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** + * #1247 big-constant contract: element counts fold in Long (the gemma3n + * embedding is 262144 x 2048 = Int.MAX_VALUE + 1 BYTES — an Int fold went + * negative and surfaced as a NegativeArraySizeException mistaken for a + * registry miss), oversized single-buffer serialization refuses with an + * actionable message instead of throwing array-size garbage, and the + * external FP32 path hands the aliased FloatArray to the packager as + * [BufferHandle.Floats] with no byte serialization at all. + */ +class BigConstantMaterializationTest { + + @Test + fun elementCountFromShape_folds_in_long() { + assertEquals(536_870_912L, elementCountFromShape(listOf(262_144, 2_048))) + assertEquals(1L, elementCountFromShape(emptyList())) + assertEquals(0L, elementCountFromShape(null)) + // 2 Gi elements — Int fold would be 0/negative; Long must be exact. + assertEquals(2_147_483_648L, elementCountFromShape(listOf(65_536, 32_768))) + } + + @Test + fun checkedIntElements_refuses_oversized_counts_with_actionable_message() { + val e = assertFailsWith { + checkedIntElements(536_870_912L * 4L) + } + assertTrue("ExternalAlways" in (e.message ?: ""), "refusal must point at the external path") + assertEquals(1024, checkedIntElements(1024L)) + } + + @Test + fun serializer_refuses_byte_overflow_instead_of_negative_array_size() { + // 536,870,912 FP32 elements = 2 GiB of bytes: previously + // ByteArray(count * 4) threw NegativeArraySizeException. + val e = assertFailsWith { + floatArrayToLittleEndianBytes(FloatArray(0), "FP32", 536_870_912) + } + assertTrue("2 GiB" in (e.message ?: ""), "refusal must state the ceiling: ${e.message}") + } + + @Test + fun external_fp32_path_aliases_the_float_array_without_serialization() { + val weights = FloatArray(12) { it * 0.25f } + val module = convertWeightGraph(weights, shape = listOf(4, 3)) + + val ref = module.externalParameters.single() + val source = ref.source + assertTrue(source is BufferHandle.Floats, "FP32 external constant must ride BufferHandle.Floats, got $source") + assertSame(weights, source.data, "the handle must alias the input array — zero copies end-to-end") + assertEquals(48L, source.sizeInBytes) + assertTrue(module.content.contains("util.global.load"), "external constant must load from a util.global") + } + + @Test + fun tied_weight_feeding_two_consumers_emits_one_global() { + val weights = FloatArray(8) { it.toFloat() } + val graph = DefaultComputeGraph() + val weightNode = GraphNode( + id = "w1", + operation = weightOp(weights), + inputs = emptyList(), + outputs = listOf(TensorSpec("tied_embed", listOf(2, 4), "FP32")) + ) + val consumerA = computeNode("relu1", "relu", inputs = 1) + val consumerB = computeNode("relu2", "relu", inputs = 1) + graph.addNode(weightNode) + graph.addNode(consumerA) + graph.addNode(consumerB) + graph.addEdge(GraphEdge("ea", weightNode, consumerA, 0, 0, weightNode.outputs[0])) + graph.addEdge(GraphEdge("eb", weightNode, consumerB, 0, 0, weightNode.outputs[0])) + + val converter = StableHloConverterFactory.createExtended( + policy = ConstantMaterializationPolicy.ExternalAlways(scope = "model") + ) + val module = converter.convert(graph, "tied") + + assertEquals(1, module.externalParameters.size, "one tied weight must register exactly one external ref") + val globalDecls = module.content.lines().count { "util.global private @tied_embed" in it } + assertEquals(1, globalDecls, "one tied weight must declare exactly one util.global:\n${module.content}") + } + + private fun convertWeightGraph(weights: FloatArray, shape: List): StableHloModule { + val graph = DefaultComputeGraph() + val weightNode = GraphNode( + id = "w1", + operation = weightOp(weights), + inputs = emptyList(), + outputs = listOf(TensorSpec("w1_spec", shape, "FP32")) + ) + val consumer = computeNode("relu1", "relu", inputs = 1) + graph.addNode(weightNode) + graph.addNode(consumer) + graph.addEdge(GraphEdge("e1", weightNode, consumer, 0, 0, weightNode.outputs[0])) + + val converter = StableHloConverterFactory.createExtended( + policy = ConstantMaterializationPolicy.ExternalAlways(scope = "model") + ) + return converter.convert(graph, "weights") + } + + private fun weightOp(values: FloatArray): Operation = object : Operation { + override val name: String = "weight" + override val type: String = "constant" + override val parameters: Map = mapOf( + "initial_value" to values, + "trainable" to false + ) + override fun execute(inputs: List>): List> = + throw UnsupportedOperationException("test fixture only") + override fun validateInputs(inputs: List): ValidationResult = ValidationResult.Valid + override fun inferOutputs(inputs: List): List = emptyList() + override fun clone(newParameters: Map): Operation = this + override fun serialize(): Map = mapOf("name" to name, "type" to type) + } + + private fun computeNode(id: String, opName: String, inputs: Int): GraphNode = GraphNode( + id = id, + operation = object : Operation { + override val name: String = opName + override val type: String = "compute" + override val parameters: Map = emptyMap() + override fun execute(inputs: List>): List> = + throw UnsupportedOperationException("test fixture only") + override fun validateInputs(inputs: List): ValidationResult = ValidationResult.Valid + override fun inferOutputs(inputs: List): List = emptyList() + override fun clone(newParameters: Map): Operation = this + override fun serialize(): Map = mapOf("name" to name, "type" to type) + }, + inputs = List(inputs) { TensorSpec("in$it", listOf(2, 4), "FP32") }, + outputs = listOf(TensorSpec("$id-out", listOf(2, 4), "FP32")) + ) +} diff --git a/skainet-io/skainet-io-iree-params/src/commonMain/kotlin/sk/ainet/io/irpa/IrpaWriter.kt b/skainet-io/skainet-io-iree-params/src/commonMain/kotlin/sk/ainet/io/irpa/IrpaWriter.kt index e3f6caed1..50ff95647 100644 --- a/skainet-io/skainet-io-iree-params/src/commonMain/kotlin/sk/ainet/io/irpa/IrpaWriter.kt +++ b/skainet-io/skainet-io-iree-params/src/commonMain/kotlin/sk/ainet/io/irpa/IrpaWriter.kt @@ -224,6 +224,7 @@ public class IrpaWriter { private fun writeBufferHandle(sink: Sink, handle: BufferHandle) { when (handle) { is BufferHandle.Owned -> writeByteArray(sink, handle.data, handle.offset, handle.sizeInBytes.toInt()) + is BufferHandle.Floats -> writeFloatArrayLe(sink, handle.data) is BufferHandle.Borrowed -> writeByteArray(sink, handle.data, handle.offset, handle.sizeInBytes.toInt()) is BufferHandle.FileBacked -> writeFileBackedBytes(sink, handle) else -> throw IllegalArgumentException( @@ -235,6 +236,29 @@ public class IrpaWriter { } } + private fun writeFloatArrayLe(sink: Sink, data: FloatArray) { + // Stream little-endian in 64 MiB chunks: a Floats handle exists + // precisely because its logical bytes can exceed what a single + // ByteArray can hold (issue #1247, the 2 GiB tied embedding) — + // never materialize the whole payload. + val chunkFloats = 16 * 1024 * 1024 + val buf = ByteArray(minOf(chunkFloats, data.size.coerceAtLeast(1)) * 4) + var i = 0 + while (i < data.size) { + val n = minOf(chunkFloats, data.size - i) + var b = 0 + for (j in i until i + n) { + val bits = data[j].toRawBits() + buf[b++] = (bits and 0xff).toByte() + buf[b++] = (bits ushr 8 and 0xff).toByte() + buf[b++] = (bits ushr 16 and 0xff).toByte() + buf[b++] = (bits ushr 24 and 0xff).toByte() + } + sink.write(buf, 0, b) + i += n + } + } + private fun writeByteArray(sink: Sink, data: ByteArray, offset: Int, length: Int) { // Bulk range write. Owned/Borrowed buffers can be large for real // LLM weights (a 262153x640 FP32 embedding = ~670MB); a diff --git a/skainet-io/skainet-io-iree-params/src/commonTest/kotlin/sk/ainet/io/irpa/IrpaFloatsParityTest.kt b/skainet-io/skainet-io-iree-params/src/commonTest/kotlin/sk/ainet/io/irpa/IrpaFloatsParityTest.kt new file mode 100644 index 000000000..010e5e382 --- /dev/null +++ b/skainet-io/skainet-io-iree-params/src/commonTest/kotlin/sk/ainet/io/irpa/IrpaFloatsParityTest.kt @@ -0,0 +1,43 @@ +package sk.ainet.io.irpa + +import kotlinx.io.Buffer +import kotlinx.io.readByteArray +import sk.ainet.compile.hlo.ExternalParameterRef +import sk.ainet.lang.tensor.storage.BufferHandle +import sk.ainet.lang.tensor.storage.TensorEncoding +import kotlin.test.Test +import kotlin.test.assertContentEquals + +/** + * #1247: a [BufferHandle.Floats] entry must produce byte-identical archive + * output to the same values serialized eagerly into [BufferHandle.Owned] — + * the Floats handle only removes the up-front byte copy (and with it the + * 2 GiB single-array ceiling), never changes the on-disk format. + */ +class IrpaFloatsParityTest { + + @Test + fun floats_entry_is_byte_identical_to_owned_serialization() { + val values = FloatArray(1027) { (it - 500) * 0.37f } // odd size: crosses chunk lane logic + val bytes = ByteArray(values.size * 4) + for (i in values.indices) { + val bits = values[i].toRawBits() + bytes[i * 4] = (bits and 0xff).toByte() + bytes[i * 4 + 1] = (bits ushr 8 and 0xff).toByte() + bytes[i * 4 + 2] = (bits ushr 16 and 0xff).toByte() + bytes[i * 4 + 3] = (bits ushr 24 and 0xff).toByte() + } + + fun ref(source: BufferHandle) = ExternalParameterRef( + scope = "model", + key = "w", + encoding = TensorEncoding.Dense(bytesPerElement = 4), + source = source + ) + + val ownedOut = Buffer().also { IrpaWriter().write(listOf(ref(BufferHandle.Owned(bytes))), it) } + val floatsOut = Buffer().also { IrpaWriter().write(listOf(ref(BufferHandle.Floats(values))), it) } + + assertContentEquals(ownedOut.readByteArray(), floatsOut.readByteArray()) + } +} diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 24e79f9c3..51766e4b5 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -7090,6 +7090,14 @@ public final class sk/ainet/lang/tensor/storage/BufferHandle$FileBacked : sk/ain public fun isMutable ()Z } +public final class sk/ainet/lang/tensor/storage/BufferHandle$Floats : sk/ainet/lang/tensor/storage/BufferHandle { + public fun ([F)V + public final fun getData ()[F + public fun getOwnership ()Lsk/ainet/lang/tensor/storage/Ownership; + public fun getSizeInBytes ()J + public fun isMutable ()Z +} + public final class sk/ainet/lang/tensor/storage/BufferHandle$Owned : sk/ainet/lang/tensor/storage/BufferHandle { public fun ([BIJ)V public synthetic fun ([BIJILkotlin/jvm/internal/DefaultConstructorMarker;)V @@ -7222,6 +7230,15 @@ public final class sk/ainet/lang/tensor/storage/DeviceKind : java/lang/Enum { public static fun values ()[Lsk/ainet/lang/tensor/storage/DeviceKind; } +public final class sk/ainet/lang/tensor/storage/FloatArrayByteViewAccessor : sk/ainet/lang/tensor/storage/BufferAccessor { + public fun ([F)V + public fun close ()V + public fun getSizeInBytes ()J + public fun readAllBytes ()[B + public fun readByte (J)B + public fun readBytes (JI)[B +} + public abstract interface annotation class sk/ainet/lang/tensor/storage/KvCache : java/lang/annotation/Annotation { public abstract fun device ()Lsk/ainet/lang/tensor/storage/DeviceKind; public abstract fun keyBits ()I diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Describe.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Describe.kt index 649788adb..251dcdba9 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Describe.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/memory/Describe.kt @@ -33,6 +33,7 @@ public fun TensorStorage.describe(id: sk.ainet.lang.tensor.TensorId? = null): St append( when (val b = buffer) { is BufferHandle.Owned -> "Owned" + is BufferHandle.Floats -> "Floats" is BufferHandle.Borrowed -> "Borrowed" is BufferHandle.Aliased -> "Aliased" is BufferHandle.FileBacked -> "Mapped" diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/BufferAccessor.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/BufferAccessor.kt index 6527b51c2..48f2b8b88 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/BufferAccessor.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/BufferAccessor.kt @@ -51,6 +51,7 @@ public class DefaultBufferResolver( override fun resolve(handle: BufferHandle): BufferAccessor = when (handle) { is BufferHandle.Owned -> ByteArrayAccessor(handle.data, handle.offset, handle.sizeInBytes) + is BufferHandle.Floats -> FloatArrayByteViewAccessor(handle.data) is BufferHandle.Borrowed -> ByteArrayAccessor(handle.data, handle.offset, handle.sizeInBytes) is BufferHandle.Aliased -> resolve(handle.parent).sliced(handle.byteOffset, handle.sizeInBytes) is BufferHandle.FileBacked -> { @@ -97,6 +98,51 @@ public class ByteArrayAccessor( } /** Helper to create a sliced accessor from any accessor. */ +/** + * Little-endian byte view over a [BufferHandle.Floats] array (issue #1247). + * + * The backing tensor can exceed 2 GiB of logical bytes — that is the whole + * point of the Floats handle — so [readAllBytes] on such a buffer throws via + * the Int narrowing; consumers must read in chunks ([readBytes] with a + * bounded length) or per element. + */ +public class FloatArrayByteViewAccessor( + private val data: FloatArray, +) : BufferAccessor { + + override val sizeInBytes: Long = data.size.toLong() * 4L + + override fun readByte(offset: Long): Byte { + require(offset in 0 until sizeInBytes) { "Offset out of bounds: $offset" } + val bits = data[(offset ushr 2).toInt()].toRawBits() + val lane = (offset and 3L).toInt() + return ((bits ushr (lane * 8)) and 0xff).toByte() + } + + override fun readBytes(offset: Long, length: Int): ByteArray { + require(length >= 0) { "Length must be non-negative: $length" } + require(offset >= 0 && offset + length <= sizeInBytes) { + "Range [$offset, ${offset + length}) out of bounds for $sizeInBytes bytes" + } + val out = ByteArray(length) + var i = 0 + while (i < length) { + val at = offset + i + val bits = data[(at ushr 2).toInt()].toRawBits() + var lane = (at and 3L).toInt() + // Copy the remaining lanes of the current float in one go. + while (lane < 4 && i < length) { + out[i] = ((bits ushr (lane * 8)) and 0xff).toByte() + lane++ + i++ + } + } + return out + } + + override fun close() {} // no-op: aliases a heap array owned elsewhere +} + private fun BufferAccessor.sliced(byteOffset: Long, size: Long): BufferAccessor { if (this is ByteArrayAccessor) return this.sliced(byteOffset, size) // Fallback: wrap in a delegating accessor diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/BufferHandle.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/BufferHandle.kt index 8e50d089d..3d1ba5342 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/BufferHandle.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/BufferHandle.kt @@ -41,6 +41,25 @@ public sealed interface BufferHandle { override val ownership: Ownership get() = Ownership.OWNED } + /** + * FP32 logical values held as a primitive [FloatArray], read-only. + * + * Exists because a single [ByteArray] caps at 2 GiB − 1 while a + * [FloatArray] holds up to 2 Gi elements (8 GiB logical) — a + * 262144x2048 FP32 embedding is exactly [Int.MAX_VALUE] + 1 bytes and + * therefore can never be serialized into one byte buffer (issue #1247). + * The array typically aliases a live model weight: consumers must not + * mutate it, and must stream it to bytes in chunks (little-endian + * [Float.toRawBits]) rather than materializing a full byte copy. + */ + public class Floats( + public val data: FloatArray, + ) : BufferHandle { + override val sizeInBytes: Long get() = data.size.toLong() * 4L + override val isMutable: Boolean get() = false + override val ownership: Ownership get() = Ownership.BORROWED + } + /** * A reference to externally-owned memory (e.g. a caller-supplied array). * The runtime must not free or resize it. Mutation is possible only if