Add gather / embedding StableHLO converter (#483)#487
Merged
michalharakal merged 2 commits intodevelopfrom Apr 13, 2026
Merged
Conversation
Adds GatherConverterTest with four cases: 1. gather_and_embedding_aliases_are_supported — `gather` must be claimed by a converter, not fall through to the registry's "No converter found" path. Red against develop today. 2. embedding_alias_routes_to_same_lowering — same for the `embedding` alias. 3. index_select_alias_routes_to_same_lowering — same for the `index_select` alias (PyTorch-style name). 4. embedding_lowering_carries_canonical_dim_numbers_and_slice_sizes — the emitted `stablehlo.gather` must carry the dim_numbers (offset_dims, collapsed_slice_dims, start_index_map) and slice_sizes attributes that downstream MLIR tools (IREE) expect for a 1-D index tensor indexing the leading dim of a 2-D embedding weight. Test fixture builds the canonical embedding-lookup shape: vocab_size=8, hidden_size=4, seq_len=3, axis=0. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Introduces GatherOperationsConverter claiming the operation
names `gather`, `embedding`, `Embedding`, and `index_select`.
Every LLM export begins with a token-id -> embedding lookup, so
without this converter a traced Llama / Mistral / Qwen / Gemma
forward pass was failing at its very first operation and never
reaching the norms, activations, softmax, or attention the other
P1 converters cover.
Target lowering is the canonical `embedding(input_ids)` shape:
1-D index tensor indexing the leading dim of a 2-D embedding
weight. For vocab_size=8, hidden_size=4, seq_len=3 and axis=0
the emitted op is:
%out = stablehlo.gather(%W, %ids)
{ dimension_numbers = #stablehlo.gather<
offset_dims = [1],
collapsed_slice_dims = [0],
start_index_map = [0],
index_vector_dim = 1>,
slice_sizes = array<i64: 1, 4>,
indices_are_sorted = false }
: (tensor<8x4xf32>, tensor<3xi32>) -> tensor<3x4xf32>
offset_dims / collapsed_slice_dims / start_index_map are derived
from a single gather axis. slice_sizes gets a 1 along the
gathered axis and the full extent of every other dim. Negative
axes are normalized against weight rank. Higher-rank gathers
(attention-side index gathers, multi-dim scatter/gather) can
be added in follow-ups when a traced model surfaces them; the
first PR deliberately targets the LLM front-door case only.
Registered in StableHloConverterFactory.createBasic /
createExtended / createFast.
Tests: 4/4 in GatherConverterTest — registration-via-missing
for each of the three aliases plus a canonical dim_numbers /
slice_sizes assertion. The canonical test also pins the exact
`stablehlo.gather(%arg0, %arg1)` operand shape so a prior
Kotlin-string-template bug that emitted
`stablehlo.gather([%arg0, %arg1][0], [%arg0, %arg1][1])` can
never regress.
Verified locally with
`./gradlew :skainet-compile:skainet-compile-hlo:allTests
-x kotlinWasmStoreYarnLock` — green across jvmTest, wasmJsTest,
wasmJsBrowserTest, wasmWasiTest, wasmWasiNodeTest, macosArm64Test,
and iosSimulatorArm64Test.
Co-Authored-By: Claude Opus 4.6 (1M context) <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
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.
Closes #483.
Summary
Adds a real converter for `gather` / `embedding` / `index_select` / `Embedding`. Every LLM export begins with a token-id \u2192 embedding lookup, and without this converter a traced Llama / Mistral / Qwen / Gemma forward pass was failing at its very first operation — it never reached the norms, softmax, RMSNorm, LayerNorm, or attention that the other P1 converters cover.
Target lowering
For the canonical `embedding(input_ids)` shape (vocab_size=8, hidden_size=4, seq_len=3, axis=0):
```mlir
%out = stablehlo.gather(%W, %ids)
{ dimension_numbers = #stablehlo.gather<
offset_dims = [1],
collapsed_slice_dims = [0],
start_index_map = [0],
index_vector_dim = 1>,
slice_sizes = array<i64: 1, 4>,
indices_are_sorted = false }
: (tensor<8x4xf32>, tensor<3xi32>) -> tensor<3x4xf32>
```
The `offset_dims` / `collapsed_slice_dims` / `start_index_map` are derived from the gather axis. `slice_sizes` gets a `1` along the gathered axis and the full extent along every other dimension. Negative axes are normalized against the weight rank.
Registered in `StableHloConverterFactory.createBasic` / `createExtended` / `createFast`.
Two commits
1. Failing test
`GatherConverterTest` with 4 cases:
```kotlin
assertTrue(mlir.contains("stablehlo.gather(%arg0, %arg1)"))
assertFalse(mlir.contains("stablehlo.gather([%"))
```
This catches a Kotlin string-template pitfall I hit in development: `"$operands[0]"` expands to `[a, b][0]` literal, producing `stablehlo.gather([%arg0, %arg1][0], [%arg0, %arg1][1])`. The assertion pins the correct shape so the bug can never regress silently.
2. The converter
`GatherOperationsConverter.kt` implementing the lowering above, plus factory registration.
Test plan
Out of scope
🤖 Generated with Claude Code