diff --git a/config/stub-allowlist.txt b/config/stub-allowlist.txt index c21183c7b..4c9dc91db 100644 --- a/config/stub-allowlist.txt +++ b/config/stub-allowlist.txt @@ -27,4 +27,3 @@ skainet-io/skainet-io-iree-params/src/wasmJsMain/kotlin/sk/ainet/io/irpa/FileBac skainet-io/skainet-io-iree-params/src/wasmWasiMain/kotlin/sk/ainet/io/irpa/FileBackedTransfer.wasmwasi.kt|1|IRPA FileBacked transfer on wasmWasi (follow-up of #1423) # --- Lang core: DSL and tensor views --------------------------------------------------------- -skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt|6|SKEEP-008 (#1433) mixed-precision dense() / stage() diff --git a/docs/modules/ROOT/nav.adoc b/docs/modules/ROOT/nav.adoc index 4ce2f0986..aabdb2bef 100644 --- a/docs/modules/ROOT/nav.adoc +++ b/docs/modules/ROOT/nav.adoc @@ -44,6 +44,7 @@ ** xref:explanation/kernel-selection.adoc[Kernel SPI and the selection algorithm] ** xref:explanation/schedules.adoc[Algorithm and schedule] ** xref:explanation/quantization-process.adoc[The quantization process] +** xref:explanation/mixed-precision.adoc[Mixed precision] ** xref:explanation/theory/index.adoc[Mathematical theory] *** xref:explanation/theory/matmul.adoc[Matrix multiplication] ** xref:explanation/examples/index.adoc[Worked examples] diff --git a/docs/modules/ROOT/pages/explanation/mixed-precision.adoc b/docs/modules/ROOT/pages/explanation/mixed-precision.adoc new file mode 100644 index 000000000..0ce5203ac --- /dev/null +++ b/docs/modules/ROOT/pages/explanation/mixed-precision.adoc @@ -0,0 +1,141 @@ += Mixed Precision: Why a Layer's Dtype Is Worth Choosing +:description: Why SKaiNET lets one layer or one stage run in a narrower dtype than the rest of the network, what that buys (and doesn't buy yet), and how the PrecisionBoundary design keeps every cast visible. + +A network built with the {framework_name} DSL has one dtype: `sequential { … }` +fixes it for every layer. Mixed precision (SKEEP-008, tracked in +https://github.com/SKaiNET-developers/SKaiNET/issues/1433[#1433]) loosens that — one layer or +one block can run in another dtype while the network around it stays unchanged: + +[source,kotlin] +---- +val model = sequential { + input(784) + dense(256, dtype = BF16) { } // one BF16 layer in an FP32 network + activation { it.relu() } + stage("head", dtype = FP16) { // a whole FP16 block + dense(64) { } + dense(10) + } +} +---- + +This page explains why that is worth having, what it costs, and where the current +implementation stops. + +== Precision is not uniformly needed + +The observation behind every mixed-precision system — PyTorch AMP, JAX's dtype promotion, +ONNX Runtime's float16 transforms — is that a network's tolerance for rounding is wildly +uneven. The matmul-heavy bulk of a model (dense and convolution layers) tolerates 16-bit +floats well: their error statistics average out over thousands of multiply-accumulates. A few +places do not: softmax and normalization (small differences of large numbers), accumulators, +and the loss. Forcing the whole network to the precision its most sensitive operation needs +wastes memory and bandwidth on the 95% that doesn't care. + +Per-layer `dense(dtype = …)` and per-block `stage(id, dtype = …)` are exactly that dial: keep +the skeleton FP32, drop the bulk into a cheaper dtype, and keep the choice *local* — a +narrow layer is opt-in per call, never a network-wide default. + +== What the two 16-bit floats trade away + +[cols="1,2,2",options="header"] +|=== +| | *FP16* | *BF16* + +| Layout +| 1 sign, 5 exponent, 10 mantissa bits +| 1 sign, 8 exponent, 7 mantissa bits + +| Range +| overflows past 65 504 +| FP32's full range + +| Precision +| ~3 decimal digits +| ~2 decimal digits + +| Failure mode +| overflow to `Inf` in activations with large magnitudes +| coarse rounding; sums of many small values drift +|=== + +BF16 keeps FP32's exponent, so anything that fits in FP32 fits in BF16 — it trades mantissa +for safety and is usually the first thing to try. FP16 is more precise per bit but needs the +values to stay small. The `convert` codecs handle overflow (to infinity) and NaN correctly, +but a model can still *silently* degrade — which is why the dtype appears in the code at the +call site, not in a config file. + +== Why it matters for SKaiNET's targets + +{framework_name} runs on Kotlin Multiplatform targets where the usual mixed-precision +arguments are amplified: + +* *Memory-constrained hosts.* A browser tab (wasm), an Android app, an iOS app — environments + where halving or quartering a model's working set decides whether it runs at all. +* *Bandwidth-bound inference.* Most inference time is spent moving weights, not multiplying + them. Fewer bytes per element means fewer bytes through the cache hierarchy per token. +* *Models that arrive mixed.* GGUF and ONNX models in the wild already mix precisions — + quantized matmul weights next to FP32 norms. A DSL that can only express one dtype has no + faithful home for them. + +[IMPORTANT] +==== +Today the memory and bandwidth benefits are *not yet realized*: FP16 and BF16 tensors are +float-backed (the values are rounded to the 16-bit format, the storage is still 4 bytes per +element). Real 2-byte storage belongs to the storage model (SKEEP-003). What mixed precision +gives you now is the *numerics* — a faithful simulation of the narrow format, the DSL surface, +and exports that carry the casts — so models can be validated against 16-bit behaviour before +the storage catches up, and the API doesn't change when it does. +==== + +== The design: one visible boundary + +The DSL wraps a mixed-precision layer or stage in a single module, `PrecisionBoundary`. Its +forward pass is two casts around the wrapped module: + +[source,kotlin] +---- +val x = ctx.ops.convert(input, innerType) // cast in +ctx.ops.convert(inner.forward(x, ctx), outerType) // cast back out +---- + +Three properties follow from using `TensorOps.convert` — an ordinary tensor op — rather than +hiding the conversion inside the layer: + +* *Free when degenerate.* `convert` returns its input unchanged when the dtypes already + match, so a boundary between equal dtypes costs nothing. +* *Visible in traces and exports.* The tracing wrappers record `convert` like any other op, + and the StableHLO converter lowers it to `stablehlo.convert` — a compiled or exported graph + shows exactly where precision changes, and a later optimisation pass can fold adjacent + cast pairs. +* *Honest about cost.* Every boundary copies its input and its output. A network that + switches dtype every layer can be *slower* than all-FP32, especially while narrow dtypes + are float-backed. Few, coarse boundaries (a `stage` around a block) beat many fine ones. + +The dtype is passed as a *value* (`dense(256, dtype = BF16)`), not as a type argument. The +earlier API sketched `dense(256)`, but an interface method's type parameter cannot be +`reified`, so the implementation had no `KClass` to build the layer's tensors with — those +overloads could never be implemented and are now deprecated. Passing the `DType` object gives +the implementation both the instance (for `convert`) and `dtype::class` (for tensor creation). + +== What it does not do (yet) + +* *Training.* `convert` is not differentiable today, so gradients stop at a boundary: a + mixed-precision layer or stage is inference-only. SKEEP-008 phase A3 makes `convert` a + `@Diff` op, which lifts this. +* *Automatic mixed precision.* There is no autocast policy, no FP32 master weights, no loss + scaling. The dtype is always an explicit, local choice in the code. +* *Integer layers.* A cast to Int8 truncates; it is not quantization. Integer inference goes + through packed encodings and the quantized kernels — see + xref:explanation/quantization-process.adoc[The quantization process]. Mixed precision and + quantization are complementary, not competing: one changes the *float format* a layer + computes in, the other changes the *encoding* weights are stored and multiplied in. + +== Where to look + +* `sk.ainet.lang.nn.PrecisionBoundary` — the boundary module. +* `NeuralNetworkDsl.dense(outputDimension, dtype, …)` and `stage(id, dtype, …)` — the DSL + entry points. +* `TensorOps.convert` — the cast, including FP16/BF16 rounding through the format codecs. +* xref:skeep:008-precision-and-encoding-conversion.adoc[SKEEP-008] — the full design, + including the storage-side `repackTo`/`requantizeTo` split and the phased rollout. diff --git a/docs/modules/ROOT/pages/reference/operators/generated/index.adoc b/docs/modules/ROOT/pages/reference/operators/generated/index.adoc index 40f914442..81f7574e4 100644 --- a/docs/modules/ROOT/pages/reference/operators/generated/index.adoc +++ b/docs/modules/ROOT/pages/reference/operators/generated/index.adoc @@ -1,6 +1,6 @@ = AI-NET Operators Reference -Generated from version `0.62.4` on 2026-10-31 +Generated from version `0.62.4` on 2026-11-01 == Operators by Modality diff --git a/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc b/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc index e0bf3086c..cbf50c761 100644 --- a/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc +++ b/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc @@ -1,7 +1,7 @@ = Operator Coverage Matrix :description: Cross-backend status for every operator function in SKaiNET. -Generated from `operators.json` version `0.62.4` on 2026-10-31. +Generated from `operators.json` version `0.62.4` on 2026-11-01. Rows are `Operator.function` pairs. A gap is marked in source with `@NotImplemented(reason = …, issue = …)` (or `@InProgress`) and shows up here as ❌ (or 🚧) linked to its tracking issue; the PR that closes the gap removes the annotation. See xref:contributing/issue-taxonomy.adoc[Issue taxonomy] to pick one up. @@ -110,6 +110,7 @@ Every concrete `sk.ainet.lang.nn` `Module` subclass. ❌ means the layer's `forw | `Linear` | `sk.ainet.lang.nn` | ✅ | — | `Lstm` | `sk.ainet.lang.nn` | ✅ | — | `MaxPool2d` | `sk.ainet.lang.nn` | ✅ | — +| `PrecisionBoundary` | `sk.ainet.lang.nn` | ✅ | — | `TransposedConv1d` | `sk.ainet.lang.nn` | ✅ | — | `TransposedConv2d` | `sk.ainet.lang.nn` | ✅ | — | `Upsample2d` | `sk.ainet.lang.nn` | ✅ | — @@ -129,5 +130,5 @@ Every concrete `sk.ainet.lang.nn` `Module` subclass. ❌ means the layer's `forw | `LayerNormalization` | `sk.ainet.lang.nn.normalization` | ✅ | — | `MLP` | `sk.ainet.lang.nn.topology` | ✅ | — -| *Done* | | *32 / 32* | +| *Done* | | *33 / 33* | |=== diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index 69cd20ef6..f8ed5d828 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt @@ -406,7 +406,7 @@ public open class DefaultCpuOpsBase( floatBinaryFast(a, b) { x, y -> x + y }?.let { return it } return elementwise(a, b) { av, bv, dtype -> when (dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val x = av as Float; val y = bv as Float; (x + y) as V } @@ -429,7 +429,7 @@ public open class DefaultCpuOpsBase( floatBinaryFast(a, b) { x, y -> x - y }?.let { return it } return elementwise(a, b) { av, bv, dtype -> when (dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val x = av as Float; val y = bv as Float; (x - y) as V } @@ -452,7 +452,7 @@ public open class DefaultCpuOpsBase( floatBinaryFast(a, b) { x, y -> x * y }?.let { return it } return elementwise(a, b) { av, bv, dtype -> when (dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val x = av as Float; val y = bv as Float; (x * y) as V } @@ -475,7 +475,7 @@ public open class DefaultCpuOpsBase( floatBinaryFast(a, b) { x, y -> x / y }?.let { return it } return elementwise(a, b) { av, bv, dtype -> when (dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val x = av as Float; val y = bv as Float; (x / y) as V } @@ -886,7 +886,7 @@ public open class DefaultCpuOpsBase( val bBatchIdx = mapBatchIndexEff(batchIdx, bEff) when (a.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { var acc = 0.0f var k = 0 while (k < kA) { @@ -1525,7 +1525,7 @@ public open class DefaultCpuOpsBase( val ow = outIdx[3] when (input.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { var acc = 0.0f val groupIdx = (oc * groups) / cOut val inCStart = groupIdx * cInPerGroup @@ -1672,7 +1672,7 @@ public open class DefaultCpuOpsBase( val ol = outIdx[2] when (input.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { var acc = 0.0f val groupIdx = (oc * groups) / cOut val inCStart = groupIdx * cInPerGroup @@ -1830,7 +1830,11 @@ public open class DefaultCpuOpsBase( require(groups >= 1) { "convTranspose2d: groups must be >= 1" } require(input.dtype == weight.dtype) { "convTranspose2d: dtype mismatch between input and weight: ${input.dtype} vs ${weight.dtype}" } bias?.let { require(it.dtype == input.dtype) { "convTranspose2d: dtype mismatch for bias" } } - require(input.dtype == sk.ainet.lang.types.FP32::class || input.dtype == sk.ainet.lang.types.FP16::class) { + require( + input.dtype == sk.ainet.lang.types.FP32::class || + input.dtype == sk.ainet.lang.types.FP16::class || + input.dtype == sk.ainet.lang.types.BF16::class + ) { "Unsupported dtype for convTranspose2d: ${input.dtype}" } @@ -2004,7 +2008,7 @@ public open class DefaultCpuOpsBase( val ow = outIdx[4] when (input.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { var acc = 0.0f val groupIdx = (oc * groups) / cOut val inCStart = groupIdx * cInPerGroup @@ -2141,7 +2145,7 @@ public open class DefaultCpuOpsBase( } UpsampleMode.Bilinear -> { - require(input.dtype == FP32::class || input.dtype == FP16::class) { + require(input.dtype == FP32::class || input.dtype == FP16::class || input.dtype == BF16::class) { "upsample2d: Bilinear mode is only implemented for float dtypes (got ${input.dtype})" } dataFactory.init(outShape, input.dtype) { idx -> @@ -2216,7 +2220,7 @@ public open class DefaultCpuOpsBase( val hBase = oh * sH - pH val wBase = ow * sW - pW when (input.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { var best = Float.NEGATIVE_INFINITY var kh = 0 while (kh < kH) { @@ -2295,7 +2299,7 @@ public open class DefaultCpuOpsBase( val hBase = oh * sH - pH val wBase = ow * sW - pW when (input.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { var sum = 0.0f var count = 0 var kh = 0 @@ -2625,7 +2629,7 @@ public open class DefaultCpuOpsBase( floatUnaryFast(tensor) { x -> if (x < 0f) 0f else x }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val v = tensor.data.get(*idx) as Float @Suppress("UNCHECKED_CAST") (if (v < 0f) 0f else v) as V @@ -2646,7 +2650,7 @@ public open class DefaultCpuOpsBase( floatUnaryFast(tensor) { x -> if (x < 0f) negativeSlope * x else x }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val v = tensor.data.get(*idx) as Float @Suppress("UNCHECKED_CAST") (if (v < 0f) negativeSlope * v else v) as V @@ -2667,7 +2671,7 @@ public open class DefaultCpuOpsBase( floatUnaryFast(tensor) { x -> if (x >= 0f) x else alpha * (kotlin.math.exp(x) - 1f) }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val v = tensor.data.get(*idx) as Float @Suppress("UNCHECKED_CAST") (if (v >= 0f) v else alpha * (kotlin.math.exp(v) - 1f)) as V @@ -2710,7 +2714,7 @@ public open class DefaultCpuOpsBase( } } when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val outData = dataFactory.init(tensor.shape, tensor.dtype) { outIdx -> // For the slice defined by outIdx except for nd, compute softmax at that position val idxMax = outIdx.copyOf() @@ -2771,7 +2775,7 @@ public open class DefaultCpuOpsBase( } } when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val outData = dataFactory.init(tensor.shape, tensor.dtype) { outIdx -> val idx = outIdx.copyOf() var maxVal = Float.NEGATIVE_INFINITY @@ -2802,7 +2806,7 @@ public open class DefaultCpuOpsBase( floatUnaryFast(tensor) { x -> 1.0f / (1.0f + kotlin.math.exp(-x)) }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val x = tensor.data.get(*idx) as Float @Suppress("UNCHECKED_CAST") (1.0f / (1.0f + kotlin.math.exp(-x))) as V @@ -2818,7 +2822,7 @@ public open class DefaultCpuOpsBase( floatUnaryFast(tensor) { x -> x * (1.0f / (1.0f + kotlin.math.exp(-x))) }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val x = tensor.data.get(*idx) as Float val s = 1.0f / (1.0f + kotlin.math.exp(-x)) @Suppress("UNCHECKED_CAST") @@ -2840,7 +2844,7 @@ public open class DefaultCpuOpsBase( // Tanh approximation of GELU val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val x = tensor.data.get(*idx) as Float val x3 = x * x * x val inner = x + 0.044715f * x3 @@ -2861,7 +2865,7 @@ public open class DefaultCpuOpsBase( floatUnaryFast(tensor) { x -> ErfKernel.geluExact(x) }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val x = tensor.data.get(*idx) as Float @Suppress("UNCHECKED_CAST") ErfKernel.geluExact(x) as V @@ -2877,7 +2881,7 @@ public open class DefaultCpuOpsBase( floatUnaryFast(tensor) { x -> ErfKernel.erf(x.toDouble()).toFloat() }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val x = tensor.data.get(*idx) as Float @Suppress("UNCHECKED_CAST") ErfKernel.erf(x.toDouble()).toFloat() as V @@ -2907,7 +2911,7 @@ public open class DefaultCpuOpsBase( tensor.data.get(*outIdx) as V } else { when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { @Suppress("UNCHECKED_CAST") 0.0f as V } @@ -2963,7 +2967,7 @@ public open class DefaultCpuOpsBase( val outShape = Shape() val outData = dataFactory.init(outShape, tensor.dtype) { when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { var acc = 0.0f val dims = tensor.shape.dimensions if (dims.isEmpty()) { @@ -3030,7 +3034,7 @@ public open class DefaultCpuOpsBase( val outShape = Shape(outDims) val outData = dataFactory.init(outShape, tensor.dtype) { outIdx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { var acc = 0.0f val inIdx = IntArray(rank) var o = 0 @@ -3150,7 +3154,7 @@ public open class DefaultCpuOpsBase( val count = if (tensor.volume == 0) 0 else tensor.volume val outData = dataFactory.init(outShape, tensor.dtype) { when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { if (count == 0) { @Suppress("UNCHECKED_CAST") (0.0f as V) } else { @@ -3218,7 +3222,7 @@ public open class DefaultCpuOpsBase( val n = inDims[nd] val outData = dataFactory.init(outShape, tensor.dtype) { outIdx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { if (n == 0) { @Suppress("UNCHECKED_CAST") (0.0f as V) } else { @@ -3272,7 +3276,7 @@ public open class DefaultCpuOpsBase( val n = tensor.volume val outData = dataFactory.init(outShape, tensor.dtype) { when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { if (n == 0) { @Suppress("UNCHECKED_CAST") (0.0f as V) } else { @@ -3354,7 +3358,7 @@ public open class DefaultCpuOpsBase( val n = inDims[nd] val outData = dataFactory.init(outShape, tensor.dtype) { outIdx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { if (n == 0) { @Suppress("UNCHECKED_CAST") (0.0f as V) } else { @@ -3411,7 +3415,8 @@ public open class DefaultCpuOpsBase( override fun sqrt(tensor: Tensor): Tensor { require( tensor.dtype == sk.ainet.lang.types.FP32::class || - tensor.dtype == sk.ainet.lang.types.FP16::class + tensor.dtype == sk.ainet.lang.types.FP16::class || + tensor.dtype == sk.ainet.lang.types.BF16::class ) { "sqrt supports only FP16/FP32, got ${tensor.dtype}" } floatUnaryFast(tensor) { x -> sqrt(x) }?.let { return it } @@ -3432,7 +3437,8 @@ public open class DefaultCpuOpsBase( override fun pow(a: Tensor, b: Tensor): Tensor { require( a.dtype == sk.ainet.lang.types.FP32::class || - a.dtype == sk.ainet.lang.types.FP16::class + a.dtype == sk.ainet.lang.types.FP16::class || + a.dtype == sk.ainet.lang.types.BF16::class ) { "pow supports only FP16/FP32, got ${a.dtype}" } require(a.shape == b.shape) { "pow requires matching shapes; got ${a.shape} and ${b.shape}" } floatBinaryFast(a, b) { x, y -> scalarPow(x, y) }?.let { return it } @@ -3453,7 +3459,8 @@ public open class DefaultCpuOpsBase( override fun powScalar(a: Tensor, n: Number): Tensor { require( a.dtype == sk.ainet.lang.types.FP32::class || - a.dtype == sk.ainet.lang.types.FP16::class + a.dtype == sk.ainet.lang.types.FP16::class || + a.dtype == sk.ainet.lang.types.BF16::class ) { "powScalar supports only FP16/FP32, got ${a.dtype}" } val nFloat = n.toFloat() val nInt = n.toInt() @@ -3493,7 +3500,8 @@ public open class DefaultCpuOpsBase( override fun log(tensor: Tensor): Tensor { require( tensor.dtype == sk.ainet.lang.types.FP32::class || - tensor.dtype == sk.ainet.lang.types.FP16::class + tensor.dtype == sk.ainet.lang.types.FP16::class || + tensor.dtype == sk.ainet.lang.types.BF16::class ) { "log supports only FP16/FP32, got ${tensor.dtype}" } floatUnaryFast(tensor) { x -> ln(x) }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> @@ -3508,7 +3516,8 @@ public open class DefaultCpuOpsBase( override fun log2(tensor: Tensor): Tensor { require( tensor.dtype == sk.ainet.lang.types.FP32::class || - tensor.dtype == sk.ainet.lang.types.FP16::class + tensor.dtype == sk.ainet.lang.types.FP16::class || + tensor.dtype == sk.ainet.lang.types.BF16::class ) { "log2 supports only FP16/FP32, got ${tensor.dtype}" } floatUnaryFast(tensor) { x -> kmLog2(x) }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> @@ -3523,7 +3532,8 @@ public open class DefaultCpuOpsBase( override fun log10(tensor: Tensor): Tensor { require( tensor.dtype == sk.ainet.lang.types.FP32::class || - tensor.dtype == sk.ainet.lang.types.FP16::class + tensor.dtype == sk.ainet.lang.types.FP16::class || + tensor.dtype == sk.ainet.lang.types.BF16::class ) { "log10 supports only FP16/FP32, got ${tensor.dtype}" } floatUnaryFast(tensor) { x -> kmLog10(x) }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> @@ -3541,7 +3551,7 @@ public open class DefaultCpuOpsBase( floatUnaryFast(tensor) { x -> kotlin.math.abs(x) }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val v = tensor.data.get(*idx) as Float @Suppress("UNCHECKED_CAST") kotlin.math.abs(v) as V @@ -3562,7 +3572,7 @@ public open class DefaultCpuOpsBase( floatUnaryFast(tensor) { x -> if (x > 0f) 1f else if (x < 0f) -1f else 0f }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val v = tensor.data.get(*idx) as Float @Suppress("UNCHECKED_CAST") (if (v > 0f) 1f else if (v < 0f) -1f else 0f) as V @@ -3584,7 +3594,7 @@ public open class DefaultCpuOpsBase( floatUnaryFast(tensor) { x -> x.coerceIn(minVal, maxVal) }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val v = tensor.data.get(*idx) as Float @Suppress("UNCHECKED_CAST") v.coerceIn(minVal, maxVal) as V @@ -3605,7 +3615,7 @@ public open class DefaultCpuOpsBase( floatUnaryFast(tensor) { x -> if (x < value) 1f else 0f }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val v = tensor.data.get(*idx) as Float @Suppress("UNCHECKED_CAST") (if (v < value) 1f else 0f) as V @@ -3626,7 +3636,7 @@ public open class DefaultCpuOpsBase( floatUnaryFast(tensor) { x -> if (x >= value) 1f else 0f }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val v = tensor.data.get(*idx) as Float @Suppress("UNCHECKED_CAST") (if (v >= value) 1f else 0f) as V @@ -3677,7 +3687,7 @@ public open class DefaultCpuOpsBase( } else { // Zero padding when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { @Suppress("UNCHECKED_CAST") (0f as V) } sk.ainet.lang.types.Int32::class -> { @@ -3919,7 +3929,7 @@ public open class DefaultCpuOpsBase( floatUnaryFast(tensor) { x -> kotlin.math.tanh(x) }?.let { return it } val outData = dataFactory.init(tensor.shape, tensor.dtype) { idx -> when (tensor.dtype) { - sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> { + sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class, sk.ainet.lang.types.BF16::class -> { val x = tensor.data.get(*idx) as Float @Suppress("UNCHECKED_CAST") kotlin.math.tanh(x) as V diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/MixedPrecisionDslTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/MixedPrecisionDslTest.kt new file mode 100644 index 000000000..8a3e34dc6 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/lang/nn/dsl/MixedPrecisionDslTest.kt @@ -0,0 +1,168 @@ +package sk.ainet.lang.nn.dsl + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.ExecutionContext +import sk.ainet.context.data +import sk.ainet.lang.nn.Module +import sk.ainet.lang.nn.PrecisionBoundary +import sk.ainet.lang.nn.definition +import sk.ainet.lang.nn.network +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.dsl.tensor +import sk.ainet.lang.types.BF16 +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +/** + * SKEEP-008 phase A1 (#1433): `dense(dtype = …)` and `stage(id, dtype = …)` run their layers + * in another dtype behind a [PrecisionBoundary], casting in and out with `TensorOps.convert`. + */ +class MixedPrecisionDslTest { + + private val ctx = DirectCpuExecutionContext() + + // Values deliberately not representable in FP16/BF16, so the boundary's rounding is observable. + private val w1 = floatArrayOf( + 0.1234567f, -1.7654321f, + 2.3456789f, 0.0123456f, + -0.9876543f, 1.1111111f, + ) + private val b1 = floatArrayOf(0.1111111f, -0.2222222f, 0.3333333f) + private val w2 = floatArrayOf( + 1.9999999f, -0.1111111f, 0.5555555f, + -1.3333333f, 0.7777777f, -0.4444444f, + ) + private val b2 = floatArrayOf(0.0101010f, -0.0202020f) + + private fun input(): Tensor = data(ctx) { + tensor { + shape(1, 2) { fromArray(floatArrayOf(0.3f, -1.1f)) } + } + } + + private fun fp32Reference(context: ExecutionContext): Module = definition { + network(context) { + input(2) + dense(3, "layer1") { + weights { fromArray(w1) } + bias { fromArray(b1) } + } + dense(2, "layer2") { + weights { fromArray(w2) } + bias { fromArray(b2) } + } + } + } + + @Test + fun denseWithDtype_runsAndRoundsThroughBf16() { + val mixed = definition { + network(ctx) { + input(2) + dense(3, dtype = BF16, id = "layer1") { + weights { fromArray(w1) } + bias { fromArray(b1) } + } + dense(2, "layer2") { + weights { fromArray(w2) } + bias { fromArray(b2) } + } + } + } + + val reference = fp32Reference(ctx).forward(input(), ctx) + val result = mixed.forward(input(), ctx) + + assertEquals(Shape(1, 2), result.shape) + assertEquals(FP32::class, result.dtype, "the boundary must cast back to the network dtype") + var maxDiff = 0f + for (c in 0 until 2) { + val diff = abs(result.data[0, c] - reference.data[0, c]) + assertTrue(diff <= 2e-2f, "BF16 layer drifted too far at $c: diff=$diff") + maxDiff = maxOf(maxDiff, diff) + } + assertTrue(maxDiff > 0f, "BF16 rounding must change the result, or the cast never happened") + } + + @Test + fun stageWithDtype_runsAndRoundsThroughFp16() { + val mixed = definition { + network(ctx) { + input(2) + stage("body", dtype = FP16) { + dense(3, "layer1") { + weights { fromArray(w1) } + bias { fromArray(b1) } + } + dense(2, "layer2") { + weights { fromArray(w2) } + bias { fromArray(b2) } + } + } + } + } + + val reference = fp32Reference(ctx).forward(input(), ctx) + val result = mixed.forward(input(), ctx) + + assertEquals(Shape(1, 2), result.shape) + assertEquals(FP32::class, result.dtype, "the boundary must cast back to the network dtype") + var maxDiff = 0f + for (c in 0 until 2) { + val diff = abs(result.data[0, c] - reference.data[0, c]) + assertTrue(diff <= 2e-2f, "FP16 stage drifted too far at $c: diff=$diff") + maxDiff = maxOf(maxDiff, diff) + } + assertTrue(maxDiff > 0f, "FP16 rounding must change the result, or the cast never happened") + } + + @Test + fun sameDtypeBoundary_isExact() { + // convert returns its input when the dtype matches, so an FP32 stage inside an + // FP32 network reproduces the reference bit-exactly. + val wrapped = definition { + network(ctx) { + input(2) + stage("body", dtype = FP32) { + dense(3, "layer1") { + weights { fromArray(w1) } + bias { fromArray(b1) } + } + dense(2, "layer2") { + weights { fromArray(w2) } + bias { fromArray(b2) } + } + } + } + } + + val reference = fp32Reference(ctx).forward(input(), ctx) + val result = wrapped.forward(input(), ctx) + + for (c in 0 until 2) { + assertEquals(reference.data[0, c], result.data[0, c], "FP32 boundary must be a no-op at $c") + } + } + + @Test + fun mixedLayers_appearAsPrecisionBoundaryModules() { + val mixed = definition { + network(ctx) { + input(2) + dense(3, dtype = BF16, id = "layer1") { + weights { fromArray(w1) } + bias { fromArray(b1) } + } + } + } + + val boundaries = mixed.modules.filterIsInstance>() + assertEquals(1, boundaries.size, "the BF16 dense must be wrapped in one PrecisionBoundary") + assertEquals("layer1", boundaries.single().name) + } +} 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 680774298..2974404cb 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 @@ -2997,6 +2997,12 @@ public final class sk/ainet/lang/nn/NeuralNetworkExecutionContextKt { public static final fun definition (Lkotlin/jvm/functions/Function2;)Lsk/ainet/lang/nn/Module; } +public final class sk/ainet/lang/nn/PrecisionBoundary : sk/ainet/lang/nn/Module { + public fun (Ljava/lang/String;Lsk/ainet/lang/nn/Module;Lsk/ainet/lang/types/DType;Lsk/ainet/lang/types/DType;)V + public fun getModules ()Ljava/util/List; + public fun getName ()Ljava/lang/String; +} + public final class sk/ainet/lang/nn/TrainingRunnerKt { public static final fun trainStep (Lsk/ainet/lang/nn/Module;Lsk/ainet/lang/nn/loss/Loss;Lsk/ainet/lang/nn/optim/Optimizer;Lsk/ainet/context/ExecutionContext;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; } @@ -3526,10 +3532,12 @@ public abstract interface class sk/ainet/lang/nn/dsl/NeuralNetworkDsl : sk/ainet public static synthetic fun conv3d$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;IZLjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public abstract fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public abstract fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)V + public abstract fun dense (ILsk/ainet/lang/types/DType;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public abstract fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public abstract fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lsk/ainet/lang/nn/Module; public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V + public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILsk/ainet/lang/types/DType;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lsk/ainet/lang/nn/Module; public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public abstract fun flatten (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V @@ -3553,6 +3561,7 @@ public abstract interface class sk/ainet/lang/nn/dsl/NeuralNetworkDsl : sk/ainet public static synthetic fun softmax$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILjava/lang/String;ILjava/lang/Object;)V public abstract fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public abstract fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V + public abstract fun stage (Ljava/lang/String;Lsk/ainet/lang/types/DType;Lkotlin/jvm/functions/Function1;)V public abstract fun upsample2d (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public abstract fun upsample2d (Lkotlin/Pair;Lsk/ainet/lang/tensor/ops/UpsampleMode;ZLjava/lang/String;)V public static synthetic fun upsample2d$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V @@ -3569,6 +3578,7 @@ public final class sk/ainet/lang/nn/dsl/NeuralNetworkDsl$DefaultImpls { public static synthetic fun conv3d$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;IZLjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lsk/ainet/lang/nn/Module; public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILjava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V + public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;ILsk/ainet/lang/types/DType;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)Lsk/ainet/lang/nn/Module; public static synthetic fun dense$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V public static synthetic fun flatten$default (Lsk/ainet/lang/nn/dsl/NeuralNetworkDsl;Ljava/lang/String;Lkotlin/jvm/functions/Function1;ILjava/lang/Object;)V @@ -3597,6 +3607,7 @@ public final class sk/ainet/lang/nn/dsl/NeuralNetworkDslImpl : sk/ainet/lang/nn/ public final fun create ()Lsk/ainet/lang/nn/Module; public fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)V + public fun dense (ILsk/ainet/lang/types/DType;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public fun flatten (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V @@ -3618,6 +3629,7 @@ public final class sk/ainet/lang/nn/dsl/NeuralNetworkDslImpl : sk/ainet/lang/nn/ public fun softmax (ILjava/lang/String;)V public fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V + public fun stage (Ljava/lang/String;Lsk/ainet/lang/types/DType;Lkotlin/jvm/functions/Function1;)V public fun upsample2d (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public fun upsample2d (Lkotlin/Pair;Lsk/ainet/lang/tensor/ops/UpsampleMode;ZLjava/lang/String;)V } @@ -3638,6 +3650,7 @@ public final class sk/ainet/lang/nn/dsl/StageImpl : sk/ainet/lang/nn/dsl/NeuralN public final fun create ()Lsk/ainet/lang/nn/Module; public fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public fun dense (ILjava/lang/String;Lkotlin/jvm/functions/Function1;)V + public fun dense (ILsk/ainet/lang/types/DType;Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public fun dense (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public fun flatten (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V @@ -3661,6 +3674,7 @@ public final class sk/ainet/lang/nn/dsl/StageImpl : sk/ainet/lang/nn/dsl/NeuralN public fun softmax (ILjava/lang/String;)V public fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)Lsk/ainet/lang/nn/Module; public fun stage (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V + public fun stage (Ljava/lang/String;Lsk/ainet/lang/types/DType;Lkotlin/jvm/functions/Function1;)V public fun upsample2d (Ljava/lang/String;Lkotlin/jvm/functions/Function1;)V public fun upsample2d (Lkotlin/Pair;Lsk/ainet/lang/tensor/ops/UpsampleMode;ZLjava/lang/String;)V } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/MixedPrecisionModule.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/MixedPrecisionModule.kt index 818993cab..4e7afd561 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/MixedPrecisionModule.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/MixedPrecisionModule.kt @@ -39,6 +39,10 @@ import sk.ainet.lang.types.DType * } * ``` */ +@Deprecated( + message = "Never implemented; replaced by PrecisionBoundary (SKEEP-008, #1433). Will be removed.", + replaceWith = ReplaceWith("PrecisionBoundary", "sk.ainet.lang.nn.PrecisionBoundary"), +) public abstract class MixedPrecisionModule( protected val inputType: TInput, protected val outputType: TOutput, @@ -210,6 +214,10 @@ public abstract class MixedPrecisionModule( * @param TInternal The internal computation precision type * @param V The value type corresponding to the DType */ +@Deprecated( + message = "Never implemented; replaced by PrecisionBoundary (SKEEP-008, #1433). Will be removed.", + replaceWith = ReplaceWith("PrecisionBoundary", "sk.ainet.lang.nn.PrecisionBoundary"), +) public abstract class InternalMixedPrecisionModule( private val ioType: TIO, private val internalType: TInternal, diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/PrecisionBoundary.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/PrecisionBoundary.kt new file mode 100644 index 000000000..fa83e6cca --- /dev/null +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/PrecisionBoundary.kt @@ -0,0 +1,37 @@ +package sk.ainet.lang.nn + +import sk.ainet.context.ExecutionContext +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.DType + +/** + * Runs [inner], a module built in another dtype, inside a network of dtype [TOuter]: + * the input is cast to [innerType], [inner] runs, and the result is cast back to + * [outerType] (SKEEP-008, #1433). + * + * The casts use [sk.ainet.lang.tensor.ops.TensorOps.convert], which returns its input + * unchanged when the dtypes already match, so a boundary between equal dtypes costs + * nothing. `convert` is recorded by the tracing wrappers and lowers to + * `stablehlo.convert`, so the boundary stays visible in traces and exports. + * + * `convert` is not differentiable (SKEEP-008 phase A3 lifts this), so gradients stop at + * the boundary: the wrapped layers do not train. + */ +public class PrecisionBoundary( + override val name: String, + private val inner: Module, + private val innerType: TInner, + private val outerType: TOuter, +) : Module() { + + @Suppress("UNCHECKED_CAST") + override val modules: List> + // The dtype parameter is erased at runtime; the inner module is listed so that + // reflection, parameter collection and checkpointing still see its subtree. + get() = listOf(inner as Module) + + override fun onForward(input: Tensor, ctx: ExecutionContext): Tensor { + val x = ctx.ops.convert(input, innerType) + return ctx.ops.convert(inner.forward(x, ctx), outerType) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt index bab5b93ef..fb910c4d0 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/dsl/NetworkBuilder.kt @@ -11,6 +11,7 @@ import sk.ainet.lang.nn.Input import sk.ainet.lang.nn.Linear import sk.ainet.lang.nn.MaxPool2d import sk.ainet.lang.nn.Module +import sk.ainet.lang.nn.PrecisionBoundary import sk.ainet.lang.nn.Upsample2d import sk.ainet.lang.nn.normalization.BatchNormalization import sk.ainet.lang.nn.normalization.GroupNormalization @@ -166,27 +167,39 @@ public interface NeuralNetworkDsl : NetworkDslItem { public fun gru(hiddenSize: Int, id: String = "", content: GRU.() -> Unit = {}) /** - * Creates a dense layer with precision override and specified output dimension. - * This allows individual layers to use different precision than the network default. + * Creates a dense layer whose weights and computation use [dtype] while the network stays + * in its own dtype: the layer's input is cast to [dtype] and its output cast back + * (SKEEP-008, #1433). The cast is `TensorOps.convert`, which is not differentiable yet, + * so a mixed-precision layer is inference-only: gradients stop at its boundary. * - * @param TLayer The precision type for this specific layer * @param outputDimension The number of neurons/output features + * @param dtype The dtype the layer's weights and computation use (e.g. `FP16`, `BF16`) * @param id Optional identifier for the layer - * @param content Configuration block for weights, bias, and activation + * @param content Configuration block for weights, bias, and activation, typed in [dtype] */ + public fun dense( + outputDimension: Int, + dtype: TLayer, + id: String = "", + content: DENSE.() -> Unit = {}, + ) + + @Deprecated( + message = "The dtype type argument cannot be implemented (no KClass at runtime, " + + "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }", + level = DeprecationLevel.ERROR, + ) public fun dense( outputDimension: Int, id: String = "", content: DENSE.() -> Unit = {} ): Module - /** - * Creates a dense layer with precision override without specifying output dimension. - * - * @param TLayer The precision type for this specific layer - * @param id Optional identifier for the layer - * @param content Configuration block where output dimension, weights, and bias are set - */ + @Deprecated( + message = "The dtype type argument cannot be implemented (no KClass at runtime, " + + "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }", + level = DeprecationLevel.ERROR, + ) public fun dense( id: String = "", content: DENSE.() -> Unit = {} @@ -438,18 +451,26 @@ public interface NeuralNetworkDsl : NetworkDslItem { public fun stage(id: String, content: NeuralNetworkDsl.() -> Unit) /** - * Creates a precision-scoped stage within the network. - * This allows grouping layers with a specific precision type that differs - * from the network default, enabling fine-grained mixed-precision control. - * - * Unlike the basic stage method, this version allows changing the precision - * type for all layers within the stage scope. + * Creates a stage whose layers all use [dtype] while the network stays in its own dtype: + * the stage's input is cast to [dtype] and its output cast back (SKEEP-008, #1433). + * The cast is `TensorOps.convert`, which is not differentiable yet, so a mixed-precision + * stage is inference-only: gradients stop at its boundary. * - * @param TStage The precision type for all layers within this stage * @param id Identifier for the stage - * @param content DSL block containing layers with TStage precision - * @return Module that handles the precision conversion automatically + * @param dtype The dtype all layers in this stage use (e.g. `FP16`, `BF16`) + * @param content DSL block containing the stage's layers, typed in [dtype] */ + public fun stage( + id: String, + dtype: TStage, + content: NeuralNetworkDsl.() -> Unit, + ) + + @Deprecated( + message = "The dtype type argument cannot be implemented (no KClass at runtime, " + + "SKEEP-008); pass the dtype as a value: stage(id, dtype = FP16) { ... }", + level = DeprecationLevel.ERROR, + ) public fun stage( id: String, content: NeuralNetworkDsl.() -> Unit @@ -1321,28 +1342,90 @@ public class StageImpl( override fun dense( outputDimension: Int, + dtype: TLayer, id: String, - content: DENSE.() -> Unit - ): Module { - // Create a mixed-precision module that handles conversion - TODO("Mixed-precision dense implementation needed") + content: DENSE.() -> Unit, + ) { + val inputDimension = lastDimension + lastDimension = outputDimension + currentShape = intArrayOf(outputDimension) + val layerName = getDefaultName(id, "linear", modules.size) + val impl = DenseImpl( + executionContext, + inputDimension = inputDimension, + _outputDimension = outputDimension, + id = layerName, + kClass = dtype.witnessOf(), + ) + impl.content() + val innerModules = impl.create() + val inner = if (innerModules.size == 1) innerModules[0] else MLP(*innerModules.toTypedArray(), name = layerName) + modules += PrecisionBoundary( + name = layerName, + inner = inner, + innerType = dtype, + outerType = dtypeInstanceOf(kClass), + ) } + override fun stage( + id: String, + dtype: TStage, + content: NeuralNetworkDsl.() -> Unit, + ) { + val stageImpl = StageImpl(executionContext, id, dtype.witnessOf()) + stageImpl.lastDimension = lastDimension + stageImpl.currentShape = currentShape?.copyOf() + stageImpl.content() + lastDimension = stageImpl.lastDimension + currentShape = stageImpl.currentShape?.copyOf() + modules += PrecisionBoundary( + name = id, + inner = stageImpl.create(), + innerType = dtype, + outerType = dtypeInstanceOf(kClass), + ) + } + + @Deprecated( + message = "The dtype type argument cannot be implemented (no KClass at runtime, " + + "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }", + level = DeprecationLevel.ERROR, + ) override fun dense( + outputDimension: Int, id: String, content: DENSE.() -> Unit - ): Module { - // Create a mixed-precision module that handles conversion - TODO("Mixed-precision dense implementation needed") - } + ): Module = error( + "Mixed-precision dense was redesigned (SKEEP-008): pass the dtype as a value, " + + "dense(outputDimension, dtype = FP16, id) { ... }" + ) + @Deprecated( + message = "The dtype type argument cannot be implemented (no KClass at runtime, " + + "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }", + level = DeprecationLevel.ERROR, + ) + override fun dense( + id: String, + content: DENSE.() -> Unit + ): Module = error( + "Mixed-precision dense was redesigned (SKEEP-008): pass the dtype as a value, " + + "dense(outputDimension, dtype = FP16, id) { ... }" + ) + + @Deprecated( + message = "The dtype type argument cannot be implemented (no KClass at runtime, " + + "SKEEP-008); pass the dtype as a value: stage(id, dtype = FP16) { ... }", + level = DeprecationLevel.ERROR, + ) override fun stage( id: String, content: NeuralNetworkDsl.() -> Unit - ): Module { - // Create a mixed-precision stage that handles conversion - TODO("Mixed-precision stage implementation needed") - } + ): Module = error( + "Mixed-precision stage was redesigned (SKEEP-008): pass the dtype as a value, " + + "stage(id, dtype = FP16) { ... }" + ) override fun batchNorm( numFeatures: Int, @@ -1715,28 +1798,90 @@ public class NeuralNetworkDslImpl( override fun dense( outputDimension: Int, + dtype: TLayer, id: String, - content: DENSE.() -> Unit - ): Module { - // Create a mixed-precision module that handles conversion - TODO("Mixed-precision dense implementation needed") + content: DENSE.() -> Unit, + ) { + val inputDimension = lastDimension + lastDimension = outputDimension + currentShape = intArrayOf(outputDimension) + val layerName = getDefaultName(id, "linear", modules.size) + val impl = DenseImpl( + executionContext, + inputDimension = inputDimension, + _outputDimension = outputDimension, + id = layerName, + kClass = dtype.witnessOf(), + ) + impl.content() + val innerModules = impl.create() + val inner = if (innerModules.size == 1) innerModules[0] else MLP(*innerModules.toTypedArray(), name = layerName) + modules += PrecisionBoundary( + name = layerName, + inner = inner, + innerType = dtype, + outerType = dtypeInstanceOf(kClass), + ) + } + + override fun stage( + id: String, + dtype: TStage, + content: NeuralNetworkDsl.() -> Unit, + ) { + val stageImpl = StageImpl(executionContext, id, dtype.witnessOf()) + stageImpl.lastDimension = lastDimension + stageImpl.currentShape = currentShape?.copyOf() + stageImpl.content() + lastDimension = stageImpl.lastDimension + currentShape = stageImpl.currentShape?.copyOf() + modules += PrecisionBoundary( + name = id, + inner = stageImpl.create(), + innerType = dtype, + outerType = dtypeInstanceOf(kClass), + ) } + @Deprecated( + message = "The dtype type argument cannot be implemented (no KClass at runtime, " + + "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }", + level = DeprecationLevel.ERROR, + ) override fun dense( + outputDimension: Int, id: String, content: DENSE.() -> Unit - ): Module { - // Create a mixed-precision module that handles conversion - TODO("Mixed-precision dense implementation needed") - } + ): Module = error( + "Mixed-precision dense was redesigned (SKEEP-008): pass the dtype as a value, " + + "dense(outputDimension, dtype = FP16, id) { ... }" + ) + @Deprecated( + message = "The dtype type argument cannot be implemented (no KClass at runtime, " + + "SKEEP-008); pass the dtype as a value: dense(outputDimension, dtype = FP16) { ... }", + level = DeprecationLevel.ERROR, + ) + override fun dense( + id: String, + content: DENSE.() -> Unit + ): Module = error( + "Mixed-precision dense was redesigned (SKEEP-008): pass the dtype as a value, " + + "dense(outputDimension, dtype = FP16, id) { ... }" + ) + + @Deprecated( + message = "The dtype type argument cannot be implemented (no KClass at runtime, " + + "SKEEP-008); pass the dtype as a value: stage(id, dtype = FP16) { ... }", + level = DeprecationLevel.ERROR, + ) override fun stage( id: String, content: NeuralNetworkDsl.() -> Unit - ): Module { - // Create a mixed-precision stage that handles conversion - TODO("Mixed-precision stage implementation needed") - } + ): Module = error( + "Mixed-precision stage was redesigned (SKEEP-008): pass the dtype as a value, " + + "stage(id, dtype = FP16) { ... }" + ) override fun batchNorm( numFeatures: Int, @@ -1995,6 +2140,14 @@ public class NeuralNetworkDslImpl( } +/** The [KClass] witness of a dtype instance, typed to the instance's own type (dtypes are objects). */ +@Suppress("UNCHECKED_CAST") +private fun T.witnessOf(): KClass = this::class as KClass + +/** The dtype object whose witness is [kClass]: the DSL impls hold the witness, [PrecisionBoundary] needs the instance. */ +@Suppress("UNCHECKED_CAST") +private fun dtypeInstanceOf(kClass: KClass): T = DType.fromWitness(kClass) as T + @NetworkDsl public class NetworkBuilder { private val modules = mutableListOf>() diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt index 59e597a46..ff31c71a1 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt @@ -521,10 +521,10 @@ public class DenseTensorDataFactory: TensorDataFactory { val data = FloatArray(shape.volume) { random.nextFloat() * range + min } createFloatTensorData(shape, data, FP32 as T) as TensorData } - FP16::class -> { + FP16::class, BF16::class -> { val range = max - min val data = FloatArray(shape.volume) { random.nextFloat() * range + min } - createFloatTensorData(shape, data, FP16 as T) as TensorData + createFloatTensorData(shape, data, narrowTag(dtype) as T) as TensorData } Int32::class -> { val intMin = min.toInt() @@ -562,7 +562,7 @@ public class DenseTensorDataFactory: TensorDataFactory { } createFloatTensorData(shape, data, FP32 as T) as TensorData } - FP16::class -> { + FP16::class, BF16::class -> { val data = FloatArray(shape.volume) { flatIndex -> val indices = IntArray(shape.dimensions.size) var remaining = flatIndex @@ -572,7 +572,7 @@ public class DenseTensorDataFactory: TensorDataFactory { } (generator(indices) as Float) } - createFloatTensorData(shape, data, FP16 as T) as TensorData + createFloatTensorData(shape, data, narrowTag(dtype) as T) as TensorData } Int32::class -> { val data = IntArray(shape.volume) { flatIndex -> @@ -615,9 +615,9 @@ public class DenseTensorDataFactory: TensorDataFactory { val data = FloatArray(shape.volume) { (generator(random) as Float) } createFloatTensorData(shape, data, FP32 as T) as TensorData } - FP16::class -> { + FP16::class, BF16::class -> { val data = FloatArray(shape.volume) { (generator(random) as Float) } - createFloatTensorData(shape, data, FP16 as T) as TensorData + createFloatTensorData(shape, data, narrowTag(dtype) as T) as TensorData } Int32::class -> { val data = IntArray(shape.volume) { (generator(random) as Int) } @@ -702,7 +702,7 @@ public class DenseTensorDataFactory: TensorDataFactory { @Suppress("UNCHECKED_CAST") return when (dtype) { FP32::class -> DenseFloatArrayTensorData(shape, data) as TensorData - FP16::class -> DenseFloatArrayTensorData(shape, data) as TensorData + FP16::class, BF16::class -> DenseFloatArrayTensorData(shape, data) as TensorData else -> throw IllegalArgumentException("wrapFloatArray only supports floating point types: $dtype") } } @@ -714,7 +714,7 @@ public class DenseTensorDataFactory: TensorDataFactory { dtype: KClass, data: FloatArray ): TensorData = when (dtype) { - FP32::class, FP16::class -> wrapFloatArray(shape, dtype, data) + FP32::class, FP16::class, BF16::class -> wrapFloatArray(shape, dtype, data) else -> fromFloatArray(shape, dtype, data) } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps.kt index 7a40dd007..f29a90fc3 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps.kt @@ -3,6 +3,10 @@ package sk.ainet.lang.tensor.ops import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.types.DType +@Deprecated( + message = "Never implemented: TensorOps.convert already covers the cast, and the DSL's " + + "mixed-precision layers use PrecisionBoundary (SKEEP-008, #1433). Will be removed.", +) public interface MixedPrecisionTensorOps : TensorOps { override fun convert( tensor: Tensor,