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 2636197a6..11deb921a 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 @@ -5,6 +5,8 @@ import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.withRequiresGrad import sk.ainet.lang.tensor.ops.Operation import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.tensor.ops.withTensorEncoding +import sk.ainet.lang.tensor.ops.withTensorId import sk.ainet.lang.types.DType import sk.ainet.lang.trace.OpTrace import sk.ainet.lang.trace.TraceToGraphBuilder @@ -80,7 +82,7 @@ public open class DefaultExecutionTape( name = ref.id, shape = inputShapes?.getOrNull(i) ?: ref.shape.dimensions.toList(), dtype = inputDTypes?.getOrNull(i) ?: ref.dtype.name, - ) + ).withTensorEncoding(ref.encoding).withTensorId(ref.tensorId) } val outputs = List(trace.outputs.size) { i -> val ref = trace.outputs[i] @@ -88,7 +90,7 @@ public open class DefaultExecutionTape( name = ref.id, shape = outputShapes?.getOrNull(i) ?: ref.shape.dimensions.toList(), dtype = outputDTypes?.getOrNull(i) ?: ref.dtype.name, - ) + ).withTensorEncoding(ref.encoding).withTensorId(ref.tensorId) } val op = object : sk.ainet.lang.tensor.ops.Operation { @@ -126,7 +128,7 @@ public open class DefaultExecutionTape( shape = tensor.shape.dimensions.toList(), dtype = tensor.dtype.toString(), requiresGrad = tensor.requiresGrad - ) + ).withTensorEncoding(ref.encoding).withTensorId(ref.tensorId) } val outputSpecs = outputs.map { tensor -> @@ -136,7 +138,7 @@ public open class DefaultExecutionTape( shape = tensor.shape.dimensions.toList(), dtype = tensor.dtype.toString(), requiresGrad = tensor.requiresGrad - ) + ).withTensorEncoding(ref.encoding).withTensorId(ref.tensorId) } val recordedOp = RecordedOperation( 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 832f254b8..719ee1743 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 @@ -7,7 +7,9 @@ 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.ops.inferTensorEncoding +import sk.ainet.lang.tensor.ops.tensorId import sk.ainet.lang.tensor.ops.withTensorEncoding +import sk.ainet.lang.tensor.ops.withTensorId /** * Shared builder to convert OpTrace streams into a ComputeGraph. @@ -291,6 +293,7 @@ public class TraceToGraphBuilder( shape = weightShape, dtype = weightDtype ).withTensorEncoding(encoding) + .withTensorId(refs.firstNotNullOfOrNull { it.spec.tensorId }) syntheticNode = GraphNode( id = nodeId, operation = op, @@ -346,6 +349,8 @@ public class TraceToGraphBuilder( val shape = shapes?.getOrNull(i) ?: effectiveInputs[i].shape.dimensions.toList() val dtype = dtypes?.getOrNull(i) ?: effectiveInputs[i].dtype::class.simpleName ?: "unknown" TensorSpec(name = name, shape = shape, dtype = dtype) + .withTensorEncoding(effectiveInputs[i].encoding) + .withTensorId(effectiveInputs[i].tensorId) } } @@ -358,6 +363,8 @@ public class TraceToGraphBuilder( val shape = shapes?.getOrNull(i) ?: trace.outputs[i].shape.dimensions.toList() val dtype = dtypes?.getOrNull(i) ?: trace.outputs[i].dtype::class.simpleName ?: "unknown" TensorSpec(name = name, shape = shape, dtype = dtype) + .withTensorEncoding(trace.outputs[i].encoding) + .withTensorId(trace.outputs[i].tensorId) } } diff --git a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/TraceIdentityToGraphTest.kt b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/TraceIdentityToGraphTest.kt new file mode 100644 index 000000000..7c34ebe14 --- /dev/null +++ b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/TraceIdentityToGraphTest.kt @@ -0,0 +1,67 @@ +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.TensorId +import sk.ainet.lang.tensor.data.DenseTensorDataFactory +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.tensor.ops.AddOperation +import sk.ainet.lang.tensor.ops.tensorEncoding +import sk.ainet.lang.tensor.ops.tensorId +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * #1178: what the tape captured survives to the graph. The trace→graph handoff used to drop all + * metadata — a packed weight arrived downstream as a name string, a parameter as `t7`. Now the + * `TensorSpec`s carry the encoding *object* and the module-path identity through the same + * untyped-metadata mechanism `tensorEncoding` proved. + */ +class TraceIdentityToGraphTest { + + private fun ctx(): DefaultGraphExecutionContext { + val dataFactory = DenseTensorDataFactory() + return DefaultGraphExecutionContext( + baseOps = DefaultCpuOps(dataFactory), + phase = Phase.TRAIN, + tensorDataFactory = dataFactory, + createTapeFactory = { _ -> DefaultGradientTape(true) }, + ) + } + + @Test + fun identityAndEncodingSurviveToTheGraphSpecs() { + val trainCtx = ctx() + val id = TensorId(listOf("model", "blk[0]"), "ffn.weight") + + 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) + trainCtx.session.identify(weight, id) + + val dense = trainCtx.fromFloatArray(Shape(32), FP32::class, FloatArray(32) { it * 0.5f }) + + val tape = DefaultExecutionTape(trainCtx.session) + tape.startRecording() + tape.recordOperation(AddOperation(), listOf(weight, dense), listOf(dense)) + tape.stopRecording() + + val graph = tape.toComputeGraph(synthesizeExternalInputs = true) + val allSpecs = graph.nodes.flatMap { it.inputs + it.outputs } + graph.edges.map { it.tensorSpec } + + val identified = allSpecs.filter { it.tensorId == id } + assertTrue(identified.isNotEmpty(), "the weight's TensorId must reach the graph, got specs: ${allSpecs.map { it.name }}") + assertEquals( + TensorEncoding.Q8_0, + identified.first { it.tensorEncoding != null }.tensorEncoding, + "the encoding object — block size intact — must ride along", + ) + } +} diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt index 661fdc066..47c06f4a9 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt @@ -62,6 +62,14 @@ public object HloGenerator { ctx: DefaultGraphExecutionContext ) { val module = model.create(ctx) + // Register each parameter's module-path identity on the trace session before recording + // (#1178): the tensor itself does not know which parameter it is, and refs are immutable + // once created, so this must happen before the forward pass touches them. Unparsable + // names stay unidentified rather than guessed. + for (parameter in module.trainableParameters()) { + val parsed = runCatching { sk.ainet.lang.tensor.TensorId.parse(parameter.name) }.getOrNull() + if (parsed != null) ctx.session.identify(parameter.value, parsed) + } model.calculate( module = module, inputValue = sampleInput, 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 e687e0d46..5ac226357 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 @@ -6563,6 +6563,12 @@ public final class sk/ainet/lang/tensor/ops/TensorSpec { public fun toString ()Ljava/lang/String; } +public final class sk/ainet/lang/tensor/ops/TensorSpecIdentities { + public static final field TENSOR_ID_METADATA_KEY Ljava/lang/String; + public static final fun getTensorId (Lsk/ainet/lang/tensor/ops/TensorSpec;)Lsk/ainet/lang/tensor/TensorId; + public static final fun withTensorId (Lsk/ainet/lang/tensor/ops/TensorSpec;Lsk/ainet/lang/tensor/TensorId;)Lsk/ainet/lang/tensor/ops/TensorSpec; +} + public final class sk/ainet/lang/tensor/ops/TensorSpecs { public static final field TENSOR_ENCODING_METADATA_KEY Ljava/lang/String; public static final fun getTensorEncoding (Lsk/ainet/lang/tensor/ops/TensorSpec;)Lsk/ainet/lang/tensor/storage/TensorEncoding; @@ -7796,16 +7802,21 @@ public final class sk/ainet/lang/trace/OpTrace { } public final class sk/ainet/lang/trace/TensorRef { - public fun (Ljava/lang/String;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;)V + public fun (Ljava/lang/String;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/tensor/storage/TensorEncoding;)V + public synthetic fun (Ljava/lang/String;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/tensor/storage/TensorEncoding;ILkotlin/jvm/internal/DefaultConstructorMarker;)V public final fun component1 ()Ljava/lang/String; public final fun component2 ()Lsk/ainet/lang/tensor/Shape; public final fun component3 ()Lsk/ainet/lang/types/DType; - public final fun copy (Ljava/lang/String;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/trace/TensorRef; - public static synthetic fun copy$default (Lsk/ainet/lang/trace/TensorRef;Ljava/lang/String;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;ILjava/lang/Object;)Lsk/ainet/lang/trace/TensorRef; + public final fun component4 ()Lsk/ainet/lang/tensor/TensorId; + public final fun component5 ()Lsk/ainet/lang/tensor/storage/TensorEncoding; + public final fun copy (Ljava/lang/String;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/tensor/storage/TensorEncoding;)Lsk/ainet/lang/trace/TensorRef; + public static synthetic fun copy$default (Lsk/ainet/lang/trace/TensorRef;Ljava/lang/String;Lsk/ainet/lang/tensor/Shape;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/tensor/TensorId;Lsk/ainet/lang/tensor/storage/TensorEncoding;ILjava/lang/Object;)Lsk/ainet/lang/trace/TensorRef; public fun equals (Ljava/lang/Object;)Z public final fun getDtype ()Lsk/ainet/lang/types/DType; + public final fun getEncoding ()Lsk/ainet/lang/tensor/storage/TensorEncoding; public final fun getId ()Ljava/lang/String; public final fun getShape ()Lsk/ainet/lang/tensor/Shape; + public final fun getTensorId ()Lsk/ainet/lang/tensor/TensorId; public fun hashCode ()I public fun toString ()Ljava/lang/String; } @@ -7813,6 +7824,7 @@ public final class sk/ainet/lang/trace/TensorRef { public class sk/ainet/lang/trace/TraceSession { public fun ()V public final fun clear ()V + public fun identify (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/TensorId;)V public fun refOf (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/trace/TensorRef; public final fun refsOf (Ljava/util/List;)Ljava/util/List; public fun resolve (Ljava/lang/String;)Lsk/ainet/lang/tensor/Tensor; diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorSpecIdentity.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorSpecIdentity.kt new file mode 100644 index 000000000..5252b682a --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorSpecIdentity.kt @@ -0,0 +1,36 @@ +@file:JvmName("TensorSpecIdentities") + +package sk.ainet.lang.tensor.ops + +import sk.ainet.lang.tensor.TensorId +import kotlin.jvm.JvmName + +/** + * Metadata key used to carry a [TensorId] on a [TensorSpec] (#1178). + * + * Same contract as [TENSOR_ENCODING_METADATA_KEY]: an untyped metadata entry with typed + * accessors, so the compile pipeline carries the fact without importing storage-model types — + * the `TensorSpecEncoding` precedent, and the standing rule for everything the tape captures. + */ +public const val TENSOR_ID_METADATA_KEY: String = "tensorId" + +/** + * The module-path identity carried on this spec (`model.layers[3].attn.q_proj.weight`), or + * `null` if no producer knew it. `null` means "unidentified", not "anonymous by design" — + * consumers keying per-weight policy on identity should skip unidentified tensors. + */ +public val TensorSpec.tensorId: TensorId? + get() = metadata[TENSOR_ID_METADATA_KEY] as? TensorId + +/** + * Return a copy of this spec with [id] stored in its metadata map. Passing `null` removes the + * entry; a non-null value adds or replaces it, leaving all other metadata untouched. + */ +public fun TensorSpec.withTensorId(id: TensorId?): TensorSpec { + val newMetadata: Map = if (id == null) { + metadata - TENSOR_ID_METADATA_KEY + } else { + metadata + (TENSOR_ID_METADATA_KEY to id) + } + return copy(metadata = newMetadata) +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TensorRef.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TensorRef.kt index 84d630b42..6b1387a54 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TensorRef.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TensorRef.kt @@ -1,14 +1,24 @@ package sk.ainet.lang.trace import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.storage.TensorEncoding import sk.ainet.lang.types.DType /** * TensorRef represents a reference to a tensor in the computation graph. * It captures the essential metadata needed for graph construction and optimization. + * + * [tensorId] and [encoding] carry what only the tape can see (#1178): the live tensor's + * module-path identity (when someone who knows it called [TraceSession.identify]) and its + * physical storage encoding as an *object* — block size intact, unlike the display-name + * strings downstream stages used to be left with. Both default to `null`, so every existing + * construction and destructuring compiles unchanged. */ public data class TensorRef( val id: String, val shape: Shape, - val dtype: DType -) \ No newline at end of file + val dtype: DType, + val tensorId: TensorId? = null, + val encoding: TensorEncoding? = null, +) diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TraceSession.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TraceSession.kt index fb610fd88..3a64150b4 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TraceSession.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/trace/TraceSession.kt @@ -1,6 +1,8 @@ package sk.ainet.lang.trace import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.ops.inferTensorEncoding import sk.ainet.lang.types.* /** @@ -11,6 +13,19 @@ public open class TraceSession { private var nextId = 0 private val tensorToRef = mutableMapOf() private val refToId = mutableMapOf>() + private val identities = mutableMapOf() + + /** + * Register [tensor]'s module-path identity so its [TensorRef] carries it (#1178). + * + * A tensor does not know which parameter it is — the module that owns it does. Whoever + * holds that knowledge (e.g. the HLO generator walking `trainableParameters()` before + * recording) calls this *before* the tensor's first [refOf]; refs are immutable and cached, + * so identities registered later do not retrofit existing refs. + */ + public open fun identify(tensor: Tensor<*, *>, id: TensorId) { + identities[unwrap(tensor)] = id + } /** * Get or create a TensorRef for the given tensor. @@ -35,7 +50,9 @@ public open class TraceSession { val ref = TensorRef( id = "t${nextId++}", shape = tensor.shape, - dtype = dtypeInstance + dtype = dtypeInstance, + tensorId = identities[key], + encoding = tensor.data.inferTensorEncoding(), ) refToId[ref.id] = tensor ref diff --git a/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/trace/TraceSessionIdentityTest.kt b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/trace/TraceSessionIdentityTest.kt new file mode 100644 index 000000000..7d0d1b421 --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/trace/TraceSessionIdentityTest.kt @@ -0,0 +1,57 @@ +package sk.ainet.lang.trace + +import sk.ainet.context.DefaultDataExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.TensorId +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.types.FP32 +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNull + +/** + * #1178: the tape captures what only it can see — the live tensor's storage encoding as an + * object (block size intact) and, when someone who knows it registered one, its module-path + * identity. Refs are immutable and cached, so identity must be registered before first use. + */ +class TraceSessionIdentityTest { + + private val ctx = DefaultDataExecutionContext() + + @Test + fun refCapturesThePackedEncodingObject() { + val packed = Q8_0BlockTensorData.fromRawBytes(Shape(32), ByteArray(34)) + @Suppress("UNCHECKED_CAST") + val tensor = ctx.fromData(packed as sk.ainet.lang.tensor.data.TensorData, FP32::class) + val ref = TraceSession().refOf(tensor) + assertEquals(TensorEncoding.Q8_0, ref.encoding, "the encoding object, not a display name") + } + + @Test + fun denseTensorHasNoEncodingAndNoIdentityByDefault() { + val tensor = ctx.fromFloatArray(Shape(4), FP32::class, floatArrayOf(1f, 2f, 3f, 4f)) + val ref = TraceSession().refOf(tensor) + assertNull(ref.encoding) + assertNull(ref.tensorId) + } + + @Test + fun identityRegisteredBeforeFirstUseIsCarried() { + val session = TraceSession() + val weight = ctx.fromFloatArray(Shape(2, 2), FP32::class, FloatArray(4)) + val id = TensorId(listOf("model", "layers[3]", "attn"), "q_proj.weight") + session.identify(weight, id) + assertEquals(id, session.refOf(weight).tensorId) + } + + @Test + fun identityRegisteredAfterFirstUseDoesNotRetrofitTheCachedRef() { + val session = TraceSession() + val weight = ctx.fromFloatArray(Shape(2), FP32::class, FloatArray(2)) + val before = session.refOf(weight) + assertNull(before.tensorId) + session.identify(weight, TensorId(listOf("model"), "weight")) + assertNull(session.refOf(weight).tensorId, "refs are immutable and cached — identify before first refOf") + } +}