JuliaDiff / JuliaDiff/ForwardDiff.jl

Unsupported operation when using CUDA

Open
#713 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Julia
Stars
1k
Forks
160
No merged PRs in 30d

Description

I previously raised this issue here, https://github.com/FluxML/Zygote.jl/issues/1532, but was recommended it would be more appropriate here.

My environment. I have seen the same issue on Linux machines
Julia Version 1.11.0
Commit 501a4f25c2 (2024-10-07 11:40 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: Windows (x86_64-w64-mingw32)
CPU: 8 × Intel(R) Core(TM) i7-1065G7 CPU @ 1.30GHz
WORD_SIZE: 64
LLVM: libLLVM-16.0.6 (ORCJIT, icelake-client)
Threads: 1 default, 0 interactive, 1 GC (on 8 virtual cores)
Status C:\Users\salbe\OneDrive\Documents\Research\JuliaBugs\Project.toml
[052768ef] CUDA v5.5.2
[7073ff75] IJulia v1.25.0
[e88e6eb3] Zygote v0.6.71

The example:

using CUDA, Zygote, LinearAlgebra  # [edit -- added using + code block]

f₁(x) = sum(abs2, exp.(log.(x) .* (1:length(x))))
f₂(x) = sum(abs2, x.^(1:length(x)))
x = randn(ComplexF64, 5);
z = CuArray{ComplexF64}(x);

# **Check the gradient calculations are consistent between the 2 functons**
test₁ = Zygote.gradient(f₁, x)[1]
test₂ = Zygote.gradient(f₂, x)[1]
norm(test₁ - test₂) / norm(test₁)
# **Output:**  2.2530284453414604e-16 <-- This is reasonable

# **Check the calculation using CUDA**
test₃ = Zygote.gradient(f₁, z)[1];
norm(test₁ - Array(test₃))/ norm(test₁)
# **Output:** 2.0454901873585542e-16

# **However, using f₂ generates an exception**
test₄ = Zygote.gradient(f₂, z)

Output:

InvalidIRError: compiling MethodInstance for (::GPUArrays.var"#34#36")(::CUDA.CuKernelContext, ::CuDeviceVector{GPUArrays.BrokenBroadcast{Union{}}, 1}, ::Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{1, CUDA.DeviceMemory}, Tuple{Base.OneTo{Int64}}, Zygote.var"#1409#1410"{typeof(^)}, Tuple{Base.Broadcast.Extruded{CuDeviceVector{ComplexF64, 1}, Tuple{Bool}, Tuple{Int64}}, Base.Broadcast.Extruded{UnitRange{Int64}, Tuple{Bool}, Tuple{Int64}}}}, ::Int64) resulted in invalid LLVM IR
Reason: unsupported dynamic function invocation (call to ≺(a, b) @ ForwardDiff [C:\Users\salbe\.julia\packages\ForwardDiff\PcZ48\src\dual.jl:54](file:///C:/Users/salbe/.julia/packages/ForwardDiff/PcZ48/src/dual.jl#line=53))
Stacktrace:
 [1] promote_rule
   @ [C:\Users\salbe\.julia\packages\ForwardDiff\PcZ48\src\dual.jl:407](file:///C:/Users/salbe/.julia/packages/ForwardDiff/PcZ48/src/dual.jl#line=406)
 [2] promote_type
   @ [.\promotion.jl:318](http://localhost:8888/promotion.jl#line=317)
 [3] ^
   @ [.\complex.jl:886](http://localhost:8888/complex.jl#line=885)
 [4] #1409
   @ [C:\Users\salbe\.julia\packages\Zygote\Tt5Gx\src\lib\broadcast.jl:276](file:///C:/Users/salbe/.julia/packages/Zygote/Tt5Gx/src/lib/broadcast.jl#line=275)
 [5] _broadcast_getindex_evalf
   @ [.\broadcast.jl:673](http://localhost:8888/broadcast.jl#line=672)
 [6] _broadcast_getindex
   @ [.\broadcast.jl:646](http://localhost:8888/broadcast.jl#line=645)
 [7] getindex
   @ [.\broadcast.jl:605](http://localhost:8888/broadcast.jl#line=604)
 [8] #34
   @ [C:\Users\salbe\.julia\packages\GPUArrays\qt4ax\src\host\broadcast.jl:59](file:///C:/Users/salbe/.julia/packages/GPUArrays/qt4ax/src/host/broadcast.jl#line=58)
Hint: catch this exception as `err` and call `code_typed(err; interactive = true)` to introspect the erronous code with Cthulhu.jl

Stacktrace:
  [1] check_ir(job::GPUCompiler.CompilerJob{GPUCompiler.PTXCompilerTarget, CUDA.CUDACompilerParams}, args::LLVM.Module)
    @ GPUCompiler [C:\Users\salbe\.julia\packages\GPUCompiler\2CW9L\src\validation.jl:147](file:///C:/Users/salbe/.julia/packages/GPUCompiler/2CW9L/src/validation.jl#line=146)
  [2] macro expansion
    @ [C:\Users\salbe\.julia\packages\GPUCompiler\2CW9L\src\driver.jl:382](file:///C:/Users/salbe/.julia/packages/GPUCompiler/2CW9L/src/driver.jl#line=381) [inlined]
  [3] macro expansion
    @ [C:\Users\salbe\.julia\packages\TimerOutputs\NRdsv\src\TimerOutput.jl:253](file:///C:/Users/salbe/.julia/packages/TimerOutputs/NRdsv/src/TimerOutput.jl#line=252) [inlined]
  [4] macro expansion
    @ [C:\Users\salbe\.julia\packages\GPUCompiler\2CW9L\src\driver.jl:381](file:///C:/Users/salbe/.julia/packages/GPUCompiler/2CW9L/src/driver.jl#line=380) [inlined]
  [5] emit_llvm(job::GPUCompiler.CompilerJob; toplevel::Bool, libraries::Bool, optimize::Bool, cleanup::Bool, validate::Bool, only_entry::Bool)
    @ GPUCompiler [C:\Users\salbe\.julia\packages\GPUCompiler\2CW9L\src\utils.jl:108](file:///C:/Users/salbe/.julia/packages/GPUCompiler/2CW9L/src/utils.jl#line=107)
  [6] emit_llvm
    @ [C:\Users\salbe\.julia\packages\GPUCompiler\2CW9L\src\utils.jl:106](file:///C:/Users/salbe/.julia/packages/GPUCompiler/2CW9L/src/utils.jl#line=105) [inlined]
  [7] codegen(output::Symbol, job::GPUCompiler.CompilerJob; toplevel::Bool, libraries::Bool, optimize::Bool, cleanup::Bool, validate::Bool, strip::Bool, only_entry::Bool, parent_job::Nothing)
    @ GPUCompiler [C:\Users\salbe\.julia\packages\GPUCompiler\2CW9L\src\driver.jl:100](file:///C:/Users/salbe/.julia/packages/GPUCompiler/2CW9L/src/driver.jl#line=99)
  [8] codegen
    @ [C:\Users\salbe\.julia\packages\GPUCompiler\2CW9L\src\driver.jl:82](file:///C:/Users/salbe/.julia/packages/GPUCompiler/2CW9L/src/driver.jl#line=81) [inlined]
  [9] compile(target::Symbol, job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
    @ GPUCompiler [C:\Users\salbe\.julia\packages\GPUCompiler\2CW9L\src\driver.jl:79](file:///C:/Users/salbe/.julia/packages/GPUCompiler/2CW9L/src/driver.jl#line=78)
 [10] compile
    @ [C:\Users\salbe\.julia\packages\GPUCompiler\2CW9L\src\driver.jl:74](file:///C:/Users/salbe/.julia/packages/GPUCompiler/2CW9L/src/driver.jl#line=73) [inlined]
 [11] #1145
    @ [C:\Users\salbe\.julia\packages\CUDA\2kjXI\src\compiler\compilation.jl:250](file:///C:/Users/salbe/.julia/packages/CUDA/2kjXI/src/compiler/compilation.jl#line=249) [inlined]
 [12] JuliaContext(f::CUDA.var"#1145#1148"{GPUCompiler.CompilerJob{GPUCompiler.PTXCompilerTarget, CUDA.CUDACompilerParams}}; kwargs::@Kwargs{})
    @ GPUCompiler [C:\Users\salbe\.julia\packages\GPUCompiler\2CW9L\src\driver.jl:34](file:///C:/Users/salbe/.julia/packages/GPUCompiler/2CW9L/src/driver.jl#line=33)
 [13] JuliaContext(f::Function)
    @ GPUCompiler [C:\Users\salbe\.julia\packages\GPUCompiler\2CW9L\src\driver.jl:25](file:///C:/Users/salbe/.julia/packages/GPUCompiler/2CW9L/src/driver.jl#line=24)
 [14] compile(job::GPUCompiler.CompilerJob)
    @ CUDA [C:\Users\salbe\.julia\packages\CUDA\2kjXI\src\compiler\compilation.jl:249](file:///C:/Users/salbe/.julia/packages/CUDA/2kjXI/src/compiler/compilation.jl#line=248)
 [15] actual_compilation(cache::Dict{Any, CuFunction}, src::Core.MethodInstance, world::UInt64, cfg::GPUCompiler.CompilerConfig{GPUCompiler.PTXCompilerTarget, CUDA.CUDACompilerParams}, compiler::typeof(CUDA.compile), linker::typeof(CUDA.link))
    @ GPUCompiler [C:\Users\salbe\.julia\packages\GPUCompiler\2CW9L\src\execution.jl:237](file:///C:/Users/salbe/.julia/packages/GPUCompiler/2CW9L/src/execution.jl#line=236)
 [16] cached_compilation(cache::Dict{Any, CuFunction}, src::Core.MethodInstance, cfg::GPUCompiler.CompilerConfig{GPUCompiler.PTXCompilerTarget, CUDA.CUDACompilerParams}, compiler::Function, linker::Function)
    @ GPUCompiler [C:\Users\salbe\.julia\packages\GPUCompiler\2CW9L\src\execution.jl:151](file:///C:/Users/salbe/.julia/packages/GPUCompiler/2CW9L/src/execution.jl#line=150)
 [17] macro expansion
    @ [C:\Users\salbe\.julia\packages\CUDA\2kjXI\src\compiler\execution.jl:380](file:///C:/Users/salbe/.julia/packages/CUDA/2kjXI/src/compiler/execution.jl#line=379) [inlined]
 [18] macro expansion
    @ [.\lock.jl:273](http://localhost:8888/lock.jl#line=272) [inlined]
 [19] cufunction(f::GPUArrays.var"#34#36", tt::Type{Tuple{CUDA.CuKernelContext, CuDeviceVector{GPUArrays.BrokenBroadcast{Union{}}, 1}, Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{1, CUDA.DeviceMemory}, Tuple{Base.OneTo{Int64}}, Zygote.var"#1409#1410"{typeof(^)}, Tuple{Base.Broadcast.Extruded{CuDeviceVector{ComplexF64, 1}, Tuple{Bool}, Tuple{Int64}}, Base.Broadcast.Extruded{UnitRange{Int64}, Tuple{Bool}, Tuple{Int64}}}}, Int64}}; kwargs::@Kwargs{})
    @ CUDA [C:\Users\salbe\.julia\packages\CUDA\2kjXI\src\compiler\execution.jl:375](file:///C:/Users/salbe/.julia/packages/CUDA/2kjXI/src/compiler/execution.jl#line=374)
 [20] cufunction
    @ [C:\Users\salbe\.julia\packages\CUDA\2kjXI\src\compiler\execution.jl:372](file:///C:/Users/salbe/.julia/packages/CUDA/2kjXI/src/compiler/execution.jl#line=371) [inlined]
 [21] macro expansion
    @ [C:\Users\salbe\.julia\packages\CUDA\2kjXI\src\compiler\execution.jl:112](file:///C:/Users/salbe/.julia/packages/CUDA/2kjXI/src/compiler/execution.jl#line=111) [inlined]
 [22] #launch_heuristic#1200
    @ [C:\Users\salbe\.julia\packages\CUDA\2kjXI\src\gpuarrays.jl:17](file:///C:/Users/salbe/.julia/packages/CUDA/2kjXI/src/gpuarrays.jl#line=16) [inlined]
 [23] launch_heuristic
    @ [C:\Users\salbe\.julia\packages\CUDA\2kjXI\src\gpuarrays.jl:15](file:///C:/Users/salbe/.julia/packages/CUDA/2kjXI/src/gpuarrays.jl#line=14) [inlined]
 [24] _copyto!
    @ [C:\Users\salbe\.julia\packages\GPUArrays\qt4ax\src\host\broadcast.jl:78](file:///C:/Users/salbe/.julia/packages/GPUArrays/qt4ax/src/host/broadcast.jl#line=77) [inlined]
 [25] copyto!
    @ [C:\Users\salbe\.julia\packages\GPUArrays\qt4ax\src\host\broadcast.jl:44](file:///C:/Users/salbe/.julia/packages/GPUArrays/qt4ax/src/host/broadcast.jl#line=43) [inlined]
 [26] copy
    @ [C:\Users\salbe\.julia\packages\GPUArrays\qt4ax\src\host\broadcast.jl:29](file:///C:/Users/salbe/.julia/packages/GPUArrays/qt4ax/src/host/broadcast.jl#line=28) [inlined]
 [27] materialize(bc::Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{1, CUDA.DeviceMemory}, Nothing, Zygote.var"#1409#1410"{typeof(^)}, Tuple{CuArray{ComplexF64, 1, CUDA.DeviceMemory}, UnitRange{Int64}}})
    @ Base.Broadcast [.\broadcast.jl:867](http://localhost:8888/broadcast.jl#line=866)
 [28] broadcast_forward(::Function, ::CuArray{ComplexF64, 1, CUDA.DeviceMemory}, ::UnitRange{Int64})
    @ Zygote [C:\Users\salbe\.julia\packages\Zygote\Tt5Gx\src\lib\broadcast.jl:282](file:///C:/Users/salbe/.julia/packages/Zygote/Tt5Gx/src/lib/broadcast.jl#line=281)
 [29] adjoint
    @ [C:\Users\salbe\.julia\packages\Zygote\Tt5Gx\src\lib\broadcast.jl:361](file:///C:/Users/salbe/.julia/packages/Zygote/Tt5Gx/src/lib/broadcast.jl#line=360) [inlined]
 [30] _pullback(::Zygote.Context{false}, ::typeof(Base.Broadcast.broadcasted), ::CUDA.CuArrayStyle{1, CUDA.DeviceMemory}, ::Function, ::CuArray{ComplexF64, 1, CUDA.DeviceMemory}, ::UnitRange{Int64})
    @ Zygote [C:\Users\salbe\.julia\packages\ZygoteRules\M4xmc\src\adjoint.jl:67](file:///C:/Users/salbe/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl#line=66)
 [31] _apply(::Function, ::Vararg{Any})
    @ Core [.\boot.jl:946](http://localhost:8888/boot.jl#line=945)
 [32] adjoint
    @ [C:\Users\salbe\.julia\packages\Zygote\Tt5Gx\src\lib\lib.jl:203](file:///C:/Users/salbe/.julia/packages/Zygote/Tt5Gx/src/lib/lib.jl#line=202) [inlined]
 [33] _pullback
    @ [C:\Users\salbe\.julia\packages\ZygoteRules\M4xmc\src\adjoint.jl:67](file:///C:/Users/salbe/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl#line=66) [inlined]
 [34] broadcasted
    @ [.\broadcast.jl:1326](http://localhost:8888/broadcast.jl#line=1325) [inlined]
 [35] f₂
    @ [.\In](http://localhost:8888/In)[3]:2 [inlined]
 [36] _pullback(ctx::Zygote.Context{false}, f::typeof(f₂), args::CuArray{ComplexF64, 1, CUDA.DeviceMemory})
    @ Zygote [C:\Users\salbe\.julia\packages\Zygote\Tt5Gx\src\compiler\interface2.jl:0](file:///C:/Users/salbe/.julia/packages/Zygote/Tt5Gx/src/compiler/interface2.jl#line=-1)
 [37] pullback(f::Function, cx::Zygote.Context{false}, args::CuArray{ComplexF64, 1, CUDA.DeviceMemory})
    @ Zygote [C:\Users\salbe\.julia\packages\Zygote\Tt5Gx\src\compiler\interface.jl:90](file:///C:/Users/salbe/.julia/packages/Zygote/Tt5Gx/src/compiler/interface.jl#line=89)
 [38] pullback
    @ [C:\Users\salbe\.julia\packages\Zygote\Tt5Gx\src\compiler\interface.jl:88](file:///C:/Users/salbe/.julia/packages/Zygote/Tt5Gx/src/compiler/interface.jl#line=87) [inlined]
 [39] gradient(f::Function, args::CuArray{ComplexF64, 1, CUDA.DeviceMemory})
    @ Zygote [C:\Users\salbe\.julia\packages\Zygote\Tt5Gx\src\compiler\interface.jl:147](file:///C:/Users/salbe/.julia/packages/Zygote/Tt5Gx/src/compiler/interface.jl#line=146)
 [40] top-level scope
    @ In[7]:1

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reproducing the CUDA/Zygote example and inspect ForwardDiff's src/dual.jl, especially the promotion_rule entries around lines 54 and 407 shown in the stack trace. Done means the f₂ gradient on a CuArray compiles without InvalidIRError and remains consistent with the f₁ result.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.