Failed to include Enzyme+CUDA test code as module

[junyixu]

I was trying to debug some CUDA AD code using my go-to workflow - putting test code in a module and including it in REPL. Something like:

module EnzymeTest
# This is from Enzyme.jl/test/cuda.jl
using CUDA
using Enzyme
using Test

function mul_kernel(A)
    i = threadIdx().x 
    if i <= length(A)
        A[i] *= A[i]
    end
    return nothing
end

function grad_mul_kernel(A, dA)
    autodiff_deferred(Reverse, Const(mul_kernel), Const, Duplicated(A, dA))
    return nothing
end

@testset "mul_kernel" begin
    A = CUDA.ones(64,)
    @cuda threads=length(A) mul_kernel(A)
    A = CUDA.ones(64,)
    dA = similar(A)
    dA .= 1
    @cuda threads=length(A) grad_mul_kernel(A, dA)
    @test all(dA .== 2)
end
end

But got a bunch of precompilation errors when doing include("EnzymeTest.jl").
Turns out removing the module EnzymeTest makes it work. Seems like there's some circular dependency stuff going on between CUDA.jl and Enzyme.jl during precompilation.
Noting this down in case I run into similar issues again. For now, just don't wrap CUDA+Enzyme test code in modules when including in REPL.

The output

julia> include("EnzymeTest.jl")
┌ Warning: Circular dependency detected. Precompilation will be skipped for:
│   SparseArraysExt [85068d23-b5fb-53f1-8204-05c2aba6942f]
│   EnzymeCoreExt [c3ed20f3-746d-5a71-99a7-df1cca8c2b90]
│   AtomixCUDAExt [13011619-4c7c-5ef0-948f-5fc81565cd05]
│   LinearAlgebraExt [66d79d19-2cc4-5b0b-ac7a-b340256d1ecd]
│   EnzymeExt [6425d1ab-ad86-5f6e-82f1-f5118e0dc800]
│   GPUArrays [0c68f7d7-f131-5f86-a1c3-88cf8149b2d7]
│   KernelAbstractions [63c18a36-062a-441e-b654-da1e3ab1ce7c]
│   CUDA [052768ef-5323-5732-b1bb-66c8b64840ba]
└ @ Pkg.API ~/.local/stow/julia-1.10.2/share/julia/stdlib/v1.10/Pkg/src/API.jl:1239
[ Info: Precompiling CUDA [052768ef-5323-5732-b1bb-66c8b64840ba]
┌ Warning: Module CUDA with build ID ffffffff-ffff-ffff-0000-0045d291662b is missing from the cache.
│ This may mean CUDA [052768ef-5323-5732-b1bb-66c8b64840ba] does not support precompilation but is imported by a module that does.
└ @ Base loading.jl:1948
┌ Error: Error during loading of extension AtomixCUDAExt of Atomix, use Base.retry_load_extensions() to retry.
│   exception =
│    1-element ExceptionStack:
│    Declaring precompile(false) is not allowed in files that are being precompiled.
│    Stacktrace:
│      [1] _require(pkg::Base.PkgId, env::Nothing)
│        @ Base ./loading.jl:1999
│      [2] require_prelocked(uuidkey::Base.PkgId, env::Nothing)
│        @ Base ./loading.jl:1812
│      [3] #invoke_in_world#3
│        @ ./essentials.jl:926 [inlined]
│      [4] invoke_in_world
│        @ ./essentials.jl:923 [inlined]
│      [5] requireprelocked
│        @ ./loading.jl:1803 [inlined]
│      [6] requireprelocked
│        @ ./loading.jl:1802 [inlined]
│      [7] run_extension_callbacks(extid::Base.ExtensionId)
│        @ Base ./loading.jl:1295
│      [8] run_extension_callbacks(pkgid::Base.PkgId)
│        @ Base ./loading.jl:1330
│      [9] run_package_callbacks(modkey::Base.PkgId)
│        @ Base ./loading.jl:1164
│     [10] tryrequirefrom_serialized(modkey::Base.PkgId, path::String, ocachepath::String, sourcepath::String, depmods::Vector{Any})
│        @ Base ./loading.jl:1487
│     [11] requiresearch_from_serialized(pkg::Base.PkgId, sourcepath::String, build_id::UInt128)
│        @ Base ./loading.jl:1574
│     [12] _require(pkg::Base.PkgId, env::String)
│        @ Base ./loading.jl:1938
│     [13] require_prelocked(uuidkey::Base.PkgId, env::String)
│        @ Base ./loading.jl:1812
│     [14] #invoke_in_world#3
│        @ ./essentials.jl:926 [inlined]│     [15] invoke_in_world
│        @ ./essentials.jl:923 [inlined]
│     [16] requireprelocked(uuidkey::Base.PkgId, env::String)
│        @ Base ./loading.jl:1803
│     [17] macro expansion
│        @ ./loading.jl:1790 [inlined]
│     [18] macro expansion
│        @ ./lock.jl:267 [inlined]
│     [19] __require(into::Module, mod::Symbol)
│        @ Base ./loading.jl:1753
│     [20] #invoke_in_world#3
│        @ ./essentials.jl:926 [inlined]
│     [21] invoke_in_world
│        @ ./essentials.jl:923 [inlined]
│     [22] require(into::Module, mod::Symbol)
│        @ Base ./loading.jl:1746
│     [23] include
│        @ ./Base.jl:495 [inlined]
│     [24] include_package_for_output(pkg::Base.PkgId, input::String, depot_path::Vector{String}, dl_load_path::Vector{String}, load_path::Vector{String}, concrete_deps::Vector{Pair{Base.PkgId, UInt128}}, source::String)
│        @ Base ./loading.jl:2222
│     [25] top-level scope
│        @ stdin:3
│     [26] eval
│        @ ./boot.jl:385 [inlined]
│     [27] include_string(mapexpr::typeof(identity), mod::Module, code::String, filename::String)
│        @ Base ./loading.jl:2076
│     [28] include_string
│        @ ./loading.jl:2086 [inlined]
│     [29] exec_options(opts::Base.JLOptions)
│        @ Base ./client.jl:316
│     [30] _start()
│        @ Base ./client.jl:552
└ @ Base loading.jl:1301
┌ Warning: Circular dependency detected. Precompilation will be skipped for:
│   SparseArraysExt [85068d23-b5fb-53f1-8204-05c2aba6942f]
│   EnzymeCoreExt [c3ed20f3-746d-5a71-99a7-df1cca8c2b90]
│   AtomixCUDAExt [13011619-4c7c-5ef0-948f-5fc81565cd05]
│   LinearAlgebraExt [66d79d19-2cc4-5b0b-ac7a-b340256d1ecd]
│   EnzymeExt [6425d1ab-ad86-5f6e-82f1-f5118e0dc800]
│   GPUArrays [0c68f7d7-f131-5f86-a1c3-88cf8149b2d7]
│   KernelAbstractions [63c18a36-062a-441e-b654-da1e3ab1ce7c]
│   CUDA [052768ef-5323-5732-b1bb-66c8b64840ba]
└ @ Pkg.API ~/.local/stow/julia-1.10.2/share/julia/stdlib/v1.10/Pkg/src/API.jl:1239
[ Info: Precompiling ConstructionBaseStaticArraysExt [8497ba20-d017-5d93-8a79-2639523b7219]
┌ Warning: Circular dependency detected. Precompilation will be skipped 

[junyixu]

It seems we have to wait trixi-gpu/TrixiCUDA.jl#89