Make it possible to control the 265-hack.

Author: maleadt

src/context.jl

@@ -1,8 +1,5 @@
-##
-# Implements contextual dispatch through Cassette.jl
-# Goals:
-# - Rewrite common CPU functions to appropriate GPU intrinsics
-# 
+# contextual dispatch using Cassette.jl
+#
 # TODO:
 # - error (erf, ...)
 # - pow
@@ -15,40 +12,50 @@
 
 using Cassette
 
-@inline function unknowably_false()
-    Base.llvmcall("ret i8 0", Bool, Tuple{})
-end
+@inline unknowably_false() = Base.llvmcall("ret i8 0", Bool, Tuple{})
+
+function generate_transform(method_redefinitions)
+    return function transform(ctx, ref)
+        CI = ref.code_info
 
-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
+        # inline everything
+        noinline = any(@nospecialize(x) ->
+                           Core.Compiler.isexpr(x, :meta) &&
+                           x.args[1] == :noinline,
+                       CI.code)
+        CI.inlineable = !noinline
 
-    if isinteractive()
-        # 265 fix, insert a call to the original method
-        # that we later will remove with LLVM's DCE
-        # TODO: We also don't want to compile these functions
-        unknowably_false = GlobalRef(@__MODULE__, :unknowably_false)
-        Cassette.insert_statements!(CI.code, CI.codelocs,
-          (x, i) -> i == 1 ?  4 : nothing,
-          (x, i) -> i == 1 ? [
-              Expr(:call, Expr(:nooverdub, unknowably_false)),
-              Expr(:gotoifnot, Core.SSAValue(i), i+3),
-              Expr(:call, Expr(:nooverdub, Core.SlotNumber(1)), (Core.SlotNumber(i) for i in 2:ref.method.nargs)...),
-              x] : nothing)
+        if method_redefinitions
+            # 265 fix, insert a call to the original method
+            # that we later will remove with LLVM's DCE
+            # TODO: We also don't want to compile these functions
+            unknowably_false = GlobalRef(@__MODULE__, :unknowably_false)
+            Cassette.insert_statements!(CI.code, CI.codelocs,
+              (x, i) -> i == 1 ?  4 : nothing,
+              (x, i) -> i == 1 ? [
+                  Expr(:call, Expr(:nooverdub, unknowably_false)),
+                  Expr(:gotoifnot, Core.SSAValue(i), i+3),
+                  Expr(:call, Expr(:nooverdub, Core.SlotNumber(1)), (Core.SlotNumber(i) for i in 2:ref.method.nargs)...),
+                  x] : nothing)
+        end
+        CI.ssavaluetypes = length(CI.code)
+
+        #Core.Compiler.validate_code(CI)
+        return CI
     end
-    CI.ssavaluetypes = length(CI.code)
-    # Core.Compiler.validate_code(CI)
-    return CI
 end
 
-const InlinePass = Cassette.@pass transform
+const StaticPass = Cassette.@pass generate_transform(false)
+const InteractivePass = Cassette.@pass generate_transform(true)
 
 Cassette.@context CUDACtx
-const cudactx = Cassette.disablehooks(CUDACtx(pass = InlinePass))
+const StaticCtx = Cassette.disablehooks(CUDACtx(pass = StaticPass))
+const InteractiveCtx = Cassette.disablehooks(CUDACtx(pass = InteractivePass))
+
+@inline function contextualize(f::F, interactive) where F
+    ctx = interactive ? InteractiveCtx : StaticCtx
+    (args...) -> Cassette.overdub(ctx, f, args...)
+end
 
 ###
 # Cassette fixes
@@ -88,6 +95,3 @@ for f in (:cos, :cospi, :sin, :sinpi, :tan,
         return CUDAnative.$f(x)
     end
 end
-
-contextualize(f::F) where F = (args...) -> Cassette.overdub(cudactx, f, args...)
-

src/execution.jl

@@ -8,7 +8,7 @@ export @cuda, cudaconvert, cufunction, dynamic_cufunction, nearest_warpsize
 # split keyword arguments to `@cuda` into ones affecting the macro itself, the compiler and
 # the code it generates, or the execution
 function split_kwargs(kwargs)
-    macro_kws    = [:dynamic]
+    macro_kws    = [:dynamic, :interactive]
     compiler_kws = [:minthreads, :maxthreads, :blocks_per_sm, :maxregs, :name]
     call_kws     = [:cooperative, :blocks, :threads, :config, :shmem, :stream]
     macro_kwargs = []
@@ -138,11 +138,15 @@ macro cuda(ex...)
 
     # handle keyword arguments that influence the macro's behavior
     dynamic = false
+    interactive = isinteractive()
     for kwarg in macro_kwargs
         key,val = kwarg.args
         if key == :dynamic
             isa(val, Bool) || throw(ArgumentError("`dynamic` keyword argument to @cuda should be a constant value"))
             dynamic = val::Bool
+        elseif key == :interactive
+            isa(val, Bool) || throw(ArgumentError("`interactive` keyword argument to @cuda should be a constant value"))
+            interactive = val::Bool
         else
             throw(ArgumentError("Unsupported keyword argument '$key'"))
         end
@@ -159,7 +163,7 @@ macro cuda(ex...)
             quote
                 # we're in kernel land already, so no need to cudaconvert arguments
                 local kernel_tt = Tuple{$((:(Core.Typeof($var)) for var in var_exprs)...)}
-                local kernel_f = contextualize($(esc(f)))
+                local kernel_f = contextualize($(esc(f)), $interactive)
                 local kernel = dynamic_cufunction(kernel_f, kernel_tt)
                 kernel($(var_exprs...); $(map(esc, call_kwargs)...))
              end)
@@ -173,7 +177,7 @@ macro cuda(ex...)
                 GC.@preserve $(vars...) begin
                     local kernel_args = cudaconvert.(($(var_exprs...),))
                     local kernel_tt = Tuple{Core.Typeof.(kernel_args)...}
-                    local kernel_f = contextualize($(esc(f)))
+                    local kernel_f = contextualize($(esc(f)), $interactive)
                     local kernel = cufunction(kernel_f, kernel_tt;
                                               $(map(esc, compiler_kwargs)...))
                     kernel(kernel_args...; $(map(esc, call_kwargs)...))

test/device/execution.jl

@@ -308,15 +308,15 @@ end
         return
     end
 
-    @cuda kernel(convert(CuPtr{Int}, arr.buf))
+    @cuda interactive=true kernel(convert(CuPtr{Int}, arr.buf))
     @test Array(arr)[] == 1
 
     function kernel(ptr)
         unsafe_store!(ptr, 2)
         return
     end
 
-    @cuda kernel(convert(CuPtr{Int}, arr.buf))
+    @cuda interactive=true kernel(convert(CuPtr{Int}, arr.buf))
     @test Array(arr)[] == 2
 end