Skip to content

Add concat / slice / cast StableHLO converters (#489)#491

Merged
michalharakal merged 2 commits intodevelopfrom
feature/489-concat-slice-cast
Apr 13, 2026
Merged

Add concat / slice / cast StableHLO converters (#489)#491
michalharakal merged 2 commits intodevelopfrom
feature/489-concat-slice-cast

Conversation

@michalharakal
Copy link
Copy Markdown
Contributor

Closes #489.

Summary

Adds the last generic structural / type tensor ops missing from mainline SKaiNET's StableHLO emitter: concat (plus `concatenate` / `cat` / `stack` aliases), slice, and cast (plus `convert` / `to` aliases). None of these are LLM- or transformer-specific; they're the standard building blocks that sit alongside the existing reshape / flatten / squeeze / unsqueeze coverage.

Emitted lowerings

```mlir
// concat
%out = stablehlo.concatenate %a, %b, dim = :

// slice
%out = stablehlo.slice %x
{start_indices = [..], limit_indices = [..], strides = [..]}
:

// cast
%out = stablehlo.convert %x : () ->
```

  • `concat` reads `axis` or `dim` (default 0), normalizes negative axes against rank.
  • `slice` reads `start_indices` / `limit_indices` / `strides` (also accepts `starts` / `limits`) and defaults strides to 1 per dim when absent. Static slice only — dynamic / runtime bounds are out of scope.
  • `cast` relies on the output spec's dtype and emits the standard MLIR type-transition signature.

Two commits

1. Failing test

`ConcatSliceCastConverterTest` with 7 cases covering all three ops plus every alias.

2. The fix

Extends `ShapeOperationsConverter` (for concat / slice — structural ops belong with the reshape family) and `MathOperationsConverter` (for cast — elementwise primitive). Also updates `ShapeOperationsConverterTest.testSupportedOperations` and `testRegistryIntegration` to cover the new ops — the prior assertions pinned exact set equality on the four reshape ops and would otherwise regress on this additive change. That regression was caught locally by `allTests` across all targets, not by my earlier narrowly-filtered `jvmTest --tests=ConcatSliceCastConverterTest` run.

Test plan

  • `ConcatSliceCastConverterTest` — 7/7 green
  • `ShapeOperationsConverterTest` — updated, green
  • `./gradlew :skainet-compile:skainet-compile-hlo:allTests -x kotlinWasmStoreYarnLock` — green across jvmTest, wasmJsTest, wasmJsBrowserTest, wasmWasiTest, wasmWasiNodeTest, macosArm64Test, iosSimulatorArm64Test (`linuxX64Test` skipped on macOS host)
  • CI: full multiplatform build

Out of scope

  • Dynamic slice with runtime bounds. Static slice is sufficient for the current emitter targets; can be added when a traced model surfaces the pattern.
  • `gather` / `scatter` — separate converter, gather already landed in Add gather / embedding StableHLO converter (#483) #487.
  • Quantization-aware cast that preserves `TensorEncoding`. The current converter just maps element types; a quant-preserving variant can come in a follow-up.

Roadmap context

This closes out the generic-op coverage needed before the IREE compile-path spike. With this merged, the mainline StableHLO emitter can reasonably handle a full transformer forward pass up to the engine boundary — any LLM-specific ops beyond this point belong in `SKaiNET-transformers`.

🤖 Generated with Claude Code

michalharakal and others added 2 commits April 13, 2026 14:04
Adds ConcatSliceCastConverterTest exercising the three generic
structural / type ops that mainline SKaiNET's StableHLO emitter
doesn't yet cover: concat (with concatenate / cat / stack
aliases), slice (with static start_indices / limit_indices /
strides), and cast (with convert / to aliases).

Asserts each op is claimed by a converter (not dropped as
"Unsupported operation"), lowers to the expected stablehlo.*
op, and carries the expected attributes: `dim = N` for
concatenate, start/limit/strides for slice, and an f16-typed
-> transition for cast(FP32 -> FP16).

Red against current StableHloConverter.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Generic structural / type tensor ops that mainline SKaiNET's
emitter was missing. Extends ShapeOperationsConverter with concat
(plus `concatenate` / `cat` / `stack` aliases) and slice, and
extends MathOperationsConverter with cast (plus `convert` / `to`
aliases). All three are companions to the existing reshape /
flatten / squeeze / unsqueeze coverage — none of them are LLM-
or transformer-specific, they're the standard building blocks
any tensor graph uses.

Emitted lowerings:

  concat: %out = stablehlo.concatenate %a, %b, dim = <axis> : <type>
  slice : %out = stablehlo.slice %x
          {start_indices = [..], limit_indices = [..], strides = [..]}
          : <type>
  cast  : %out = stablehlo.convert %x : (<from>) -> <to>

concat reads `axis` or `dim` (default 0), normalizes negative axes
against rank. slice reads `start_indices` / `limit_indices` /
`strides` (also accepts `starts` / `limits`) and defaults strides
to 1 per dim when absent. cast relies on the output spec's dtype
and emits the standard MLIR type-transition signature.

Also updates ShapeOperationsConverterTest.testSupportedOperations
and testRegistryIntegration to cover the new ops — the old
tests pinned exact set equality on the four reshape ops and would
otherwise regress on the additive change.

Tests: 7/7 in ConcatSliceCastConverterTest green + both updated
shape registration tests green.

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>
@michalharakal michalharakal merged commit 43c0040 into develop Apr 13, 2026
4 checks passed
@michalharakal michalharakal deleted the feature/489-concat-slice-cast branch April 13, 2026 12:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add concat / slice / cast StableHLO converters (P1)

1 participant