EnzymeTestUtils fails for FFT plans / structs with undefined references

[maximilian-gelbrecht]

I am currently in the process of writing some custom rules for code that involves FFT plans. While testing those EnzymeTestUtils fails.

EnzymeTestUtils tests with test_approx if two variables are approximately equal. For structs like FFT it checks all of their fields. For most FFT plans this is bound to fail because FFT plans have an uninitialised field pinv for the inverse plan:

using AbstractFFTs, FFTW 
x = rand(10)
plan = plan_fft(x)
plan.pinv 

returns

ERROR: UndefRefError: access to undefined reference

pinv gets assigned once the adjoint/inverse is called the first time (e.g. with \). But the inverse again has a field pinv which has a field pinv and so on. For most purposes that's not a problem, but for EnzymeTestUtils it is currently.

So, as far as I can see it test_approx will always return an
ERROR: UndefRefError: access to undefined reference if an FFT plan is handed over

isdefined and isassigned don't work in this case unfortunately. For my personal use I just fixed that by defining

function EnzymeTestUtils.test_approx(x::AbstractFFTs.Plan, y::AbstractFFTs.Plan, msg; kwargs...)
    EnzymeTestUtils.@test_msg "$msg: types must match" typeof(x) == typeof(y)
    names = fieldnames(typeof(x))[1:end-1]
    if isempty(names)
        EnzymeTestUtils.@test_msg msg x == y
    else
        for k in names
            if k isa Symbol && hasproperty(x, k)
                msg_new = "$msg: ::$(typeof(x)).$k"
            else
                msg_new = "$msg: getfield(::$(typeof(x)), $k)"
            end
            EnzymeTestUtils.test_approx(getfield(x, k), getfield(y, k), msg_new; kwargs...)
        end
    end
    return nothing
end 

That's the same as the generic test_approx, it just excludes the last field, which is pinv. That' a quite hacky solution, I am sure there's a way to solve this in a better, more elegant way.

[wsmoses]

@sethaxen

[sethaxen]

Yeah I guess maybe we should not recur into these undefined fields.

[sethaxen]

I looked into this a little, and more changes are necessary to properly handle a case like this:

mutable struct MutatedStructWithUninitializedField
    x
    y  # field can be uninitialized
    MutatedStructWithUninitializedField(x) = new(x)
end

function f_uninit!(s::MutatedStructWithUninitializedField)
    x = s.x * 3
    if !isdefined(s, :y)
        s.y = s.x * 10
    end
    s.x = x
    return s
end

Even though f_uninit! returns the same object, the dimension of the object for the purposes of finite differences (i.e. the number of floats to reconstruct the object) is different. Even if we update EnzymeTestUtils._test_fields_approx to assert that both or neither object defines a given field, we still get an error like this:

julia> x = MutatedStructWithUninitializedField(1.0)
MutatedStructWithUninitializedField(1.0, #undef)

julia> test_reverse(f_uninit!, Duplicated, (x, Duplicated))
test_reverse: f_uninit! with return activity Duplicated on (::MutatedStructWithUninitializedField, Duplicated): Error During Test at /home/sethaxen/projects/Enzyme.jl/lib/EnzymeTestUtils/src/test_reverse.jl:85
  Got exception outside of a @test
  DimensionMismatch: second dimension of A, 4, does not match length of x, 3
  Stacktrace:
    [1] gemv!
      @ ~/.julia/juliaup/julia-1.11.1+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/matmul.jl:438 [inlined]
    [2] generic_matvecmul!
      @ ~/.julia/juliaup/julia-1.11.1+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/matmul.jl:79 [inlined]
    [3] _mul!
      @ ~/.julia/juliaup/julia-1.11.1+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/matmul.jl:73 [inlined]
    [4] mul!
      @ ~/.julia/juliaup/julia-1.11.1+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/matmul.jl:70 [inlined]
    [5] mul!
      @ ~/.julia/juliaup/julia-1.11.1+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/matmul.jl:253 [inlined]
    [6] *(A::LinearAlgebra.Transpose{Float64, Matrix{Float64}}, x::Vector{Float64})
      @ LinearAlgebra ~/.julia/juliaup/julia-1.11.1+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/matmul.jl:56
    [7] _j′vp(fdm::FiniteDifferences.AdaptedFiniteDifferenceMethod{5, 1, FiniteDifferences.UnadaptedFiniteDifferenceMethod{7, 5}}, f::Function, ȳ::Vector{Float64}, x::Vector{Float64})
      @ FiniteDifferences ~/.julia/packages/FiniteDifferences/IPGFN/src/grad.jl:84
    [8] j′vp(fdm::FiniteDifferences.AdaptedFiniteDifferenceMethod{5, 1, FiniteDifferences.UnadaptedFiniteDifferenceMethod{7, 5}}, f::ComposedFunction{ComposedFunction{ComposedFunction{typeof(first), typeof(EnzymeTestUtils.to_vec)}, Base.Splat{EnzymeTestUtils.var"#fnew#38"{Bool, EnzymeTestUtils.var"#call_with_captured_kwargs#56"{@NamedTuple{}}, Tuple{typeof(f_uninit!), MutatedStructWithUninitializedField}, Tuple{Bool, Bool}}}}, EnzymeTestUtils.var"#from_vec#3"{EnzymeTestUtils.var"#Tuple_from_vec#6"{Tuple{MutatedStructWithUninitializedField}, EnzymeTestUtils.var"#Array_from_vec#5"{Vector{MutatedStructWithUninitializedField}, Bool, Bool}}}}, ȳ::Vector{Float64}, x::Vector{Float64})
      @ FiniteDifferences ~/.julia/packages/FiniteDifferences/IPGFN/src/grad.jl:77
    [9] _fd_reverse(fdm::FiniteDifferences.AdaptedFiniteDifferenceMethod{5, 1, FiniteDifferences.UnadaptedFiniteDifferenceMethod{7, 5}}, f::Function, ȳ::MutatedStructWithUninitializedField, activities::Tuple{Const{typeof(f_uninit!)}, Duplicated{MutatedStructWithUninitializedField}}, active_return::Bool)
      @ EnzymeTestUtils ~/projects/Enzyme.jl/lib/EnzymeTestUtils/src/finite_difference_calls.jl:91
   [10] macro expansion
      @ ~/projects/Enzyme.jl/lib/EnzymeTestUtils/src/test_reverse.jl:104 [inlined]
   [11] macro expansion
      @ ~/.julia/juliaup/julia-1.11.1+0.x64.linux.gnu/share/julia/stdlib/v1.11/Test/src/Test.jl:1700 [inlined]
   [12] test_reverse(f::typeof(f_uninit!), ret_activity::Type, args::Tuple{MutatedStructWithUninitializedField, UnionAll}; rng::Random.TaskLocalRNG, fdm::FiniteDifferences.AdaptedFiniteDifferenceMethod{5, 1, FiniteDifferences.UnadaptedFiniteDifferenceMethod{7, 5}}, fkwargs::@NamedTuple{}, rtol::Float64, atol::Float64, testset_name::Nothing, runtime_activity::Bool)
      @ EnzymeTestUtils ~/projects/Enzyme.jl/lib/EnzymeTestUtils/src/test_reverse.jl:87
   [13] test_reverse(f::Function, ret_activity::Type, args::Tuple{MutatedStructWithUninitializedField, UnionAll})
      @ EnzymeTestUtils ~/projects/Enzyme.jl/lib/EnzymeTestUtils/src/test_reverse.jl:69
   [14] top-level scope
      @ REPL[19]:1
   [15] eval
      @ ./boot.jl:430 [inlined]
   [16] eval_user_input(ast::Any, backend::REPL.REPLBackend, mod::Module)
      @ REPL ~/.julia/juliaup/julia-1.11.1+0.x64.linux.gnu/share/julia/stdlib/v1.11/REPL/src/REPL.jl:245
   [17] repl_backend_loop(backend::REPL.REPLBackend, get_module::Function)
      @ REPL ~/.julia/juliaup/julia-1.11.1+0.x64.linux.gnu/share/julia/stdlib/v1.11/REPL/src/REPL.jl:342
   [18] start_repl_backend(backend::REPL.REPLBackend, consumer::Any; get_module::Function)
      @ REPL ~/.julia/juliaup/julia-1.11.1+0.x64.linux.gnu/share/julia/stdlib/v1.11/REPL/src/REPL.jl:327
   [19] run_repl(repl::REPL.AbstractREPL, consumer::Any; backend_on_current_task::Bool, backend::Any)
      @ REPL ~/.julia/juliaup/julia-1.11.1+0.x64.linux.gnu/share/julia/stdlib/v1.11/REPL/src/REPL.jl:483
   [20] run_repl(repl::REPL.AbstractREPL, consumer::Any)
      @ REPL ~/.julia/juliaup/julia-1.11.1+0.x64.linux.gnu/share/julia/stdlib/v1.11/REPL/src/REPL.jl:469
   [21] (::Base.var"#1139#1141"{Bool, Symbol, Bool})(REPL::Module)
      @ Base ./client.jl:446
   [22] #invokelatest#2
      @ ./essentials.jl:1055 [inlined]
   [23] invokelatest
      @ ./essentials.jl:1052 [inlined]
   [24] run_main_repl(interactive::Bool, quiet::Bool, banner::Symbol, history_file::Bool, color_set::Bool)
      @ Base ./client.jl:430
   [25] repl_main
      @ ./client.jl:567 [inlined]
   [26] _start()
      @ Base ./client.jl:541
Test Summary:                                                                                                  | Error  Total  Time
test_reverse: f_uninit! with return activity Duplicated on (::MutatedStructWithUninitializedField, Duplicated) |     1      1  1.7s
ERROR: Some tests did not pass: 0 passed, 0 failed, 1 errored, 0 broken.

Clearly, the vectorization code for FiniteDifferences makes an assumption that is being violated here. Might be an easy fix or not, but unfortunately I won't have the bandwidth to dig into this for some time.

@maximilian-gelbrecht in the meantime can you use the testers on a closure that uses a plan internally but does not take it as an argument? I think in general for an FFT rule you would want the plan to be a Const anyways.

[maximilian-gelbrecht]

@maximilian-gelbrecht in the meantime can you use the testers on a closure that uses a plan internally but does not take it as an argument? I think in general for an FFT rule you would want the plan to be a Const anyways.

Don't worry about me. This temporary/hacky solution that I posted above works for me. I just wanted to share that it's an issue that exists, and my personal solution isn't fit for a PR.