search

EnzymeAD / Enzyme.jl

Julia bindings for the Enzyme automatic differentiator
https://enzyme.mit.edu
MIT License
443 stars 62 forks source link

Concatenating AD for multiple GPU kernels #415

Closed renatobellotti closed 2 years ago

renatobellotti commented 2 years ago

Is it currently possible to concatenate AD calls over multiple GPU kernels? I'm currently trying to implement this on a toy problem and can't figure out how to do it.

using ChainRulesCore
using CUDA
using CUDAKernels
using Enzyme
using KernelAbstractions
using KernelGradients
using Zygote

@kernel function example_kernel(x, y, z)
    i = @index(Global)
    if(i == 1)
        z[i] = 2 * x[i] + y[i]
    elseif (i == 2)
        z[i] = 3 * x[i] + y[i]
    elseif (i == 3)
        z[i] = 4 * x[i] + y[i]
    elseif (i == 4)
        z[i] = 5 * x[i] + y[i]
    end
    nothing
end

@kernel function example_kernel2(z, a, result)
    i = @index(Global)
    result[i] = 3 * z[i] + a[i]
    nothing
end

function my_call!(x, y, a)
    z = cu(zeros(4))
    result = cu(zeros(4))

    kernel = example_kernel(CUDADevice())
    kernel2 = example_kernel2(CUDADevice())

    event = kernel(x, y, z, ndrange=4)
    wait(event)

    event = kernel2(z, a, result, ndrange=4)
    wait(event)

    return result
end

dx = Duplicated(x, cu(zeros(Float32, 4)))
dy = Const(y)
dz = Duplicated(cu(zeros(Float32, 4)), cu(ones(Float32, 4)))
da = Const(a)

Enzyme.autodiff_deferred(my_call!, Const, dx, dy, da)

Running the code above fails and kills the kernel on a Jupyter notebook. Here is the error message:

InvalidIRError: compiling function #my_call!(CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}) resulted in invalid LLVM IR
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:59
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuDeviceGetAttribute
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] attribute
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/devices.jl:175
 [5] capability
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/devices.jl:194
 [6] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:41
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuDeviceGetName
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] name
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/devices.jl:86
 [5] macro expansion
   @ ./logging.jl:360
 [6] context
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:215
 [7] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] unsafe_cuEventQuery
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:1081
 [2] isdone
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/events.jl:90
 [3] isdone
   @ ~/.julia/packages/CUDAKernels/YWLqR/src/CUDAKernels.jl:105
 [4] #22
   @ ~/.julia/packages/KernelAbstractions/1ZLga/src/cpu.jl:42
 [5] ntuple
   @ ./ntuple.jl:19
 [6] wait
   @ ~/.julia/packages/KernelAbstractions/1ZLga/src/cpu.jl:38
 [7] wait (repeats 2 times)
   @ ~/.julia/packages/KernelAbstractions/1ZLga/src/cpu.jl:29
 [8] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:23
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuDriverGetVersion
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] version
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/version.jl:24
 [5] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:59
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuDeviceGetAttribute
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] attribute
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/devices.jl:175
 [5] memory_pools_supported
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/devices.jl:198
 [6] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/src/pool.jl:72
 [7] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/memoization.jl:78
 [8] stream_ordered
   @ ~/.julia/packages/CUDA/DfvRa/src/pool.jl:71
 [9] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] unsafe_cuCtxGetDevice
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:150
 [2] current_device
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/devices.jl:23
 [3] context!
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:150
 [4] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:138
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuCtxSetCurrent
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] activate
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:181
 [5] context!
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:149
 [6] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] unsafe_cuCtxGetDevice
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:150
 [2] current_device
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/devices.jl:23
 [3] context!
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:144
 [4] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:138
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuCtxSetCurrent
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] activate
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:181
 [5] context!
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:143
 [6] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:144
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuCtxGetCurrent
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] current_context
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:63
 [5] CuStream
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:26
 [6] CuStream
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:18
 [7] create_stream
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:354
 [8] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:934
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuStreamCreate
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] CuStream
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:20
 [5] CuStream
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:18
 [6] create_stream
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:354
 [7] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] nvtxNameCuStreamA
   @ ~/.julia/packages/CUDA/DfvRa/lib/nvtx/libnvtx.jl:244
 [2] create_stream
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:360
 [3] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:144
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuCtxGetCurrent
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] current_context
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:63
 [5] memory_pool
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/pool.jl:48
 [6] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:2277
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuDeviceGetMemPool
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] memory_pool
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/pool.jl:46
 [5] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
  [1] macro expansion
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:591
  [2] macro expansion
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
  [3] cuMemcpyHtoDAsync_v2
    @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
  [4] #unsafe_copyto!#9
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/memory.jl:397
  [5] #183
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:395
  [6] #context!#63
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:164
  [7] context!
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:161
  [8] unsafe_copyto!
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:388
  [9] copyto!
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:346
 [10] copyto!
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:350
 [11] CuArray
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:292
 [12] adapt_storage
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:539
 [13] multiple call sites
    @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
  [1] macro expansion
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:1254
  [2] macro expansion
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
  [3] cuLaunchHostFunc
    @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
  [4] #launch#46
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/execution.jl:159
  [5] nonblocking_synchronize
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:158
  [6] #183
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:392
  [7] #context!#63
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:164
  [8] context!
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:161
  [9] unsafe_copyto!
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:388
 [10] copyto!
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:346
 [11] copyto!
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:350
 [12] CuArray
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:292
 [13] adapt_storage
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:539
 [14] multiple call sites
    @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
  [1] unsafe_cuStreamQuery
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:1025
  [2] isdone
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:103
  [3] nonblocking_synchronize
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:152
  [4] #183
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:392
  [5] #context!#63
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:164
  [6] context!
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:161
  [7] unsafe_copyto!
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:388
  [8] copyto!
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:346
  [9] copyto!
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:350
 [10] CuArray
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:292
 [11] adapt_storage
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:539
 [12] multiple call sites
    @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
  [1] unsafe_cuStreamQuery
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:1025
  [2] isdone
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:103
  [3] nonblocking_synchronize
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:139
  [4] #183
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:392
  [5] #context!#63
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:164
  [6] context!
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:161
  [7] unsafe_copyto!
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:388
  [8] copyto!
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:346
  [9] copyto!
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:350
 [10] CuArray
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:292
 [11] adapt_storage
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:539
 [12] multiple call sites
    @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
  [1] unsafe_cuPointerGetAttribute
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:883
  [2] is_pinned
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/memory.jl:849
  [3] #183
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:392
  [4] #context!#63

    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:164
  [5] context!
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:161
  [6] unsafe_copyto!
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:388
  [7] copyto!
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:346
  [8] copyto!
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:350
  [9] CuArray
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:292
 [10] adapt_storage
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:539
 [11] multiple call sites
    @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:138
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuCtxSetCurrent
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] activate
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:181
 [5] prepare_cuda_state
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:96
 [6] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:144
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuCtxGetCurrent
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] prepare_cuda_state
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:94
 [5] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:345
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuMemGetInfo_v2
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] info
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/memory.jl:754
 [5] MemoryInfo
   @ ~/.julia/packages/CUDA/DfvRa/src/pool.jl:154
 [6] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:35
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuDeviceGetCount
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:84
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuDevicePrimaryCtxRetain
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] CuContext
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:57
 [5] context
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:222
 [6] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:138
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuCtxSetCurrent
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] activate
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:181
 [5] device!
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:264
 [6] device!
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:244
 [7] validate_task_local_state
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:61
 [8] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:29
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuDeviceGet
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] CuDevice
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/devices.jl:17
 [5] TaskLocalState
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:50
 [6] task_local_state!
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/state.jl:73
 [7] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:144
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuCtxGetCurrent
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] current_context
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:63
 [5] CuEvent
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/events.jl:21
 [6] Event
   @ ~/.julia/packages/CUDAKernels/YWLqR/src/CUDAKernels.jl:109
 [7] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:1067
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuEventCreate
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] CuEvent
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/events.jl:19
 [5] Event
   @ ~/.julia/packages/CUDAKernels/YWLqR/src/CUDAKernels.jl:109
 [6] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:1074
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuEventRecord
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] record
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/events.jl:44
 [5] Event
   @ ~/.julia/packages/CUDAKernels/YWLqR/src/CUDAKernels.jl:110
 [6] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:352
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuMemAlloc_v2
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] #alloc#1
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/memory.jl:86
 [5] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/src/pool.jl:41
 [6] macro expansion
   @ ./timing.jl:299
 [7] #actual_alloc#154
   @ ~/.julia/packages/CUDA/DfvRa/src/pool.jl:39
 [8] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:2267
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuMemAllocAsync

   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] #alloc#1
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/memory.jl:83
 [5] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/src/pool.jl:41
 [6] macro expansion
   @ ./timing.jl:299
 [7] #actual_alloc#154
   @ ~/.julia/packages/CUDA/DfvRa/src/pool.jl:39
 [8] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:132
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuCtxPopCurrent_v2
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] pop!
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:171
 [5] device
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:288
 [6] check_exceptions
   @ ~/.julia/packages/CUDA/DfvRa/src/compiler/exceptions.jl:33
 [7] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] unsafe_cuCtxGetDevice
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:150
 [2] current_device
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/devices.jl:23
 [3] device
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:287
 [4] check_exceptions
   @ ~/.julia/packages/CUDA/DfvRa/src/compiler/exceptions.jl:33
 [5] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:126
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuCtxPushCurrent_v2
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] push!
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:162
 [5] device
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:286
 [6] check_exceptions
   @ ~/.julia/packages/CUDA/DfvRa/src/compiler/exceptions.jl:33
 [7] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:1254
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuLaunchHostFunc
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] #launch#46
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/execution.jl:159
 [5] nonblocking_synchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:158
 [6] #synchronize#12
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:128
 [7] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] unsafe_cuStreamQuery
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:1025
 [2] isdone
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:103
 [3] nonblocking_synchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:152
 [4] #synchronize#12
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:128
 [5] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:1032
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuStreamSynchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] #synchronize#12
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:132
 [5] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] unsafe_cuStreamQuery
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:1025
 [2] isdone
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:103
 [3] nonblocking_synchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:139
 [4] #synchronize#12
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:128
 [5] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:1254
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuLaunchHostFunc
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] #launch#46
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/execution.jl:159
 [5] nonblocking_synchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:158
 [6] nonblocking_synchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:325
 [7] device_synchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:319
 [8] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] unsafe_cuStreamQuery
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:1025
 [2] isdone
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:103
 [3] nonblocking_synchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:152
 [4] nonblocking_synchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:325
 [5] device_synchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:319
 [6] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:164
 [2] macro expansion
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
 [3] cuCtxSynchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
 [4] nonblocking_synchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:329
 [5] device_synchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:319
 [6] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
 [1] unsafe_cuStreamQuery
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:1025
 [2] isdone
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:103
 [3] nonblocking_synchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/stream.jl:139
 [4] nonblocking_synchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:325
 [5] device_synchronize
   @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/context.jl:319
 [6] multiple call sites
   @ unknown:0
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
  [1] macro expansion
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/libcuda.jl:2272
  [2] macro expansion
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/error.jl:95
  [3] cuMemPoolTrimTo
    @ ~/.julia/packages/CUDA/DfvRa/lib/utils/call.jl:26
  [4] trim (repeats 2 times)
    @ ~/.julia/packages/CUDA/DfvRa/lib/cudadrv/pool.jl:66
  [5] macro expansion

    @ ~/.julia/packages/CUDA/DfvRa/src/pool.jl:257
  [6] macro expansion
    @ ./timing.jl:299
  [7] #_alloc#170
    @ ~/.julia/packages/CUDA/DfvRa/src/pool.jl:313
  [8] #alloc#169
    @ ~/.julia/packages/CUDA/DfvRa/src/pool.jl:299
  [9] alloc
    @ ~/.julia/packages/CUDA/DfvRa/src/pool.jl:295
 [10] CuArray
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:42
 [11] CuArray
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:291
 [12] adapt_storage
    @ ~/.julia/packages/CUDA/DfvRa/src/array.jl:539
 [13] multiple call sites
    @ unknown:0
Hint: catch this exception as `err` and call `code_typed(err; interactive = true)` to introspect the erronous code
Stacktrace:
  [1] check_ir
    @ ~/.julia/packages/Enzyme/di3zM/src/compiler/validation.jl:135 [inlined]
  [2] codegen(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams, GPUCompiler.FunctionSpec{typeof(my_call!), Tuple{CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, ctx::LLVM.Context, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/di3zM/src/compiler.jl:3901
  [3] _thunk
    @ ~/.julia/packages/Enzyme/di3zM/src/compiler.jl:4599 [inlined]
  [4] (::Enzyme.Compiler.JIT.var"#materialize#7"{GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams, GPUCompiler.FunctionSpec{typeof(my_call!), Tuple{CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}}, String, LLVM.LLJIT})(mr::LLVM.MaterializationResponsibility)
    @ Enzyme.Compiler.JIT ~/.julia/packages/Enzyme/di3zM/src/compiler/orcv2.jl:160
  [5] __materialize(ctx::Ptr{Nothing}, mr::Ptr{LLVM.API.LLVMOrcOpaqueMaterializationResponsibility})
    @ LLVM ~/.julia/packages/LLVM/WjSQG/src/orcv2.jl:300
  [6]

How can I run AD over multiple kernels? My goal is to call multiple kernels to calculate a vector, then calculate a scalar loss function and optimise it with gradient descent.

renatobellotti commented 2 years ago

Is chaining of kernels supported?

renatobellotti commented 2 years ago

I found out how to do it: https://github.com/JuliaDiff/ChainRules.jl/issues/665#issuecomment-1220772450