Skip to content

Metal: expand llvm.powi into multiplies - #946

Merged
maleadt merged 1 commit into
JuliaGPU:mainfrom
timesselens:pr/metal-expand-powi
Sep 25, 2026
Merged

maleadt merged 1 commit into
JuliaGPU:mainfrom
timesselens:pr/metal-expand-powi

Conversation

@timesselens

@timesselens timesselens commented Sep 24, 2026 •

Copy link
Copy Markdown
Contributor

With an integer n, @fastmath x^n compiles to llvm.powi (Base.FastMath.pow_fast), which Metal cannot compile:

julia> f(y, x) = (i = thread_position_in_grid().x; @inbounds y[i] = @fastmath(x[i]^3); nothing)
julia> @metal threads=2 f(y, x)
ERROR: Compilation to native code failed; see below for details.
caused by: NSError: Compilation failed due to an interrupted connection: XPC_ERROR_CONNECTION_INTERRUPTED.

AIR has no integer power, and Apple's back-end can't select llvm.powi either, not even with a constant exponent. Apple's own front-end emits the intrinsic for MSL's __builtin_powif, and metal-tt then aborts with unable 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 of powi. So this PR lowers it in GPUCompiler, next to the other intrinsics AIR lacks, and doesn't add a pow_fast override to Metal.jl.

The expansion uses exponentiation by squaring, as LLVM does for CPUs:

  • A constant exponent is unrolled into multiplies, like SelectionDAG's ExpandPowI.
  • Any other exponent calls an alwaysinline loop 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^n there bit for bit, except that the GPU flushes subnormal results to zero. LLVM's own GPU back-ends that lack powi (DXIL, Vulkan SPIR-V, AMDGPU's GlobalISel) lower it to pow(x, sitofp(n)) instead. That isn't an option here, because MSL leaves pow undefined for negative bases.

Tests: FileCheck tests on code_native cover a constant exponent, a run-time exponent and a vector overload with an i64 exponent. They fail on main. 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_fast overrides that go through the float power (__nv_fast_powf), so IR-level powi isn't covered there either.

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
maleadt force-pushed the pr/metal-expand-powi branch from 6b55c19 to 7de6cec Compare September 25, 2026 09:38
@maleadt

maleadt commented Sep 25, 2026

Copy link
Copy Markdown
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

codecov Bot commented Sep 25, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 86.54%. Comparing base (8496909) to head (7de6cec).
⚠️ Report is 1 commits behind head on main.

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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@maleadt
maleadt merged commit 0d06ee4 into JuliaGPU:main Sep 25, 2026
33 checks passed
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.

2 participants