clamp & sign

Author: christiangnrd

src/device/intrinsics/math.jl

@@ -8,6 +8,14 @@ using Base.Math: throw_complex_domainerror
 # - add support for vector types
 # - consider emitting LLVM intrinsics and lowering those in the back-end
 
+### Common Intrinsics
+@device_function clamp_fast(x::Float32, minval::Float32, maxval::Float32) = ccall("extern air.fast_clamp.f32", llvmcall, Cfloat, (Cfloat, Cfloat, Cfloat), x, minval, maxval)
+@device_override Base.clamp(x::Float32, minval::Float32, maxval::Float32) = ccall("extern air.clamp.f32", llvmcall, Cfloat, (Cfloat, Cfloat, Cfloat), x, minval, maxval)
+@device_override Base.clamp(x::Float16, minval::Float16, maxval::Float16) = ccall("extern air.clamp.f16", llvmcall, Float16, (Float16, Float16, Float16), x, minval, maxval)
+
+@device_override Base.sign(x::Float32) = ccall("extern air.sign.f32", llvmcall, Cfloat, (Cfloat,), x)
+@device_override Base.sign(x::Float16) = ccall("extern air.sign.f16", llvmcall, Float16, (Float16,), x)
+
 ### Floating Point Intrinsics
 
 ## Metal only supports single and half-precision floating-point types (and their vector counterparts)

test/device/intrinsics.jl

@@ -107,7 +107,7 @@ end
 MATH_INTR_FUNCS_1_ARG = [
     # Common functions
     # saturate, # T saturate(T x) Clamp between 0.0 and 1.0
-    # sign, # T sign(T x) returns 0.0 if x is NaN. Not tested because intrinsic not yet defined
+    sign, # T sign(T x) returns 0.0 if x is NaN
 
     # float math
     acos, # T acos(T x)
@@ -166,7 +166,6 @@ MATH_INTR_FUNCS_2_ARG = [
 
 MATH_INTR_FUNCS_3_ARG = [
     # Common functions
-    # clamp, # T clamp(T x, T minval, T maxval). Not tested because intrinsic not yet defined
     # mix, # T mix(T x, T y, T a) # x+(y-x)*a
     # smoothstep, # T smoothstep(T edge0, T edge1, T x)
     fma, # T fma(T a, T b, T c)
@@ -268,6 +267,27 @@ end
         end
     end
 
+    let # clamp
+        N = 4
+        in = randn(T, N)
+        minval = fill(T(-0.6), N)
+        maxval = fill(T(0.6), N)
+
+        mtlin = MtlArray(in)
+        mtlminval = MtlArray(minval)
+        mtlmaxval = MtlArray(maxval)
+
+        mtlout = fill!(similar(mtlin), 0)
+
+        function kernel(res, x, y, z)
+            idx = thread_position_in_grid_1d()
+            res[idx] = clamp(x[idx], y[idx], z[idx])
+            return nothing
+        end
+        Metal.@sync @metal threads = N kernel(mtlout, mtlin, mtlminval, mtlmaxval)
+        @test Array(mtlout) == clamp.(in, minval, maxval)
+    end
+
     let #pow
         N = 4
         arr1 = rand(T, N)