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:
|
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.