Error differentiating past FFT

[vpuri3]

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

[vchuravy]

Custom rules are not exposed to the user in a nice and friendly way see #177 for some early thoughts on that topic.

[wsmoses]

@vpuri3 can you attach the entire test code which causes the issue. Additionally, can you retry with latest main?

[wsmoses]

Closing from insufficient information. Please reopen if it persists.