From 7de6ceca7f0ce51571660a295d7694e4c49e9bae Mon Sep 17 00:00:00 2001 From: Tim Esselens Date: Tue, 22 Sep 2026 23:05:36 +0200 Subject: [PATCH] Metal: expand llvm.powi into multiplies 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) --- src/metal.jl | 85 +++++++++++++++++++++++++++++++++++++++++++++++++-- test/metal.jl | 42 +++++++++++++++++++++++++ 2 files changed, 125 insertions(+), 2 deletions(-) diff --git a/src/metal.jl b/src/metal.jl index cc0ab724..ce4ae7b0 100644 --- a/src/metal.jl +++ b/src/metal.jl @@ -2243,6 +2243,84 @@ function lower_minimum_maximum!(builder::IRBuilder, call::LLVM.CallBase, minmax: return promote_bf ? fptrunc!(builder, new, typ) : new end +# integer power, which AIR lacks (MSL has no `pown`, and `pow` is undefined for negative +# bases), by exponentiation by squaring as LLVM's back-ends expand it: a constant exponent +# is unrolled into multiplies like SelectionDAG's `ExpandPowI`, any other calls a loop over +# the exponent's bits like compiler-rt's `__powisf2`. A negative exponent takes the +# reciprocal, and `x^0` is 1 (even for NaN). The exponent of a vector `powi` is a scalar. +function lower_powi!(builder::IRBuilder, call::LLVM.CallBase) + x, n = arguments(call) + if n isa LLVM.ConstantInt && width(value_type(n)) <= 64 + return expand_powi!(builder, x, convert(Int, n)) + end + + # the loop halves the exponent, which doesn't work for an `i1` (where 2 wraps to 0) + width(value_type(n)) < 32 && (n = sext!(builder, n, LLVM.Int32Type())) + + mod = LLVM.parent(LLVM.parent(position(builder))) + typ, ntyp = value_type(x), value_type(n) + fn ="air.powi.$(type_suffix(typ)).$(type_suffix(ntyp))" + f = haskey(functions(mod), fn) ? functions(mod)[fn] : build_powi!(mod, fn, typ, ntyp) + return call!(builder, function_type(f), f, LLVM.Value[x, n]) +end + +# 1.0 of a floating-point type, splat across the lanes of a vector type +fp_one(typ::LLVMType) = LLVM.Value(LLVM.API.LLVMConstReal(typ, 1.0)) + +function expand_powi!(builder::IRBuilder, x::LLVM.Value, n::Int) + m = unsigned(abs(n)) # also for `typemin(n)`, which `abs` returns as is + res = nothing # 1.0, until the first set bit + sq = x # x^(2^i) for the current bit i + while m != 0 + isodd(m) && (res = res === nothing ? sq : fmul!(builder, res, sq)) + m >>= 1 + m != 0 && (sq = fmul!(builder, sq, sq)) + end + res = something(res, fp_one(value_type(x))) + return n < 0 ? fdiv!(builder, fp_one(value_type(x)), res) : res +end + +function build_powi!(mod::LLVM.Module, fn::String, typ::LLVMType, ntyp::LLVMType) + f = LLVM.Function(mod, fn, LLVM.FunctionType(typ, LLVMType[typ, ntyp])) + linkage!(f, LLVM.API.LLVMInternalLinkage) + push!(function_attributes(f), EnumAttribute("alwaysinline")) + x, n = parameters(f) + one = fp_one(typ) + zero = LLVM.ConstantInt(ntyp, 0) + + bb_entry = BasicBlock(f, "entry") + bb_loop = BasicBlock(f, "loop") + bb_done = BasicBlock(f, "done") + @dispose builder=IRBuilder() begin + position!(builder, bb_entry) + br!(builder, icmp!(builder, LLVM.API.LLVMIntEQ, n, zero), bb_done, bb_loop) + + # multiply the squares selected by the bits of `n`, least significant first. like + # `__powisf2`, shift the signed `n` by halving it (rounding towards zero), so as not + # to take `abs(n)`, which InstCombine would turn into an `llvm.abs` after that has + # already been lowered. + position!(builder, bb_loop) + acc = phi!(builder, typ, "acc") + sq = phi!(builder, typ, "sq") + rest = phi!(builder, ntyp, "rest") + bit = trunc!(builder, rest, LLVM.Int1Type()) + acc′ = select!(builder, bit, fmul!(builder, acc, sq), acc) + sq′ = fmul!(builder, sq, sq) + rest′ = sdiv!(builder, rest, LLVM.ConstantInt(ntyp, 2)) + br!(builder, icmp!(builder, LLVM.API.LLVMIntEQ, rest′, zero), bb_done, bb_loop) + append!(incoming(acc), [(one, bb_entry), (acc′, bb_loop)]) + append!(incoming(sq), [(x, bb_entry), (sq′, bb_loop)]) + append!(incoming(rest), [(n, bb_entry), (rest′, bb_loop)]) + + position!(builder, bb_done) + pow = phi!(builder, typ, "pow") + append!(incoming(pow), [(one, bb_entry), (acc′, bb_loop)]) + negative = icmp!(builder, LLVM.API.LLVMIntSLT, n, zero) + ret!(builder, select!(builder, negative, fdiv!(builder, one, pow), pow)) + end + return f +end + function build_minimum_maximum!(mod::LLVM.Module, fn::String, op_ft::LLVM.FunctionType, jltyp::Type, minmax::String) optyp = return_type(op_ft) @@ -2328,8 +2406,9 @@ function lower_llvm_intrinsics!(@nospecialize(job::CompilerJob), fun::LLVM.Funct removable = Set(LLVM.Intrinsic.(REMOVABLE_INTRINSICS)) value_intrinsics = intrinsic_table(AIR_VALUE_INTRINSICS) bit_intrinsics = intrinsic_table(AIR_BIT_INTRINSICS) - is_fpclass, copysign, minimum, maximum = - LLVM.Intrinsic.(("llvm.is.fpclass", "llvm.copysign", "llvm.minimum", "llvm.maximum")) + is_fpclass, copysign, minimum, maximum, powi = + LLVM.Intrinsic.(("llvm.is.fpclass", "llvm.copysign", "llvm.minimum", "llvm.maximum", + "llvm.powi")) changed |= lower_intrinsic_calls!(fun) do builder, call, intr if intr in removable :erase @@ -2346,6 +2425,8 @@ function lower_llvm_intrinsics!(@nospecialize(job::CompilerJob), fun::LLVM.Funct lower_copysign!(builder, call) elseif intr == minimum || intr == maximum lower_minimum_maximum!(builder, call, intr == minimum ? "min" : "max") + elseif intr == powi + lower_powi!(builder, call) end end diff --git a/test/metal.jl b/test/metal.jl index b31a26c8..9b69b434 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -1246,6 +1246,48 @@ end end end +@testset "integer power lowering" begin + # AIR has no integer power, so `llvm.powi` (as emitted for `@fastmath x^n`) is expanded + # into multiplies: unrolled for a constant exponent, and a loop over its bits otherwise. + mod = @eval module $(gensym()) + pow_const(x) = @fastmath x^-5 + pow_var(x, n) = @fastmath x^n + pow_vec(x, n) = ccall("llvm.powi.v2f32.i64", llvmcall, NTuple{2, VecElement{Float32}}, + (NTuple{2, VecElement{Float32}}, Int64), x, n) + end + + # x^-5 = 1 / (x * (x^2)^2) + @test @filecheck begin + @check_label "define float @{{(julia|j)_pow_const_[0-9]+}}" + @check_not "@llvm.powi" + @check_count 3 "fmul float" + @check "fdiv float 1.000000e+00" + @check_not "fmul" + @check "ret float" + Metal.code_native(mod.pow_const, Tuple{Float32}) + end + + @test @filecheck begin + @check_label "define float @{{(julia|j)_pow_var_[0-9]+}}" + @check_not "@llvm.powi" + @check "phi float" + @check "fmul float" + @check "fdiv float 1.000000e+00" + @check_not "@{{(llvm|air)\\.[a-z_]*pow}}" + Metal.code_native(mod.pow_var, Tuple{Float32, Int32}) + end + + @test @filecheck begin + @check_label "define <2 x float> @{{(julia|j)_pow_vec_[0-9]+}}" + @check_not "@llvm.powi" + @check "phi <2 x float>" + @check "fmul <2 x float>" + @check "fdiv <2 x float>" + @check_not "@{{(llvm|air)\\.[a-z_]*pow}}" + Metal.code_native(mod.pow_vec, Tuple{NTuple{2, VecElement{Float32}}, Int64}) + end +end + @testset "integer intrinsic lowering" begin # The integer ops Julia emits as llvm.* are lowered to their AIR builtins, so Metal.jl need # not wrap them. Names/signatures verified against Apple's frontend: