autodiff(...) fails with simple product sum including CuArrays and Float64

[simenhu]

During my testing of Enzyme and getting to know the library I tried to make a simple function to autodiff. Running this code makes whole Julia executable to crash.

using CUDA
using Enzyme

arr_type = CuArray

f(x) =  x'*x + 1.0

x = arr_type([1.0, 2.0])

dx = arr_type([0.0, 0.0])

@info "Running function"
res = f(x)
display(res)

@info "Calculating gradients"
grad = autodiff(Reverse, f, Active, Duplicated(x, dx))

display(grad)
display(dx)

Error log (from terminal):

  Activating project at `~/.julia/environments/v1.10`
  Activating project at `~/programming/birdsview/BirdsviewGPR.jl`
[ Info: Loading package Revise
[ Info: Loading package Plots
[ Info: Loading package OhMyREPL
[ Info: Loading package Infiltrator
[ Info: Running function
6.0
[ Info: Calculating gradients
JIT session error: Symbols not found: [ cudaMemset ]
JIT session error: Symbols not found: [ cudaMemset ]
JIT session error: Symbols not found: [ cudaMemset ]
JIT session error: Symbols not found: [ cudaMemset ]
JIT session error: Symbols not found: [ cudaMemset ]
JIT session error: Symbols not found: [ cudaMemset ]
ERROR: LoadError: LLVM error: Failed to materialize symbols: { (JuliaExternal, { libname_cuMemPoolSetAccess_15072, libname_cublasDdot_v2_13899, libname_cuDeviceGetAttribute_14558, libname_cuCtxGetCurrent_13963, libname_cuMemPoolTrimTo_13919, libname_cublasDdot_v2_13915, libname_cublasSetStream_v2_15235, libname_cublasGetProperty_15204, libname_cublasSetStream_v2_15225, libname_cuStreamGetCaptureInfo_15135, libname_cuCtxGetId_14601, diffejulia_f_13864wrap, libname_cuStreamQuery_15049, libname_cuStreamCreate_14866, libname_cuMemPoolTrimTo_15154, libname_cuMemPoolSetAttribute_14710, libname_cudaRuntimeGetVersion_14128, libname_cuCtxGetCurrent_14642, libname_cublasSetStream_v2_15249, libname_cuDevicePrimaryCtxRetain_14003, libname_cuMemPoolTrimTo_15238, libname_cuDeviceGetName_14136, libname_cuDeviceGet_14574, libname_cuCtxGetId_14011, libname_cublasSetMathMode_15172, libname_cuMemPoolCreate_14721, libname_cuMemPoolTrimTo_15186, libname_cuCtxGetDevice_15044, libname_cuCtxSynchronize_14889, libname_cuCtxGetApiVersion_14603, libname_cuCtxSetCurrent_13965, libname_cuCtxSynchronize_14883, libname_cublasSetMathMode_15183, libname_cuDeviceGetCount_13980, libname_cuCtxSetCurrent_14588, libname_cuStreamSynchronize_15112, libname_cublasDdot_v2_64_15151, libname_cublasGetProperty_15195, libname_cuMemPoolTrimTo_15208, libname_cuMemPoolGetAttribute_13955, libname_cuDeviceSetMemPool_14715, libname_cuStreamSynchronize_15117, libname_cuCtxPushCurrent_v2_15038, libname_cuStreamQuery_15050, libname_cublasDdot_v2_64_13894, libname_cuDeviceCanAccessPeer_15083, libname_cuMemGetInfo_v2_14799, libname_cuMemPoolCreate_14636, libname_cuCtxPopCurrent_v2_15041 }) }
Stacktrace:
  [1] macro expansion
    @ ~/.julia/packages/LLVM/wMjUU/src/executionengine/utils.jl:28 [inlined]
  [2] lookup
    @ ~/.julia/packages/LLVM/wMjUU/src/orc.jl:440 [inlined]
  [3] lookup
    @ ~/.julia/packages/LLVM/wMjUU/src/orc.jl:439 [inlined]
  [4] lookup
    @ ~/.julia/packages/Enzyme/ydGh2/src/compiler/orcv2.jl:269 [inlined]
  [5] _link(job::GPUCompiler.CompilerJob{<:Enzyme.Compiler.EnzymeTarget}, mod::LLVM.Module, edges::Vector{Any}, adjoint_name::String, primal_name::Union{Nothing, String}, TapeType::Any, prepost::String)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:5255
  [6] cached_compilation
    @ ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:5341 [inlined]
  [7] thunkbase(mi::Core.MethodInstance, World::UInt64, FA::Type{<:Annotation}, A::Type{<:Annotation}, TT::Type, Mode::Enzyme.API.CDerivativeMode, width::Int64, ModifiedBetween::Tuple{Vararg{Bool, N}} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, edges::Vector{Any})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:5451
  [8] thunk_generator(world::UInt64, source::LineNumberNode, FA::Type, A::Type, TT::Type, Mode::Enzyme.API.CDerivativeMode, Width::Int64, ModifiedBetween::Tuple{Vararg{Bool, N}} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, self::Any, fakeworld::Any, fa::Type, a::Type, tt::Type, mode::Type, width::Type, modifiedbetween::Type, returnprimal::Type, shadowinit::Type, abi::Type, erriffuncwritten::Type, runtimeactivity::Type)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/ydGh2/src/compiler.jl:5636
  [9] autodiff
    @ ~/.julia/packages/Enzyme/ydGh2/src/Enzyme.jl:485 [inlined]
 [10] autodiff(mode::ReverseMode{false, false, FFIABI, false, false}, f::typeof(f), ::Type{Active}, args::Duplicated{CuArray{Float64, 1, CUDA.DeviceMemory}})
    @ Enzyme ~/.julia/packages/Enzyme/ydGh2/src/Enzyme.jl:524
 [11] top-level scope
    @ ~/programming/birdsview/BirdsviewGPR.jl/MWE/simple_cuda_enzyme_mwe.jl:17
in expression starting at /home/simen/programming/birdsview/BirdsviewGPR.jl/MWE/simple_cuda_enzyme_mwe.jl:17

Relevant packages:

  [7da242da] Enzyme v0.13.24
  [052768ef] CUDA v5.5.2

Version info:

Julia Version 1.10.7
Commit 4976d05258e (2024-11-26 15:57 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 256 × AMD EPYC 7H12 64-Core Processor
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, znver2)
Threads: 1 default, 0 interactive, 1 GC (on 256 virtual cores)
Environment:
  JULIA_SSL_NO_VERIFY_HOSTS = github.com, julialang-s3.julialang.org, pkg.julialang.org
  LD_LIBRARY_PATH = /usr/local/cuda-11.8/lib64

[wsmoses]

@vchuravy so we appear to have a symbol resolution problem with CUDA.jl it seems