Introduce function barrier in higher dimensional transform loop to avoid a lot of allocations.

.github/workflows/ci.yml

@@ -17,16 +17,17 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.6' # Replace this with the minimum Julia version that your package supports. E.g. if your package requires Julia 1.5 or higher, change this to '1.5'.
-          - '1' # Leave this line unchanged. '1' will automatically expand to the latest stable 1.x release of Julia.
-          - 'nightly'
+          - 'min'
+          - 'lts'
+          - '1'
+          - 'pre'
         os:
           - ubuntu-latest
         arch:
           - x64
     steps:
       - uses: actions/checkout@v4
-      - uses: julia-actions/setup-julia@v1
+      - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}

Project.toml

@@ -18,7 +18,7 @@ DocStringExtensions = "0.9"
 MuladdMacro = "0.2"
 Primes = "0.5"
 Reexport = "1"
-julia = "1"
+julia = "1.6"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

src/plan.jl

@@ -112,6 +112,23 @@ function LinearAlgebra.mul!(y::AbstractArray{U,N}, p::FFTAPlan{T,2}, x::Abstract
     R3 = CartesianIndices(size(x)[p.region[2]+1:end])
     y_tmp = similar(y, axes(y)[p.region])
     rows,cols = size(x)[p.region]
+    # Introduce function barrier here since the variables used in the loop ranges aren't inferred. This
+    # is partly because the region field of the plan is abstractly typed but even if that wasn't the case,
+    # it might be a bit tricky to construct the Rxs in an inferred way.
+    _mul_loop(y_tmp, y, x, p, R1, R2, R3, rows, cols)
+end
+
+function _mul_loop(
+    y_tmp::AbstractArray,
+    y::AbstractArray,
+    x::AbstractArray,
+    p::FFTAPlan,
+    R1::CartesianIndices,
+    R2::CartesianIndices,
+    R3::CartesianIndices,
+    rows::Int,
+    cols::Int
+)
     for I1 in R1
         for I2 in R2
             for I3 in R3

test/onedim/complex_backward.jl

@@ -1,18 +1,23 @@
 using FFTA, Test
-test_nums = [8, 11, 15, 16, 27, 100]
-@testset "backward" begin
-    for N in test_nums
-        x = ones(ComplexF64, N)
-        y = bfft(x)
-        y_ref = 0*y
-        y_ref[1] = N
-        @test y ≈ y_ref atol=1e-12
-    end
+
+@testset "backward. N=$N" for N in [8, 11, 15, 16, 27, 100]
+    x = ones(ComplexF64, N)
+    y = bfft(x)
+    y_ref = 0*y
+    y_ref[1] = N
+    @test y ≈ y_ref atol=1e-12
 end
 
-@testset verbose = true "against naive implementation. Size: $n" for n in 1:64
+@testset "More backward tests. Size: $n" for n in 1:64
     x = complex.(randn(n), randn(n))
-    @test naive_1d_fourier_transform(x, FFTA.FFT_BACKWARD) ≈ bfft(x)
+
+    @testset "against naive implementation" begin
+        @test naive_1d_fourier_transform(x, FFTA.FFT_BACKWARD) ≈ bfft(x)
+    end
+
+    @testset "allocation regression" begin
+        @test (@test_allocations bfft(x)) <= 44
+    end
 end
 
 @testset "error messages" begin

test/onedim/complex_forward.jl

@@ -1,18 +1,23 @@
 using FFTA, Test
-test_nums = [8, 11, 15, 16, 27, 100]
-@testset verbose = true " forward" begin
-    for N in test_nums
-        x = ones(ComplexF64, N)
-        y = fft(x)
-        y_ref = 0*y
-        y_ref[1] = N
-        @test y ≈ y_ref atol=1e-12
-    end
+
+@testset verbose = true " forward. N=$N" for N in [8, 11, 15, 16, 27, 100]
+    x = ones(ComplexF64, N)
+    y = fft(x)
+    y_ref = 0*y
+    y_ref[1] = N
+    @test y ≈ y_ref atol=1e-12
 end
 
-@testset verbose = true "against naive implementation. Size: $n" for n in 1:64
+@testset "More forward tests. Size: $n" for n in 1:64
     x = complex.(randn(n), randn(n))
-    @test naive_1d_fourier_transform(x, FFTA.FFT_FORWARD) ≈ fft(x)
+
+    @testset "against naive implementation" begin
+        @test naive_1d_fourier_transform(x, FFTA.FFT_FORWARD) ≈ fft(x)
+    end
+
+    @testset "allocation regression" begin
+        @test (@test_allocations fft(x)) <= 44
+    end
 end
 
 @testset "error messages" begin

test/onedim/real_backward.jl

@@ -1,19 +1,17 @@
 using FFTA, Test, LinearAlgebra
-test_nums = [8, 11, 15, 16, 27, 100]
-@testset "backward" begin
-    for N in test_nums
-        x = ones(Float64, N)
-        y = brfft(x, 2*(N-1))
-        y_ref = 0*y
-        y_ref[1] = 2*(N-1)
-        if !isapprox(y_ref, y, atol=1e-12)
-            println(norm(y_ref - y))
-        end
-        @test y_ref ≈ y atol=1e-12
+
+@testset "backward. N=$N" for N in [8, 11, 15, 16, 27, 100]
+    x = ones(Float64, N)
+    y = brfft(x, 2*(N-1))
+    y_ref = 0*y
+    y_ref[1] = 2*(N-1)
+    if !isapprox(y_ref, y, atol=1e-12)
+        println(norm(y_ref - y))
     end
+    @test y_ref ≈ y atol=1e-12
 end
 
-@testset verbose = true "against naive implementation. Size: $n" for n in 1:64
+@testset "More backward tests. Size: $n" for n in 1:64
     x = complex.(randn(n ÷ 2 + 1), randn(n ÷ 2 + 1))
     x[begin] = real(x[begin])
     if iseven(n)
@@ -22,7 +20,14 @@ end
     else
         xe = [x; conj.(reverse(x[begin + 1:end]))]
     end
-    @test naive_1d_fourier_transform(xe, FFTA.FFT_BACKWARD) ≈ brfft(x, n)
+
+    @testset "against naive implementation" begin
+        @test naive_1d_fourier_transform(xe, FFTA.FFT_BACKWARD) ≈ brfft(x, n)
+    end
+
+    @testset "allocation regression" begin
+        @test (@test_allocations brfft(x, n)) <= 50
+    end
 end
 
 @testset "error messages" begin

test/onedim/real_forward.jl

@@ -1,24 +1,29 @@
 using FFTA, Test
-test_nums = [8, 11, 15, 16, 27, 100]
-@testset verbose = true " forward" begin
-    for N in test_nums
-        x = ones(Float64, N)
-        y = rfft(x)
-        y_ref = 0*y
-        y_ref[1] = N
-        @test y ≈ y_ref atol=1e-12
-    end
+
+@testset verbose = true " forward. N=$N" for N in [8, 11, 15, 16, 27, 100]
+    x = ones(Float64, N)
+    y = rfft(x)
+    y_ref = 0*y
+    y_ref[1] = N
+    @test y ≈ y_ref atol=1e-12
 end
 
-@testset verbose = true "against naive implementation. Size: $n" for n in 1:64
+@testset "More forward tests. Size: $n" for n in 1:64
     x = randn(n)
-    y = rfft(x)
-    @test naive_1d_fourier_transform(x, FFTA.FFT_FORWARD)[1:(n ÷ 2 + 1)] ≈ y
 
-    @testset "temporarily test real dft separately until used by rfft" begin
-        y_dft = similar(y)
-        FFTA.fft_dft!(y_dft, x, n, 1, 1, 1, 1, cispi(-2/n))
-        @test y ≈ y_dft
+    @testset "against naive implementation" begin
+        y = rfft(x)
+        @test naive_1d_fourier_transform(x, FFTA.FFT_FORWARD)[1:(n ÷ 2 + 1)] ≈ y
+
+        @testset "temporarily test real dft separately until used by rfft" begin
+            y_dft = similar(y)
+            FFTA.fft_dft!(y_dft, x, n, 1, 1, 1, 1, cispi(-2/n))
+            @test y ≈ y_dft
+        end
+    end
+
+    @testset "allocation regression" begin
+        @test (@test_allocations rfft(x)) <= 48
     end
 end
 

test/runtests.jl

@@ -1,5 +1,13 @@
 using Test, Random, FFTA
 
+macro test_allocations(args)
+    if Base.VERSION >= v"1.9"
+        :(@allocations($(esc(args))))
+    else
+        :(0)
+    end
+end
+
 function naive_1d_fourier_transform(x::Vector, d::FFTA.Direction)
     n = length(x)
     y = zeros(Complex{Float64}, n)

test/twodim/complex_backward.jl

@@ -1,20 +1,24 @@
 using FFTA, Test
-test_nums = [8, 11, 15, 16, 27, 100]
-@testset "backward" begin
-    for N in test_nums
-        x = ones(ComplexF64, N, N)
-        y = bfft(x)
-        y_ref = 0*y
-        y_ref[1] = length(x)
-        @test y ≈ y_ref
-    end
+
+@testset "backward. N=$N" for N in [8, 11, 15, 16, 27, 100]
+    x = ones(ComplexF64, N, N)
+    y = bfft(x)
+    y_ref = 0*y
+    y_ref[1] = length(x)
+    @test y ≈ y_ref
 end
 
-@testset verbose = true "against naive implementation" for n in 1:64
+@testset "More backward tests" for n in 1:64
     @testset "size: ($m, $n)" for m in n:(n + 1)
         X = complex.(randn(m, n), randn(m, n))
-        Y = similar(X)
-        @test naive_2d_fourier_transform(X, FFTA.FFT_BACKWARD) ≈ bfft(X)
+
+        @testset "against naive implementation" begin
+            @test naive_2d_fourier_transform(X, FFTA.FFT_BACKWARD) ≈ bfft(X)
+        end
+
+        @testset "allocations" begin
+            @test (@test_allocations bfft(X)) <= 111
+        end
     end
 end
 

test/twodim/complex_forward.jl

@@ -1,20 +1,24 @@
 using FFTA, Test
-test_nums = [8, 11, 15, 16, 27, 100]
-@testset " forward" begin
-    for N in test_nums
-        x = ones(ComplexF64, N, N)
-        y = fft(x)
-        y_ref = 0*y
-        y_ref[1] = length(x)
-        @test y ≈ y_ref
-    end
+
+@testset " forward. N=$N" for N in [8, 11, 15, 16, 27, 100]
+    x = ones(ComplexF64, N, N)
+    y = fft(x)
+    y_ref = 0*y
+    y_ref[1] = length(x)
+    @test y ≈ y_ref
 end
 
-@testset verbose = true "against naive implementation" for n in 1:64
+@testset "More forward tests" for n in 1:64
     @testset "size: ($m, $n)" for m in n:(n + 1)
         X = complex.(randn(m, n), randn(m, n))
-        Y = similar(X)
-        @test naive_2d_fourier_transform(X, FFTA.FFT_FORWARD) ≈ fft(X)
+
+        @testset "against naive implementation" begin
+            @test naive_2d_fourier_transform(X, FFTA.FFT_FORWARD) ≈ fft(X)
+        end
+
+        @testset "allocations" begin
+            @test (@test_allocations fft(X)) <= 111
+        end
     end
 end
 

test/twodim/real_backward.jl

@@ -1,11 +1,15 @@
 using FFTA, Test
-test_nums = [8]
-@testset "backward" begin
-    for N in test_nums
-        x = ones(Float64, N, N)
-        y = brfft(x, 2(N-1))
-        y_ref = 0*y
-        y_ref[1] = N*(2(N-1))
-        @test y_ref ≈ y atol=1e-12
-    end
-end
+
+@testset "backward. N=$N" for N in [8]
+    x = ones(Float64, N, N)
+    y = brfft(x, 2(N-1))
+    y_ref = 0*y
+    y_ref[1] = N*(2(N-1))
+    @test y_ref ≈ y atol=1e-12
+end
+
+@testset "allocations" begin
+    X = randn(256, 256)
+    rfft(X) # compile
+    @test (@test_allocations rfft(X)) <= 51
+end

test/twodim/real_forward.jl

@@ -1,11 +1,16 @@
 using FFTA, Test
-test_nums = [8, 11, 15, 16, 27, 100]
-@testset " forward" begin 
-    for N in test_nums
-        x = ones(Float64, N, N)
-        y = rfft(x)
-        y_ref = 0*y
-        y_ref[1] = length(x)
-        @test y ≈ y_ref
-    end
-end
+
+@testset " forward. N=$N" for N in [8, 11, 15, 16, 27, 100]
+    x = ones(Float64, N, N)
+    y = rfft(x)
+    y_ref = 0*y
+    y_ref[1] = length(x)
+    @test y ≈ y_ref
+end
+
+@testset "allocations" begin
+    X = randn(256, 256)
+    Y = rfft(X)
+    brfft(Y, 256) # compile
+    @test (@test_allocations brfft(Y, 256)) <= 54
+end