Fixes for and tests using JET. (#1577)

maleadt · web-flow · commit dbc767b1b75f · 2022-08-17T13:28:17.000+02:00
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -83,7 +83,8 @@ steps:
     plugins:
       - JuliaCI/julia#v1:
           version: 1.6
-      - JuliaCI/julia-test#v1: ~
+      - JuliaCI/julia-test#v1:
+          test_args: "--thorough"
       - JuliaCI/julia-coverage#v1:
           codecov: true
           dirs:
diff --git a/Project.toml b/Project.toml
@@ -24,7 +24,6 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
-TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [compat]
 AbstractFFTs = "0.4, 0.5, 1.0"
@@ -40,5 +39,4 @@ RandomNumbers = "1.5.3"
 Reexport = "0.2, 1.0"
 Requires = "0.5, 1.0"
 SpecialFunctions = "1.3, 2"
-TimerOutputs = "0.5.9"
 julia = "1.6"
diff --git a/deps/Deps.jl b/deps/Deps.jl
@@ -1,6 +1,6 @@
 module Deps
 
-Base.Experimental.@compiler_options compile=min optimize=0 infer=false
+Base.Experimental.@compiler_options compile=min optimize=0
 
 import ..CUDA
 import ..LLVM
diff --git a/lib/cudadrv/CUDAdrv.jl b/lib/cudadrv/CUDAdrv.jl
@@ -43,7 +43,8 @@ include("libcuda_deprecated.jl")
 
     # query the system driver version
     function get_version(driver)
-        library_handle = Libdl.dlopen(driver)
+        library_handle = Libdl.dlopen(driver; throw_error=true)
+        library_handle = library_handle::Ptr{Nothing} # doesn't const-propthrow_error
         try
             function_handle = Libdl.dlsym(library_handle, "cuDriverGetVersion")
             version_ref = Ref{Cint}()
@@ -82,7 +83,8 @@ include("libcuda_deprecated.jl")
 
     # if we're using an older driver; consider using forward compatibility
     function do_init(driver)
-        library_handle = Libdl.dlopen(driver)
+        library_handle = Libdl.dlopen(driver; throw_error=true)
+        library_handle = library_handle::Ptr{Nothing} # doesn't const-propthrow_error
         try
             function_handle = Libdl.dlsym(library_handle, "cuInit")
             @check ccall(function_handle, CUresult, (UInt32,), 0)
diff --git a/lib/cudadrv/state.jl b/lib/cudadrv/state.jl
@@ -50,7 +50,8 @@ mutable struct TaskLocalState
         math_mode = something(default_math_mode[],
                               Base.JLOptions().fast_math==1 ? FAST_MATH : DEFAULT_MATH)
         math_precision = something(default_math_precision[], :TensorFloat32)
-        new(dev, ctx, Base.fill(nothing, ndevices()), math_mode, math_precision)
+        new(dev, ctx, Union{Nothing,CuStream}[nothing for _ in 1:ndevices()],
+            math_mode, math_precision)
     end
 end
 
@@ -159,7 +160,7 @@ end
 @inline function context!(f::Function, ctx::CuContext; skip_destroyed::Bool=false)
     # @inline so that the kwarg method is inlined too and we can const-prop skip_destroyed
     if isvalid(ctx)
-        old_ctx = context!(ctx)
+        old_ctx = context!(ctx)::Union{CuContext,Nothing}
         try
             f()
         finally
@@ -187,7 +188,7 @@ end
 
 const __device_contexts = LazyInitialized{Vector{Union{Nothing,CuContext}}}()
 device_contexts() = get!(__device_contexts) do
-    [nothing for _ in 1:ndevices()]
+    Union{Nothing,CuContext}[nothing for _ in 1:ndevices()]
 end
 function device_context(i::Int)
     contexts = device_contexts()
@@ -419,8 +420,8 @@ function PerDevice{T}() where {T}
     PerDevice{T}(ReentrantLock(), values)
 end
 
-get_values(x::PerDevice) = get!(x.values) do
-    Base.fill(nothing, ndevices())
+get_values(x::PerDevice{T}) where {T} = get!(x.values) do
+    Union{Nothing,Tuple{CuContext,T}}[nothing for _ in 1:ndevices()]
 end
 
 function Base.get(x::PerDevice, dev::CuDevice, val)
@@ -437,20 +438,20 @@ function Base.get(x::PerDevice, dev::CuDevice, val)
     end
 end
 
-function Base.get!(constructor::F, x::PerDevice, dev::CuDevice) where {F}
+function Base.get!(constructor::F, x::PerDevice{T}, dev::CuDevice) where {F, T}
     y = get_values(x)
     id = deviceid(dev)+1
     ctx = device_context(id)    # may be nothing
     @inbounds begin
         # test-lock-test
-        if y[id] === nothing || y[id][1] !== ctx
+        if y[id] === nothing || (y[id]::Tuple)[1] !== ctx
             Base.@lock x.lock begin
-                if y[id] === nothing || y[id][1] !== ctx
+                if y[id] === nothing || (y[id]::Tuple)[1] !== ctx
                     y[id] = (context(), constructor())
                 end
             end
         end
-        y[id][2]
+        (y[id]::Tuple)[2]
     end
 end
 
diff --git a/lib/cudadrv/stream.jl b/lib/cudadrv/stream.jl
@@ -11,7 +11,7 @@ Create a CUDA stream.
 """
 mutable struct CuStream
     handle::CUstream
-    ctx::Union{CuContext,Nothing}
+    ctx::CuContext
 
     function CuStream(; flags::CUstream_flags=STREAM_DEFAULT,
                         priority::Union{Nothing,Integer}=nothing)
@@ -29,11 +29,11 @@ mutable struct CuStream
         return obj
     end
 
-    global default_stream() = new(convert(CUstream, C_NULL), nothing)
+    global default_stream() = new(convert(CUstream, C_NULL))
 
-    global legacy_stream() = new(convert(CUstream, 1), nothing)
+    global legacy_stream() = new(convert(CUstream, 1))
 
-    global per_thread_stream() = new(convert(CUstream, 2), nothing)
+    global per_thread_stream() = new(convert(CUstream, 2))
 end
 
 """
diff --git a/lib/utils/memoization.jl b/lib/utils/memoization.jl
@@ -47,7 +47,7 @@ macro memoize(ex...)
         global_cache_eltyp = :(Union{Nothing,$rettyp})
         ex = quote
             cache = get!($(esc(global_cache))) do
-                [nothing for _ in 1:Threads.nthreads()]
+                $global_cache_eltyp[nothing for _ in 1:Threads.nthreads()]
             end
             cached_value = @inbounds cache[Threads.threadid()]
             if cached_value !== nothing
@@ -64,7 +64,7 @@ macro memoize(ex...)
         global_init = :(Union{Nothing,$rettyp}[nothing for _ in 1:$(esc(options[:maxlen]))])
         ex = quote
             cache = get!($(esc(global_cache))) do
-                [$global_init for _ in 1:Threads.nthreads()]
+                $global_cache_eltyp[$global_init for _ in 1:Threads.nthreads()]
             end
             local_cache = @inbounds begin
                 tid = Threads.threadid()
@@ -86,7 +86,7 @@ macro memoize(ex...)
         global_init = :(Dict{$(key.typ),$rettyp}())
         ex = quote
             cache = get!($(esc(global_cache))) do
-                [$global_init for _ in 1:Threads.nthreads()]
+                $global_cache_eltyp[$global_init for _ in 1:Threads.nthreads()]
             end
             local_cache = @inbounds begin
                 tid = Threads.threadid()
diff --git a/lib/utils/threading.jl b/lib/utils/threading.jl
@@ -29,11 +29,11 @@ end
     x.value[]
 end
 
-@noinline function initialize!(x::LazyInitialized, constructor::F1, hook::F2) where {F1, F2}
+@noinline function initialize!(x::LazyInitialized{T}, constructor::F1, hook::F2) where {T, F1, F2}
     status = Threads.atomic_cas!(x.guard, 0, 1)
     if status == 0
         try
-          x.value[] = constructor()
+          x.value[] = constructor()::T
           x.guard[] = 2
         catch
           x.guard[] = 0
diff --git a/src/CUDA.jl b/src/CUDA.jl
@@ -24,21 +24,6 @@ using ExprTools: splitdef, combinedef
 include("../deps/Deps.jl")
 using .Deps
 
-# only use TimerOutputs on non latency-critical CI, in part because
-# @timeit_debug isn't truely zero-cost (KristofferC/TimerOutputs.jl#120)
-if getenv("CI", false) && !getenv("BENCHMARKS", false)
-    using TimerOutputs
-    const to = TimerOutput()
-
-    macro timeit_ci(args...)
-        TimerOutputs.timer_expr(CUDA, false, :($CUDA.to), args...)
-    end
-else
-    macro timeit_ci(args...)
-        esc(args[end])
-    end
-end
-
 
 ## source code includes
 
diff --git a/src/array.jl b/src/array.jl
@@ -69,14 +69,14 @@ function unsafe_free!(xs::CuArray, stream::CuStream=stream())
   # this call should only have an effect once, because both the user and the GC can call it
   if xs.storage === nothing
     return
-  elseif xs.storage.refcount[] < 0
+  elseif (xs.storage::ArrayStorage).refcount[] < 0
     throw(ArgumentError("Cannot free an unmanaged buffer."))
   end
 
-  refcount = Threads.atomic_add!(xs.storage.refcount, -1)
+  refcount = Threads.atomic_add!((xs.storage::ArrayStorage).refcount, -1)
   if refcount == 1
     context!(context(xs); skip_destroyed=true) do
-      free(xs.storage.buffer; stream)
+      free((xs.storage::ArrayStorage).buffer; stream)
     end
   end
 
@@ -236,7 +236,7 @@ Base.sizeof(x::CuArray) = Base.elsize(x) * length(x)
 
 function context(A::CuArray)
   A.storage === nothing && throw(UndefRefError())
-  return A.storage.buffer.ctx
+  return (A.storage::ArrayStorage).buffer.ctx
 end
 
 function device(A::CuArray)
@@ -318,8 +318,10 @@ Base.convert(::Type{T}, x::T) where T <: CuArray = x
 
 Base.unsafe_convert(::Type{Ptr{T}}, x::CuArray{T}) where {T} =
   throw(ArgumentError("cannot take the CPU address of a $(typeof(x))"))
-Base.unsafe_convert(::Type{CuPtr{T}}, x::CuArray{T}) where {T} =
-  convert(CuPtr{T}, x.storage.buffer) + x.offset*Base.elsize(x)
+function Base.unsafe_convert(::Type{CuPtr{T}}, x::CuArray{T}) where {T}
+  x.storage === nothing && throw(UndefRefError())
+  convert(CuPtr{T}, (x.storage::ArrayStorage).buffer) + x.offset*Base.elsize(x)
+end
 
 
 ## interop with device arrays
diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl
@@ -51,11 +51,11 @@ macro cuda(ex...)
     dynamic = false
     launch = true
     for kwarg in macro_kwargs
-        key,val = kwarg.args
-        if key == :dynamic
+        key::Symbol, 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 == :launch
+        elseif key === :launch
             isa(val, Bool) || throw(ArgumentError("`launch` keyword argument to @cuda should be a constant value"))
             launch = val::Bool
         else
@@ -289,7 +289,7 @@ The output of this function is automatically cached, i.e. you can simply call `c
 in a hot path without degrading performance. New code will be generated automatically, when
 when function changes, or when different types or keyword arguments are provided.
 """
-@timeit_ci function cufunction(f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where {F,TT}
+function cufunction(f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where {F,TT}
     cuda = active_state()
     cache = cufunction_cache(cuda.context)
     source = FunctionSpec(f, tt, true, name)
@@ -347,11 +347,11 @@ function cufunction_compile(@nospecialize(job::CompilerJob))
         cufunction_compile(job, ctx)
     end
 end
-@timeit_ci "compile" function cufunction_compile(@nospecialize(job::CompilerJob), ctx)
+function cufunction_compile(@nospecialize(job::CompilerJob), ctx)
     # lower to PTX
-    mi, mi_meta = @timeit_ci "emit_julia" GPUCompiler.emit_julia(job)
-    ir, ir_meta = @timeit_ci "emit_llvm" GPUCompiler.emit_llvm(job, mi; ctx)
-    asm, asm_meta = @timeit_ci "emit_asm" GPUCompiler.emit_asm(job, ir; format=LLVM.API.LLVMAssemblyFile)
+    mi, mi_meta = GPUCompiler.emit_julia(job)
+    ir, ir_meta = GPUCompiler.emit_llvm(job, mi; ctx)
+    asm, asm_meta = GPUCompiler.emit_asm(job, ir; format=LLVM.API.LLVMAssemblyFile)
 
     # remove extraneous debug info on lower debug levels
     if Base.JLOptions().debug_level < 2
@@ -414,7 +414,7 @@ end
         "--output-file", ptxas_output,
         ptx_input
     ])
-    proc, log = @timeit_ci "ptxas" run_and_collect(`$(ptxas()) $ptxas_opts`)
+    proc, log = run_and_collect(`$(ptxas()) $ptxas_opts`)
     log = strip(log)
     if !success(proc)
         reason = proc.termsignal > 0 ? "ptxas received signal $(proc.termsignal)" :
@@ -446,7 +446,7 @@ end
             "--output-file", nvlink_output,
             ptxas_output
         ])
-        proc, log = @timeit_ci "nvlink" run_and_collect(`$(nvlink()) $nvlink_opts`)
+        proc, log = run_and_collect(`$(nvlink()) $nvlink_opts`)
         log = strip(log)
         if !success(proc)
             reason = proc.termsignal > 0 ? "nvlink received signal $(proc.termsignal)" :
@@ -473,10 +473,10 @@ end
 end
 
 # link into an executable kernel
-@timeit_ci "link" function cufunction_link(@nospecialize(job::CompilerJob), compiled)
+function cufunction_link(@nospecialize(job::CompilerJob), compiled)
     # load as an executable kernel object
     ctx = context()
-    mod = @timeit_ci "CuModule" CuModule(compiled.image)
+    mod = CuModule(compiled.image)
     CuFunction(mod, compiled.entry)
 end
 
diff --git a/src/pool.jl b/src/pool.jl
@@ -32,17 +32,11 @@ AllocStats(b::AllocStats, a::AllocStats) =
 
 ## CUDA allocator
 
-@timeit_ci function actual_alloc(bytes::Integer; async::Bool=false,
-                                 stream::Union{CuStream,Nothing}=nothing)
+function actual_alloc(bytes::Integer; async::Bool=false,
+                      stream::Union{CuStream,Nothing}=nothing)
   # try the actual allocation
   buf = try
-    time = Base.@elapsed begin
-      buf = @timeit_ci "Mem.alloc" begin
-        Mem.alloc(Mem.Device, bytes; async, stream)
-      end
-    end
-
-    buf
+    Mem.alloc(Mem.Device, bytes; async, stream)
   catch err
     isa(err, OutOfGPUMemoryError) || rethrow()
     return nothing
@@ -51,12 +45,10 @@ AllocStats(b::AllocStats, a::AllocStats) =
   return buf
 end
 
-@timeit_ci function actual_free(buf::Mem.DeviceBuffer;
-                                stream::Union{CuStream,Nothing}=nothing)
+function actual_free(buf::Mem.DeviceBuffer;
+                     stream::Union{CuStream,Nothing}=nothing)
   # free the memory
-  time = Base.@elapsed begin
-    @timeit_ci "Mem.free" Mem.free(buf; stream)
-  end
+  Mem.free(buf; stream)
 
   return
 end
@@ -82,8 +74,6 @@ end
 function pool_mark(dev::CuDevice)
   status = pool_status(dev)
   if status[] === nothing
-      pool = memory_pool(dev)
-
       # allow the pool to use up all memory of this device
       attribute!(memory_pool(dev), MEMPOOL_ATTR_RELEASE_THRESHOLD, typemax(UInt64))
 
@@ -111,7 +101,7 @@ function pool_cleanup()
       status = pool_status(dev)
       status[] === nothing && continue
 
-      if status[]
+      if status[]::Bool
         idle_counters[i] = 0
       else
         idle_counters[i] += 1
@@ -289,8 +279,8 @@ end
 Allocate a number of bytes `sz` from the memory pool. Returns a buffer object; may throw
 an [`OutOfGPUMemoryError`](@ref) if the allocation request cannot be satisfied.
 """
-@inline @timeit_ci alloc(sz::Integer; kwargs...) = alloc(Mem.DeviceBuffer, sz; kwargs...)
-@inline @timeit_ci function alloc(::Type{B}, sz; stream::Union{Nothing,CuStream}=nothing) where {B<:Mem.AbstractBuffer}
+@inline alloc(sz::Integer; kwargs...) = alloc(Mem.DeviceBuffer, sz; kwargs...)
+@inline function alloc(::Type{B}, sz; stream::Union{Nothing,CuStream}=nothing) where {B<:Mem.AbstractBuffer}
   # 0-byte allocations shouldn't hit the pool
   sz == 0 && return B()
 
@@ -334,8 +324,8 @@ end
 
 Releases a buffer `buf` to the memory pool.
 """
-@inline @timeit_ci function free(buf::Mem.AbstractBuffer;
-                                 stream::Union{Nothing,CuStream}=nothing)
+@inline function free(buf::Mem.AbstractBuffer;
+                      stream::Union{Nothing,CuStream}=nothing)
   # XXX: have @timeit use the root timer, since we may be called from a finalizer
 
   # 0-byte allocations shouldn't hit the pool
diff --git a/test/Project.toml b/test/Project.toml
@@ -7,6 +7,7 @@ Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
diff --git a/test/jet.jl b/test/jet.jl
diff --git a/test/setup.jl b/test/setup.jl
