Use Cassette for contextual dispatch.

Author: maleadt

Manifest.toml

@@ -26,6 +26,11 @@ git-tree-sha1 = "1fce616fa0806c67c133eb1d2f68f0f1a7504665"
 uuid = "c5f51814-7f29-56b8-a69c-e4d8f6be1fde"
 version = "5.0.1"
 
+[[Cassette]]
+git-tree-sha1 = "36bd4e0088652b0b2d25a03e531f0d04258feb78"
+uuid = "7057c7e9-c182-5462-911a-8362d720325c"
+version = "0.3.0"
+
 [[DataStructures]]
 deps = ["InteractiveUtils", "OrderedCollections"]
 git-tree-sha1 = "b7720de347734f4716d1815b00ce5664ed6bbfd4"

Project.toml

@@ -7,6 +7,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
 CUDAapi = "3895d2a7-ec45-59b8-82bb-cfc6a382f9b3"
 CUDAdrv = "c5f51814-7f29-56b8-a69c-e4d8f6be1fde"
+Cassette = "7057c7e9-c182-5462-911a-8362d720325c"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"

src/CUDAnative.jl

@@ -34,6 +34,7 @@ const ptxas = Ref{String}()
 include("utils.jl")
 
 # needs to be loaded _before_ the compiler infrastructure, because of generated functions
+isdevice() = false
 include("device/tools.jl")
 include("device/pointer.jl")
 include("device/array.jl")
@@ -44,6 +45,7 @@ include("device/runtime.jl")
 include("init.jl")
 
 include("compiler.jl")
+include("context.jl")
 include("execution.jl")
 include("exceptions.jl")
 include("reflection.jl")

src/compiler/driver.jl

@@ -62,11 +62,12 @@ function codegen(target::Symbol, job::CompilerJob;
     @timeit_debug to "validation" check_method(job)
 
     @timeit_debug to "Julia front-end" begin
+        f = contextualize(job.f)
 
         # get the method instance
         world = typemax(UInt)
-        meth = which(job.f, job.tt)
-        sig = Base.signature_type(job.f, job.tt)::Type
+        meth = which(f, job.tt)
+        sig = Base.signature_type(f, job.tt)::Type
         (ti, env) = ccall(:jl_type_intersection_with_env, Any,
                           (Any, Any), sig, meth.sig)::Core.SimpleVector
         if VERSION >= v"1.2.0-DEV.320"

src/context.jl

@@ -0,0 +1,75 @@
+##
+# Implements contextual dispatch through Cassette.jl
+# Goals:
+# - Rewrite common CPU functions to appropriate GPU intrinsics
+#
+# TODO:
+# - error (erf, ...)
+# - pow
+# - min, max
+# - mod, rem
+# - gamma
+# - bessel
+# - distributions
+# - unsorted
+
+using Cassette
+
+function transform(ctx, ref)
+    CI = ref.code_info
+    noinline = any(@nospecialize(x) ->
+                       Core.Compiler.isexpr(x, :meta) &&
+                       x.args[1] == :noinline,
+                   CI.code)
+    CI.inlineable = !noinline
+
+    CI.ssavaluetypes = length(CI.code)
+    # Core.Compiler.validate_code(CI)
+    return CI
+end
+
+const InlinePass = Cassette.@pass transform
+
+Cassette.@context CUDACtx
+const cudactx = Cassette.disablehooks(CUDACtx(pass = InlinePass))
+
+###
+# Cassette fixes
+###
+
+# kwfunc fix
+Cassette.overdub(::CUDACtx, ::typeof(Core.kwfunc), f) = return Core.kwfunc(f)
+
+# the functions below are marked `@pure` and by rewritting them we hide that from
+# inference so we leave them alone (see https://github.com/jrevels/Cassette.jl/issues/108).
+@inline Cassette.overdub(::CUDACtx, ::typeof(Base.isimmutable), x)     = return Base.isimmutable(x)
+@inline Cassette.overdub(::CUDACtx, ::typeof(Base.isstructtype), t)    = return Base.isstructtype(t)
+@inline Cassette.overdub(::CUDACtx, ::typeof(Base.isprimitivetype), t) = return Base.isprimitivetype(t)
+@inline Cassette.overdub(::CUDACtx, ::typeof(Base.isbitstype), t)      = return Base.isbitstype(t)
+@inline Cassette.overdub(::CUDACtx, ::typeof(Base.isbits), x)          = return Base.isbits(x)
+
+@inline Cassette.overdub(::CUDACtx, ::typeof(datatype_align), ::Type{T}) where {T} = datatype_align(T)
+
+###
+# Rewrite functions
+###
+Cassette.overdub(ctx::CUDACtx, ::typeof(isdevice)) = true
+
+# libdevice.jl
+for f in (:cos, :cospi, :sin, :sinpi, :tan,
+          :acos, :asin, :atan,
+          :cosh, :sinh, :tanh,
+          :acosh, :asinh, :atanh,
+          :log, :log10, :log1p, :log2,
+          :exp, :exp2, :exp10, :expm1, :ldexp,
+          :isfinite, :isinf, :isnan,
+          :signbit, :abs,
+          :sqrt, :cbrt,
+          :ceil, :floor,)
+    @eval function Cassette.overdub(ctx::CUDACtx, ::typeof(Base.$f), x::Union{Float32, Float64})
+        @Base._inline_meta
+        return CUDAnative.$f(x)
+    end
+end
+
+contextualize(f::F) where F = (args...) -> Cassette.overdub(cudactx, f, args...)

test/device/execution.jl

@@ -463,7 +463,7 @@ end
     val_dev = CuArray(val)
     cuda_ptr = pointer(val_dev)
     ptr = CUDAnative.DevicePtr{Int}(cuda_ptr)
-    for i in (1, 10, 20, 35)
+    for i in (1, 10, 20, 32)
         variables = ('a':'z'..., 'A':'Z'...)
         params = [Symbol(variables[j]) for j in 1:i]
         # generate a kernel
@@ -553,11 +553,11 @@ let (code, out, err) = julia_script(script, `-g2`)
     @test occursin("ERROR: KernelException: exception thrown during kernel execution on device", err)
     @test occursin("ERROR: a exception was thrown during kernel execution", out)
     if VERSION < v"1.3.0-DEV.270"
-        @test occursin("[1] Type at float.jl", out)
+        @test occursin(r"\[.\] Type at float.jl", out)
     else
-        @test occursin("[1] Int64 at float.jl", out)
+        @test occursin(r"\[.\] Int64 at float.jl", out)
     end
-    @test occursin("[2] kernel at none:2", out)
+    @test occursin(r"\[.\] kernel at none:2", out)
 end
 
 end