Merge pull request #130 from jakebolewski/jcb/fix_julia_15

Rewrite expoenent implementation for the CUDA backend with Cassette for Julia 1.5 compatibility

Author: vchuravy

Project.toml

@@ -14,7 +14,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 [compat]
 Adapt = "0.4, 1.0, 2.0"
 CUDA = "~1.0, ~1.1, ~1.2, 1.3"
-Cassette = "0.3.2"
+Cassette = "0.3.3"
 MacroTools = "0.5"
 SpecialFunctions = "0.10"
 StaticArrays = "0.12"

src/backends/cuda.jl

@@ -139,7 +139,7 @@ function (obj::Kernel{CUDADevice})(args...; ndrange=nothing, dependencies=nothin
         ndrange = (ndrange,)
     end
     if workgroupsize isa Integer
-        workgroupsize = (workgroupsize, )
+        workgroupsize = (workgroupsize,)
     end
 
     if KernelAbstractions.workgroupsize(obj) <: DynamicSize && workgroupsize === nothing
@@ -256,6 +256,23 @@ end
 @inline Cassette.overdub(::CUDACtx, ::typeof(SpecialFunctions.erf), x::Union{Float32, Float64}) = CUDA.erf(x)
 @inline Cassette.overdub(::CUDACtx, ::typeof(SpecialFunctions.erfc), x::Union{Float32, Float64}) = CUDA.erfc(x)
 
+@inline function Cassette.overdub(::CUDACtx, ::typeof(exponent), x::Union{Float32, Float64})
+    T = typeof(x)
+    xs = reinterpret(Unsigned, x) & ~Base.sign_mask(T)
+    if xs >= Base.exponent_mask(T)
+        throw(DomainError(x, "Cannot be Nan of Inf."))
+    end
+    k = Int(xs >> Base.significand_bits(T))
+    if k == 0 # x is subnormal
+        if xs == 0
+            throw(DomainError(x, "Cannot be subnormal converted to 0."))
+	end
+        m = Base.leading_zeros(xs) - Base.exponent_bits(T)
+        k = 1 - m
+    end
+    return k - Base.exponent_bias(T)
+end
+
 @static if Base.isbindingresolved(CUDA, :emit_shmem) && Base.isdefined(CUDA, :emit_shmem)
     const emit_shmem = CUDA.emit_shmem
 else

src/compiler/pass.jl

@@ -41,6 +41,7 @@ function transform(ctx, ref)
     end
 
     # overdubbing IntrinsicFunctions removes our ability to profile code
+    # Base.eltype is special because it is used for type inference
     newstmt = (x, i) -> begin
         isassign = Base.Meta.isexpr(x, :(=))
         stmt = isassign ? x.args[2] : x
@@ -61,7 +62,7 @@ function transform(ctx, ref)
                 name = f.name
                 if Base.isbindingresolved(mod, name) && Base.isdefined(mod, name)
                     ff = getfield(f.mod, f.name)
-                    if ff isa Core.IntrinsicFunction || ff isa Core.Builtin
+                    if ff isa Core.IntrinsicFunction || ff isa Core.Builtin || ff === Base.eltype
                         stmt.args[fidx] = Expr(:nooverdub, f)
                     end
                 end
@@ -76,4 +77,4 @@ function transform(ctx, ref)
     return CI
 end
 
-const CompilerPass = Cassette.@pass transform
+const CompilerPass = Cassette.@pass transform