Enzyme.jl
Enzyme.jl copied to clipboard
EnzymeAD
Error differentiating past FFT
Hello, I was trying to AD some code that leverages FFTW.jl and cam across this error. How would one go about writing custom rule for FFTs? Note that chain rules are defined: https://github.com/JuliaMath/AbstractFFTs.jl/blob/master/src/chainrules.jl
julia> include("examples/opt.jl")
fwd
1280.0
bwd
ERROR: LoadError: InvalidIRError: compiling function #mul!(Vector{ComplexF64}, FFTW.rFFTWPlan{F
loat64, -1, false, 1, UnitRange{Int64}}, Vector{Float64}) resulted in invalid LLVM IR
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
[1] alignment_of
@ ~/.julia/packages/FFTW/sfy1o/src/fft.jl:236
[2] assert_applicable
@ ~/.julia/packages/FFTW/sfy1o/src/fft.jl:464
[3] mul!
@ ~/.julia/packages/FFTW/sfy1o/src/fft.jl:858
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
[1] alignment_of
@ ~/.julia/packages/FFTW/sfy1o/src/fft.jl:236
[2] assert_applicable
@ ~/.julia/packages/FFTW/sfy1o/src/fft.jl:453
[3] assert_applicable
@ ~/.julia/packages/FFTW/sfy1o/src/fft.jl:459
[4] mul!
@ ~/.julia/packages/FFTW/sfy1o/src/fft.jl:858
Reason: unsupported jl_lazy_load_and_lookup
Stacktrace:
[1] unsafe_execute!
@ ~/.julia/packages/FFTW/sfy1o/src/fft.jl:500
[2] mul!
@ ~/.julia/packages/FFTW/sfy1o/src/fft.jl:859
Hint: catch this exception as `err` and call `code_typed(err; interactive = true)` to introspec
t the erronous code
Stacktrace:
[1] check_ir
@ ~/.julia/packages/Enzyme/Wanbg/src/compiler/validation.jl:135 [inlined]
[2] codegen(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.
Compiler.EnzymeCompilerParams, GPUCompiler.FunctionSpec{typeof(mul!), Tuple{Vector{ComplexF64},
FFTW.rFFTWPlan{Float64, -1, false, 1, UnitRange{Int64}}, Vector{Float64}}}}; libraries::Bool,
deferred_codegen::Bool, optimize::Bool, ctx::LLVM.Context, strip::Bool, validate::Bool, only_en
try::Bool, parent_job::Nothing)
@ Enzyme.Compiler ~/.julia/packages/Enzyme/Wanbg/src/compiler.jl:3901
[3] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCo
mpilerParams, GPUCompiler.FunctionSpec{typeof(mul!), Tuple{Vector{ComplexF64}, FFTW.rFFTWPlan{Float64, -1, false, 1, UnitRange{Int64}}, Vector{Float64}}}})
@ Enzyme.Compiler ~/.julia/packages/Enzyme/Wanbg/src/compiler.jl:4562
Custom rules are not exposed to the user in a nice and friendly way see #177 for some early thoughts on that topic.
@vpuri3 can you attach the entire test code which causes the issue. Additionally, can you retry with latest main?
