Increasing robustness of adjoint plan optimizations (#123)

* Add some safeguards to prioritize type stability in adjoint plans

* Test adjoint plans with float32's

* Add subtype

* Apply suggestions from code review

* Rename dummy -> wrapper

Author: gaurav-arya

src/definitions.jl

@@ -681,13 +681,8 @@ function adjoint_mul(p::Plan{T}, x::AbstractArray, ::FFTAdjointStyle) where {T}
     dims = fftdims(p)
     N = normalization(T, size(p), dims)
     pinv = inv(p)
-    # Optimization: when pinv is a ScaledPlan, check if we can avoid a loop over x.
-    # Even if not, ensure that we do only one pass by combining the normalization with the plan.
-    if pinv isa ScaledPlan && pinv.scale == N
-        return pinv.p * x
-    else
-        return (1/N * pinv) * x
-    end
+    # Ensure that we do only one pass over the array by combining the normalization with the plan.
+    return (inv(N) * pinv) * x
 end
 
 function adjoint_mul(p::Plan{T}, x::AbstractArray, ::RFFTAdjointStyle) where {T<:Real}
@@ -698,7 +693,7 @@ function adjoint_mul(p::Plan{T}, x::AbstractArray, ::RFFTAdjointStyle) where {T<
     pinv = inv(p)
     n = size(pinv, halfdim)
     # Optimization: when pinv is a ScaledPlan, fuse the scaling into our map to ensure we do not loop over x twice.
-    scale = pinv isa ScaledPlan ? pinv.scale / 2N : 1 / 2N
+    scale = pinv isa ScaledPlan ? pinv.scale / 2N : inv(2N)
     twoscale = 2 * scale 
     unscaled_pinv = pinv isa ScaledPlan ? pinv.p : pinv 
     y = map(x, CartesianIndices(x)) do xj, j
@@ -721,7 +716,7 @@ function adjoint_mul(p::Plan{T}, x::AbstractArray, ::IRFFTAdjointStyle) where {T
     pinv = inv(p)
     d = size(pinv, halfdim)
     # Optimization: when pinv is a ScaledPlan, fuse the scaling into our map to ensure we do not loop over x twice.
-    scale = pinv isa ScaledPlan ? pinv.scale / N : 1 / N
+    scale = pinv isa ScaledPlan ? pinv.scale / N : inv(N)
     twoscale = 2 * scale 
     unscaled_pinv = pinv isa ScaledPlan ? pinv.p : pinv 
     y = unscaled_pinv * x

test/TestPlans.jl

@@ -278,4 +278,20 @@ AbstractFFTs.plan_inv(p::InplaceTestPlan) = InplaceTestPlan(AbstractFFTs.plan_in
 # Don't cache inverse of inplace wrapper plan (only inverse of inner plan)
 Base.inv(p::InplaceTestPlan) = InplaceTestPlan(inv(p.plan))
 
+# A wrapper plan whose inverse is not an instance of AbstractFFTs.ScaledPlan, for testing purposes
+
+struct WrapperTestPlan{T,P<:Plan{T}} <: Plan{T}
+    plan::P
+end
+
+Base.size(p::WrapperTestPlan) = size(p.plan)
+Base.ndims(p::WrapperTestPlan) = ndims(p.plan)
+AbstractFFTs.fftdims(p::WrapperTestPlan) = fftdims(p.plan)
+AbstractFFTs.AdjointStyle(p::WrapperTestPlan) = AbstractFFTs.AdjointStyle(p.plan)
+
+Base.:*(p::WrapperTestPlan, x::AbstractArray) = p.plan * x
+
+AbstractFFTs.plan_inv(p::WrapperTestPlan) = WrapperTestPlan(AbstractFFTs.plan_inv(p.plan))
+Base.inv(p::WrapperTestPlan) = WrapperTestPlan(inv(p.plan))
+
 end

test/runtests.jl

@@ -145,6 +145,39 @@ end
     end
 end
 
+@testset "Adjoint plan on single-precision" begin
+    # fft
+    p = plan_fft(zeros(ComplexF32, 3))
+    u = rand(ComplexF32, 3)
+    @test eltype(p' * (p * u)) == eltype(u)
+    # rfft
+    p = plan_rfft(zeros(Float32, 3))
+    u = rand(Float32, 3)
+    @test eltype(p' * (p * u)) == eltype(u)
+    # brfft
+    p = plan_brfft(zeros(ComplexF32, 3), 5)
+    u = rand(ComplexF32, 3)
+    @test eltype(p' * (p * u)) == eltype(u)
+end
+
+@testset "Adjoint plan application when plan inverse is not a ScaledPlan" begin
+    # fft
+    p0 = plan_fft(zeros(ComplexF64, 3))
+    p = TestPlans.WrapperTestPlan(p0)
+    u = rand(ComplexF64, 3)
+    @test p' * u ≈ p0' * u 
+    # rfft
+    p0 = plan_rfft(zeros(3))
+    p = TestPlans.WrapperTestPlan(p0)
+    u = rand(ComplexF64, 2)
+    @test p' * u ≈ p0' * u 
+    # brfft
+    p0 = plan_brfft(zeros(ComplexF64, 3), 5)
+    p = TestPlans.WrapperTestPlan(p0)
+    u = rand(Float64, 5)
+    @test p' * u ≈ p0' * u 
+end
+
 @testset "ChainRules" begin
     @testset "shift functions" begin
         for x in (randn(3), randn(3, 4), randn(3, 4, 5))