bf16-native DSL → StableHLO export + pluggable per-target optimization#791
Merged
Conversation
Add the seam that keeps hardware knowledge OUT of the agnostic compiler core:
backends register their own passes from outside core and the pipeline asks the
registry which passes to run for the selected target at each phase.
- CompilePhase { TAPE, DAG, STABLE_HLO }
- TargetOptimizer: per-phase pass provider (dagPasses() today; tape / StableHLO
hooks follow the same shape)
- TargetOptimizers: pluggable registry — register(...) / registerDagPasses(target){}
(callback) / forTarget(target)
- dagPipelineFor(target, corePasses): builds a GraphOptimizationPipeline of the
agnostic core passes plus whatever is plugged in for the target
The transforms a target plugs in still emit standard, portable StableHLO — the
shared IR emitter and model definitions carry no target-specific code.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
mapTensorType/createTensorType unconditionally inserted the `x` shape separator, so an empty (rank-0 scalar) shape produced the malformed `tensor<xbf16>` that iree-compile rejects. Emit `tensor<elem>` when the shape is empty. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
A bf16 variance (a sum of many bf16 squares) loses enough precision that sqrt(var + eps) can overflow to NaN, and some accelerator backends miscompile the decomposed low-precision reduce/normalize outright. Upcast the mean / variance / std / divide to f32 (standard PyTorch/JAX LayerNorm practice); the affine (gamma/beta) stays in the model dtype. No-op for f32 models.
Thread an optional target id and OpGranularityPolicy through toStableHlo -> StableHloConverterFactory -> StableHloConverter -> ConversionContext, so per-target emission decisions (fused vs decomposed ops = target legalization) are possible without any hardware knowledge in the agnostic core. OpGranularityPolicy (+ FusedOpAllowList) lives in skainet-compile-dag, the module both the emitter (skainet-compile-hlo) and the TargetOptimizer registry (skainet-compile-opt) already depend on, so no new dependency edges and no cycle. TargetOptimizer gains granularity(); TargetOptimizers gains granularityFor(target). The caller resolves the policy and passes it in, keeping the emitter decoupled from the registry. Everything defaults to null (decompose all), so emission is byte-identical today.
|
📖 Documentation Preview The documentation has been built successfully for this PR. Generated Files:
Artifacts:
This comment will be updated automatically when the PR is updated. |
michalharakal
changed the title
michalharakal
added a commit
that referenced
this pull request
Jul 5, 2026
…nt 2.07× Bumps VERSION_NAME 0.33.0 -> 0.34.0. Bundles the develop changes since 0.33.0: the new skainet-data-source module (URI-backed sources, HF auth, raw format parsers, suspend data pipeline DSL) + dataset operation views and richer batches (#784/#785), the bf16-native DSL -> StableHLO export path and the pluggable per-phase/per-target compile-optimization seam (#788/#791), NEON K-quant matmul perf (block-outer order + fused Q8 int8 dot, 2.07x Q4_K on Cortex-A55) with aarch64 board verification (#786/#787), LayerNorm f32 normalization + rank-0 tensor-type emission fixes, macOS host build fix (#789), Code of Conduct (#790), and the offline markup-antora docs image (#781). Minor bump (not patch): new published module skainet-data-source; all data-api additions are default-bearing (no source-incompatible changes). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a bf16-native path from the SKaiNET DSL to StableHLO, the plumbing for per-target
optimization/emission, and two numerical/emission correctness fixes. Together this is the
export path used to compile DSL models (e.g. Moonshine) for accelerator backends while
keeping the core hardware-agnostic.
What's in this branch (bottom → top)
1a1cedbc) — trace and emit bf16 models toStableHLO directly.
199a9c9f) —TargetOptimizer/
TargetOptimizersregistry so a backend registers its own DAG passes from outsidecore; the core queries by target string and stays agnostic.
6d036a88) — emittensor<elem>(nottensor<xelem>) forscalars.
006c5ab4) — compute mean/variance/std/divide in f32 for non-f32models, cast back before the affine (
LayerNormalization.forward+convertLayerNorm).Matches PyTorch/JAX; a bf16 variance can overflow
sqrt(var+eps)to NaN.b2a1e167) — thread an optionaltargetandOpGranularityPolicythroughtoStableHlo → StableHloConverter → ConversionContext, andextend
TargetOptimizerwithgranularity(). Defaults to null (decompose all) →emission byte-identical today. This is the seam for "target legalization" (choosing
fused vs decomposed emission per target) without leaking hardware knowledge into core.
Design writeup under
docs/skainet-target-legalization/.Design notes for review
OpGranularityPolicylives inskainet-compile-dag(shared by the emitter and theoptimizer registry) → no new dependency edges, no cycle; the emitter never depends on the
optimizer registry. The caller resolves the policy (
TargetOptimizers.granularityFor) andpasses it into
toStableHlo, so the target string / vendor specifics stay out of core.Testing
skainet-compile-optandskainet-compile-hlocompile;LayerNormConverterTest,ReductionOperationsConverterTestpass.Review notes / risk
converts) and adds a small compute cost. Deliberate correctness choice; if the team prefers
it gated, the op-granularity seam in commit 5 is the mechanism to make it per-target. Happy
to split
006c5ab4into its own PR if you'd rather land it independently.realization) — seam only.