JuliaMath / AbstractFFTs.jl

A Julia framework for implementing FFTs
MIT License
126 stars 32 forks source link

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

Closed gaurav-arya closed 1 year ago

gaurav-arya commented 1 year ago

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 https://github.com/JuliaMath/AbstractFFTs.jl/pull/113#issuecomment-1688865004.

codecov[bot] commented 1 year ago

Codecov Report

Patch coverage: 96.15% and project coverage change: -0.17% :warning:

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

Additional details and impacted files ```diff @@ 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](https://app.codecov.io/gh/JuliaMath/AbstractFFTs.jl/pull/113?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaMath) | Coverage Δ | | |---|---|---| | [src/definitions.jl](https://app.codecov.io/gh/JuliaMath/AbstractFFTs.jl/pull/113?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaMath#diff-c3JjL2RlZmluaXRpb25zLmps) | `83.75% <95.65%> (-0.33%)` | :arrow_down: | | [ext/AbstractFFTsTestExt.jl](https://app.codecov.io/gh/JuliaMath/AbstractFFTs.jl/pull/113?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaMath#diff-ZXh0L0Fic3RyYWN0RkZUc1Rlc3RFeHQuamw=) | `100.00% <100.00%> (ø)` | |

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

stevengj commented 1 year ago

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 commented 1 year ago

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!: https://github.com/JuliaMath/AbstractFFTs.jl/blob/f9b5de31ff2d5ab9837fd302ca2f484d7db865fa/src/definitions.jl#L268. 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 commented 1 year ago

Bump on merge