EnzymeAD / EnzymeAD/Enzyme.jl

Missing rules for FFTW `unsafe_execute`

Open
#1,717 8 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
586
Forks
108
Avg merge
1d 5h
Merged PRs (30d)
44

Description

Here is a MWE

```julia
using Enzyme
using FFTW

Enzyme.EnzymeRules.inactive(::typeof(FFTW.assert_applicable), args...) = nothing

function test(p, x)
p*x
return sum(abs2, x)
end

x = rand(ComplexF64, 256)
p = plan_fft!(x)

test(p, x)

dx = similar(x)
autodiff(Enzyme.Reverse, test, Active, Const(p), Duplicated(x, dx))

Enzyme compilation failed.

No augmented forward pass found for ijl_lazy_load_and_lookup
at context: %fftw_execute_dft.found = call void ()* @ijl_lazy_load_and_lookup({} addrspace(10)* nonnull %11, i8* noundef getelementptr inbounds ([17 x i8], [17 x i8]* @_j_str1, i32 0, i32 0)) #21, !dbg !41

Stacktrace:
[1] unsafe_execute!
@ ~/.julia/packages/FFTW/6nZei/src/fft.jl:518
[2] *
@ ~/.julia/packages/FFTW/6nZei/src/fft.jl:835
[3] *
@ ~/.julia/packages/AbstractFFTs/4iQz5/src/definitions.jl:224
[4] test
@ ~/Research/EnzymeTest/fft.jl:8

Stacktrace:
[1] throwerr(cstr::Cstring)
@ Enzyme.Compiler ~/.julia/packages/Enzyme/aEyGD/src/compiler.jl:1696
[2] unsafe_execute!
@ ~/.julia/packages/FFTW/6nZei/src/fft.jl:518 [inlined]
[3] *
@ ~/.julia/packages/FFTW/6nZei/src/fft.jl:835 [inlined]
[4] *
@ ~/.julia/packages/AbstractFFTs/4iQz5/src/definitions.jl:224 [inlined]
[5] test
@ ~/Research/EnzymeTest/fft.jl:8 [inlined]
[6] diffejulia_test_5412wrap
@ ~/Research/EnzymeTest/fft.jl:0
[7] macro expansion
@ ~/.julia/packages/Enzyme/aEyGD/src/compiler.jl:6673 [inlined]
[8] enzyme_call
@ ~/.julia/packages/Enzyme/aEyGD/src/compiler.jl:6273 [inlined]
[9] CombinedAdjointThunk
@ ~/.julia/packages/Enzyme/aEyGD/src/compiler.jl:6150 [inlined]
[10] autodiff
@ ~/.julia/packages/Enzyme/aEyGD/src/Enzyme.jl:314 [inlined]
[11] autodiff(::ReverseMode{…}, ::typeof(test), ::Type{…}, ::Const{…}, ::Duplicated{…})
@ Enzyme ~/.julia/packages/Enzyme/aEyGD/src/Enzyme.jl:326
[12] top-level scope
@ ~/Research/EnzymeTest/fft.jl:20
Some type information was truncated. Use `show(err)` to see complete types.
```

The inactive rule is to prevent another `ijl_lazy_load_and_lookup` error, but that's easy to beat since it shouldn't be differentiated.

Unless I misunderstand something, this is another missing rule. If so, I'll take a crack at writing some rules in the next couple of days.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.