Eagerly invert plan on formation of AdjointPlan: correct eltype and remove output_size

[gaurav-arya]

This PR resolves #110. After some thinking I realized that the right way to do this should be to actually call inv(p) in the formation of the adjoint plan. This makes adjoint(p) work the same way as inv(p), in that caching occurs in the plan formation rather than plan application.

As a consequence, I realized that output_size should not actually be job of an adjoint style. TPlan implementers are actually already required to manually figure out output size for their plan when implementing size(plan_inv(p)), and hence it didn't actually make sense to have the adjoint style trait provide this. If we want to make to make the output size available for free without forming the inverse plan, that's a separate line of work (that could also simplify the lives of the original plan implementation) But it shouldn't actually be the job of adjoint traits -- the goal of adjoint traits is to make adjoints easy given that all inverse functionality is implemented, and indeed inverses already give the output size for free.

Note that this PR is technically breaking for implementers of new adjoint styles (by removing 1 required function). but does not affect users of existing adjoint styles.

Edit: this PR also includes the fusion optimization suggested in #113 (comment).

[codecov]

Codecov Report

Patch coverage: 96.15% and project coverage change: -0.17% ⚠️

Comparison is base (fae1170) 93.87% compared to head (d194cbe) 93.70%.

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #113      +/-   ##
==========================================
- Coverage   93.87%   93.70%   -0.17%     
==========================================
  Files           5        5              
  Lines         441      445       +4     
==========================================
+ Hits          414      417       +3     
- Misses         27       28       +1     

Files Changed	Coverage Δ
src/definitions.jl	83.75% <95.65%> (-0.33%)	⬇️
ext/AbstractFFTsTestExt.jl	100.00% <100.00%> (ø)

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[stevengj]

In the common case where inv(p) returns a ScaledPlan, it would be better to pull out the scale factor and combine it with the scaling loop in adjoint_mul, or eliminate the scale factor entirely if it cancels (as it will for complex FFTs).

For example:

function adjoint_mul(p::Plan{T}, x::AbstractArray, ::FFTAdjointStyle) where {T}
    dims = fftdims(p)
    N = normalization(T, size(p), dims)
    pinv = inv(p)
    if pinv isa ScaledPlan
        return pinv.scale == N ? pinv.p * x : rmul!(pinv.p * x, pinv.scale / N)
    else
        return rmul!(pinv * x, inv(N))
    end
end

[gaurav-arya]

Implemented in f9b5de3. In the FFT case, a neat way to do the case where the scaling does not cancel was to combine the extra scaling into the scaled plan, which will then do the same rmul!:

AbstractFFTs.jl/src/definitions.jl

Line 268 in f9b5de3

*(p::ScaledPlan, x::AbstractArray) = LinearAlgebra.rmul!(p.p * x, p.scale)

. In the RFFT/IRFFT case, it seems we always need at least one pass through the array for the map, so we take the scaling out of the scaled plan and into the map.

[gaurav-arya]

Bump on merge

[devmotion]

@gaurav-arya This line introduced a Float64 promotion that was not (necessarily) present in the previous implementation.

[gaurav-arya]

I don't quite follow, could you explain or give an example?

[devmotion]

1 / N will be a Float64 which will promote everything. In the old code this did not happen.

[gaurav-arya]

Apologies that I still don't follow: why does where we divide by N affect the promotion behavior?

[gaurav-arya]

And why is N a Float64 in the first place, rather than a quantity of type T?

[gaurav-arya]

julia> using FFTW

julia> P = plan_fft(rand(Float32, 3, 3));

julia> P' * rand(Float32, 3, 3)
3×3 Matrix{ComplexF32}:
    3.48416+0.0im       -0.264397-0.859624im  -0.264397+0.859624im
 -0.0523138-0.781731im  -0.462729-0.257876im  -0.902054+0.397495im
 -0.0523138+0.781731im  -0.902054-0.397495im  -0.462729+0.257876im

Could you provide an example of the issue?

[gaurav-arya]

Bump

[gaurav-arya]

Bump -- I still don't follow your initial comment, since as Int64 and Float32 promote to Float32. An example would be much appreciated.

[gaurav-arya]

Bump. If this a problem that exists on the main branch, could you make an issue for it, with a reproducer?

[devmotion]

Sorry, somehow this conversation was lost in my stream of Github notifications.

And why is N a Float64 in the first place, rather than a quantity of type T?

I assumed it would be an Int (and hence 1/N would be a Float64) but you're right of course that it is of type T. Maybe one could even use T = eltype(pinv) here since the computations are performed with pinv but I think it should not matter. Possibly inv(N) would be a tiny bit safer since then no Ints would be involved but I can't think of any relevant case where it should matter that the inverse is computed in a way that involves an Int.

[devmotion]

Same here, this introduces Float64 promotions.

[gaurav-arya]

julia> using FFTW

julia> P = plan_rfft(rand(Float32, 3, 3));

julia> P' * (P * rand(Float32, 3, 3))
3×3 Matrix{Float32}:
 1.35893  3.17063  3.99985
 1.15958  6.03116  3.67565
 3.10326  5.80633  3.89977