Support pow with Int exponent (#557)

Author: christiangnrd

src/device/intrinsics/math.jl

@@ -230,6 +230,24 @@ end
 @device_override Base.:(^)(x::Float32, y::Float32) = ccall("extern air.pow.f32", llvmcall, Cfloat, (Cfloat, Cfloat), x, y)
 @device_override Base.:(^)(x::Float16, y::Float16) = ccall("extern air.pow.f16", llvmcall, Float16, (Float16, Float16), x, y)
 
+# Avoid use of Float64 in `pow`
+@device_override @inline function Base.:(^)(x::Float32, y::Integer)
+    y == -1 && return inv(x)
+    y == 0 && return one(x)
+    y == 1 && return x
+    y == 2 && return x * x
+    y == 3 && return x * x * x
+    x^Float32(y)
+end
+@device_override @inline function Base.:(^)(x::Float16, y::Integer)
+    y == -1 && return inv(x)
+    y == 0 && return one(x)
+    y == 1 && return x
+    y == 2 && return x * x
+    y == 3 && return x * x * x
+    x^Float16(y)
+end
+
 @device_function powr_fast(x::Float32, y::Float32) = ccall("extern air.fast_powr.f32", llvmcall, Cfloat, (Cfloat, Cfloat), x, y)
 @device_function powr(x::Float32, y::Float32) = ccall("extern air.powr.f32", llvmcall, Cfloat, (Cfloat, Cfloat), x, y)
 @device_function powr(x::Float16, y::Float16) = ccall("extern air.powr.f16", llvmcall, Float16, (Float16, Float16), x, y)

test/device/intrinsics.jl

@@ -300,6 +300,24 @@ end
         @test Array(mtlout) ≈ arr1 .^ arr2
     end
 
+    let #pow with Integer exponent (Issue 552)
+        N = 4
+        arr2 = [-1, 0, 1, 2, 3, rand(-10:10, N)...]
+        arr1 = rand(T, length(arr2))
+        mtlarr1 = MtlArray(arr1)
+        mtlarr2 = MtlArray(arr2)
+
+        mtlout = fill!(similar(mtlarr1), 0)
+
+        function kernel(res, x, y)
+            idx = thread_position_in_grid_1d()
+            res[idx] = x[idx]^y[idx]
+            return nothing
+        end
+        Metal.@sync @metal threads = length(mtlout) kernel(mtlout, mtlarr1, mtlarr2)
+        @test Array(mtlout) ≈ arr1 .^ arr2
+    end
+
     let #powr
         N = 4
         arr1 = rand(T, N)