Repository navigation
Metal: expand llvm.powi into multiplies - #946
Merged
Merged
Conversation
AIR has no integer power, and Apple's back-end cannot select `llvm.powi` (`metal-tt`: "unable to legalize instruction ... G_FPOWI", even for a constant exponent), so `@fastmath x^n` with an integer `n` failed to compile. The intrinsic also reaches the back-end from IR-level sources a Julia override cannot intercept, e.g. Enzyme's derivative of `powi`. Expand it by exponentiation by squaring, like LLVM's CPU back-ends do: a constant exponent is unrolled into multiplies (SelectionDAG's `ExpandPowI`), any other calls an `alwaysinline` loop over the exponent's bits (compiler-rt's `__powisf2`). This multiplies in the same order, so the results match `@fastmath x^n` on the CPU. `air.pow` is not an option: it is undefined for negative bases. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
maleadt
force-pushed
the
pr/metal-expand-powi
branch
from
September 25, 2026 09:38
6b55c19 to
7de6cec
Compare
Member
|
LGTM, thanks! Slightly abbreviated and rebased. I'll open a counterpart PR on Metal.jl with execution tests once this has been merged and tagged. |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #946 +/- ##
==========================================
+ Coverage 86.41% 86.54% +0.13%
==========================================
Files 29 29
Lines 5674 5729 +55
==========================================
+ Hits 4903 4958 +55
Misses 771 771 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
maleadt
added a commit
that referenced
this pull request
Oct 7, 2026
* Use identity comparison during type mangling (#910) (cherry picked from commit 1c5bef9) * Mangle a `Vararg` whose length is still a `TypeVar` (#921) (cherry picked from commit 370be42) * Throw a KernelError for non-singleton type arguments (#954) (cherry picked from commit 87076e0) * Look up the global logger's level in the latest world (#977) (cherry picked from commit c2f106b) * PTX: Restore FMA contraction. (#881) (cherry picked from commit 96fb6f6) * [SPIRV] add workaround for khronos translator (#879) Co-authored-by: Tim Besard <tim.besard@gmail.com> (cherry picked from commit 081c892) * Handle allocation failure when boxing (#919) (cherry picked from commit ea53814) * Fix invalid IR from small type tags on Julia 1.10 and 1.11 (#972) Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit b0dd255) * SPIR-V: lower llvm.minimum and llvm.maximum (#937) (cherry picked from commit 8eecfc8) * Metal: lower `llvm.sin` and `llvm.cos` to AIR (#967) Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> (cherry picked from commit cc26075) * Metal: lower bare LLVM fences to air.atomic.fence (#932) (cherry picked from commit cae55dd) * Metal: expand `llvm.powi` into multiplies (#946) [1.x: re-expressed in 1.23's worklist style, without #941's helpers.] Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> (cherry picked from commit 0d06ee4) * Metal: expand vector reductions into scalar chains (#962) [1.x: re-expressed in 1.23's worklist style, without #941's helpers.] Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> (cherry picked from commit 9f0d0b8) --------- Co-authored-by: Joachim Brand <j.brand@massey.ac.nz> Co-authored-by: Valentin Churavy <v.churavy@gmail.com> Co-authored-by: Simeon David Schaub <simeon@schaub.rocks> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Tim Esselens <tim.esselens@gmail.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.
With an integer
n,@fastmath x^ncompiles tollvm.powi(Base.FastMath.pow_fast), which Metal cannot compile:AIR has no integer power, and Apple's back-end can't select
llvm.powieither, not even with a constant exponent. Apple's own front-end emits the intrinsic for MSL's__builtin_powif, andmetal-ttthen aborts withunable to legalize instruction: ... G_FPOWI. The intrinsic also reaches the back-end from IR that a Julia override can't intercept, such as Enzyme's derivative ofpowi. So this PR lowers it in GPUCompiler, next to the other intrinsics AIR lacks, and doesn't add apow_fastoverride to Metal.jl.The expansion uses exponentiation by squaring, as LLVM does for CPUs:
ExpandPowI.alwaysinlineloop over the exponent's bits, like compiler-rt's__powisf2.lower_air!then inlines it.A negative exponent takes the reciprocal, and
x^0 = 1, also for NaN. Because the multiplies happen in the same order as on the CPU, the GPU results match@fastmath x^nthere bit for bit, except that the GPU flushes subnormal results to zero. LLVM's own GPU back-ends that lackpowi(DXIL, Vulkan SPIR-V, AMDGPU's GlobalISel) lower it topow(x, sitofp(n))instead. That isn't an option here, because MSL leavespowundefined for negative bases.Tests: FileCheck tests on
code_nativecover a constant exponent, a run-time exponent and a vector overload with ani64exponent. They fail onmain. A Metal.jl integration test that compares against the CPU will follow once this is released.Related: JuliaGPU/CUDA.jl#3065 is the same failure on CUDA. JuliaGPU/CUDA.jl#3098 fixed it with
pow_fastoverrides that go through the float power (__nv_fast_powf), so IR-levelpowiisn't covered there either.