feat(ops): argMax(dim) DSL op with StableHLO lowering + CPU kernel#800
Merged
Conversation
…rnel argMax is a real, HW-agnostic DSL op — the index of the max along a dim (ties -> lowest index), returning Int32 indices. Non-differentiable. Like scaledDotProductAttention it stays a single op and is lowered at the StableHLO stage rather than having a dedicated hardware primitive: ArgMaxOperationsConverter composes iota + reduce-max + broadcast + compare + select + reduce-min (the same stablehlo primitives the causal-mask code already emits), so no variadic reducer region is needed. Registered in StableHloConverterFactory (basic/extended/fast). Eager path is a scalar reduction-to-index kernel in DefaultCpuOps. Shape floor in VoidTensorOps; delegation added to RecordingTensorOpsDecorator and the dag TestTensorOps. Enables the FunctionGemma compiled export to fold its logits->token-ids tail into the DSL trace (retiring an external Python MLIR rewrite). Tested: eager numeric (last-dim, ties, middle-dim) + StableHLO lowering (the [1,24,262153]->[1,24] gemma tail); full compile-hlo suite green. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
📖 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. |
`./gradlew apiDump` after adding the public `argMax` op. Updates the jvm ABI dumps of every module that re-exposes it via `TensorOps` (interface method + DSL/graph re-exports): lang-core, lang-dag, backend-cpu, compile-hlo, compile-dag, compile-opt. Also captures pre-existing untracked public API in compile-opt/compile-hlo (FusedOpAllowList / OpGranularityPolicy / a converter ctor) that had drifted out of the golden dumps on develop — apiDump is generated from the actual code, so this only reconciles them. Verified: apiCheck green; js/wasmJs/wasmWasi compile clean; yarn.lock unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
📖 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. |
…/wasm) The eager DefaultCpuOps.argMax returned an Int32 payload inside a <FP32,Float> tensor. JVM tolerated the boxing, but Kotlin/Native (linuxX64) and Wasm threw `ClassCastException: Int cannot be cast to Float` when reading .data — a real consumer bug, not just a test issue (surfaced by `allTests` in CI). Fix: the eager kernel now materializes the indices as index-valued floats in the input dtype, so the result is a consistent Tensor<T,V> readable on every target. The traced/compiled form is unchanged — still a real i32 tensor (VoidTensorOps + ArgMaxOperationsConverter), so the board vmfb path is unaffected. Verified: DefaultCpuOpsArgMaxTest green on linuxX64 + jvm. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
📖 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. |
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.
Adds
argMax(tensor, dim)as a real, HW-agnostic tensor op — the index of the maxalong a dim (ties → lowest index), returning Int32 indices. Non-differentiable.
Design (mirrors
scaledDotProductAttention)A single first-class op lowered at the StableHLO stage rather than needing a bespoke
primitive:
ArgMaxOperationsConvertercomposesiota+reduce-max+broadcast_in_dimcompare EQ+select+reduce-min— the samestablehloprimitives the causal-maskcode already emits, so no variadic reducer region is required. The compose happens at the
StableHLO-primitive level, so no new DSL primitives are added.
Changes
TensorOps.argMax(no@Diff, non-nullabledim) → KSP tracing wrapper routesdimto op params.VoidTensorOpsshape floor (reduced shape, Int32 output).DefaultCpuOpseager kernel (scalar reduction-to-index).ArgMaxOperationsConverter+ registration inStableHloConverterFactory(basic/extended/fast).RecordingTensorOpsDecoratorand the dagTestTensorOps.Why
Lets an LLM export fold its
logits → token-idstail into the DSL trace (retiring anexternal Python MLIR rewrite in FunctionGemma).
argMaxis now reusable for any model.Tests
[1,24,262153]→[1,24]FunctionGemma tail emits iota/reduce-max/reduce-min/i32.compile-hlosuite green.🤖 Generated with Claude Code