Skip to content

Commit 55a9703

Browse files
dannys4dannys4
authored
Merge pull request #62 from JuliaMath/an/d2
Introduce function barrier in higher dimensional transform loop to avoid a lot of allocations.
2 parents fbf7435 + bb8f6a7 commit 55a9703

File tree

12 files changed

+163
-100
lines changed

12 files changed

+163
-100
lines changed

‎.github/workflows/ci.yml‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,16 +17,17 @@ jobs:
1717
fail-fast: false
1818
matrix:
1919
version:
20-
- '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'.
21-
- '1' # Leave this line unchanged. '1' will automatically expand to the latest stable 1.x release of Julia.
22-
- 'nightly'
20+
- 'min'
21+
- 'lts'
22+
- '1'
23+
- 'pre'
2324
os:
2425
- ubuntu-latest
2526
arch:
2627
- x64
2728
steps:
2829
- uses: actions/checkout@v4
29-
- uses: julia-actions/setup-julia@v1
30+
- uses: julia-actions/setup-julia@v2
3031
with:
3132
version: ${{ matrix.version }}
3233
arch: ${{ matrix.arch }}

‎Project.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ DocStringExtensions = "0.9"
1818
MuladdMacro = "0.2"
1919
Primes = "0.5"
2020
Reexport = "1"
21-
julia = "1"
21+
julia = "1.6"
2222

2323
[extras]
2424
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

‎src/plan.jl‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,23 @@ function LinearAlgebra.mul!(y::AbstractArray{U,N}, p::FFTAPlan{T,2}, x::Abstract
112112
R3 = CartesianIndices(size(x)[p.region[2]+1:end])
113113
y_tmp = similar(y, axes(y)[p.region])
114114
rows,cols = size(x)[p.region]
115+
# Introduce function barrier here since the variables used in the loop ranges aren't inferred. This
116+
# is partly because the region field of the plan is abstractly typed but even if that wasn't the case,
117+
# it might be a bit tricky to construct the Rxs in an inferred way.
118+
_mul_loop(y_tmp, y, x, p, R1, R2, R3, rows, cols)
119+
end
120+
121+
function _mul_loop(
122+
y_tmp::AbstractArray,
123+
y::AbstractArray,
124+
x::AbstractArray,
125+
p::FFTAPlan,
126+
R1::CartesianIndices,
127+
R2::CartesianIndices,
128+
R3::CartesianIndices,
129+
rows::Int,
130+
cols::Int
131+
)
115132
for I1 in R1
116133
for I2 in R2
117134
for I3 in R3

‎test/onedim/complex_backward.jl‎

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,23 @@
11
using FFTA, Test
2-
test_nums = [8, 11, 15, 16, 27, 100]
3-
@testset "backward" begin
4-
for N in test_nums
5-
x = ones(ComplexF64, N)
6-
y = bfft(x)
7-
y_ref = 0*y
8-
y_ref[1] = N
9-
@test y ≈ y_ref atol=1e-12
10-
end
2+
3+
@testset "backward. N=$N" for N in [8, 11, 15, 16, 27, 100]
4+
x = ones(ComplexF64, N)
5+
y = bfft(x)
6+
y_ref = 0*y
7+
y_ref[1] = N
8+
@test y ≈ y_ref atol=1e-12
119
end
1210

13-
@testset verbose = true "against naive implementation. Size: $n" for n in 1:64
11+
@testset "More backward tests. Size: $n" for n in 1:64
1412
x = complex.(randn(n), randn(n))
15-
@test naive_1d_fourier_transform(x, FFTA.FFT_BACKWARD) ≈ bfft(x)
13+
14+
@testset "against naive implementation" begin
15+
@test naive_1d_fourier_transform(x, FFTA.FFT_BACKWARD) ≈ bfft(x)
16+
end
17+
18+
@testset "allocation regression" begin
19+
@test (@test_allocations bfft(x)) <= 44
20+
end
1621
end
1722

1823
@testset "error messages" begin

‎test/onedim/complex_forward.jl‎

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,23 @@
11
using FFTA, Test
2-
test_nums = [8, 11, 15, 16, 27, 100]
3-
@testset verbose = true " forward" begin
4-
for N in test_nums
5-
x = ones(ComplexF64, N)
6-
y = fft(x)
7-
y_ref = 0*y
8-
y_ref[1] = N
9-
@test y ≈ y_ref atol=1e-12
10-
end
2+
3+
@testset verbose = true " forward. N=$N" for N in [8, 11, 15, 16, 27, 100]
4+
x = ones(ComplexF64, N)
5+
y = fft(x)
6+
y_ref = 0*y
7+
y_ref[1] = N
8+
@test y ≈ y_ref atol=1e-12
119
end
1210

13-
@testset verbose = true "against naive implementation. Size: $n" for n in 1:64
11+
@testset "More forward tests. Size: $n" for n in 1:64
1412
x = complex.(randn(n), randn(n))
15-
@test naive_1d_fourier_transform(x, FFTA.FFT_FORWARD) ≈ fft(x)
13+
14+
@testset "against naive implementation" begin
15+
@test naive_1d_fourier_transform(x, FFTA.FFT_FORWARD) ≈ fft(x)
16+
end
17+
18+
@testset "allocation regression" begin
19+
@test (@test_allocations fft(x)) <= 44
20+
end
1621
end
1722

1823
@testset "error messages" begin

‎test/onedim/real_backward.jl‎

Lines changed: 18 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,17 @@
11
using FFTA, Test, LinearAlgebra
2-
test_nums = [8, 11, 15, 16, 27, 100]
3-
@testset "backward" begin
4-
for N in test_nums
5-
x = ones(Float64, N)
6-
y = brfft(x, 2*(N-1))
7-
y_ref = 0*y
8-
y_ref[1] = 2*(N-1)
9-
if !isapprox(y_ref, y, atol=1e-12)
10-
println(norm(y_ref - y))
11-
end
12-
@test y_ref ≈ y atol=1e-12
2+
3+
@testset "backward. N=$N" for N in [8, 11, 15, 16, 27, 100]
4+
x = ones(Float64, N)
5+
y = brfft(x, 2*(N-1))
6+
y_ref = 0*y
7+
y_ref[1] = 2*(N-1)
8+
if !isapprox(y_ref, y, atol=1e-12)
9+
println(norm(y_ref - y))
1310
end
11+
@test y_ref ≈ y atol=1e-12
1412
end
1513

16-
@testset verbose = true "against naive implementation. Size: $n" for n in 1:64
14+
@testset "More backward tests. Size: $n" for n in 1:64
1715
x = complex.(randn(n ÷ 2 + 1), randn(n ÷ 2 + 1))
1816
x[begin] = real(x[begin])
1917
if iseven(n)
@@ -22,7 +20,14 @@ end
2220
else
2321
xe = [x; conj.(reverse(x[begin + 1:end]))]
2422
end
25-
@test naive_1d_fourier_transform(xe, FFTA.FFT_BACKWARD) ≈ brfft(x, n)
23+
24+
@testset "against naive implementation" begin
25+
@test naive_1d_fourier_transform(xe, FFTA.FFT_BACKWARD) ≈ brfft(x, n)
26+
end
27+
28+
@testset "allocation regression" begin
29+
@test (@test_allocations brfft(x, n)) <= 50
30+
end
2631
end
2732

2833
@testset "error messages" begin

‎test/onedim/real_forward.jl‎

Lines changed: 21 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,29 @@
11
using FFTA, Test
2-
test_nums = [8, 11, 15, 16, 27, 100]
3-
@testset verbose = true " forward" begin
4-
for N in test_nums
5-
x = ones(Float64, N)
6-
y = rfft(x)
7-
y_ref = 0*y
8-
y_ref[1] = N
9-
@test y ≈ y_ref atol=1e-12
10-
end
2+
3+
@testset verbose = true " forward. N=$N" for N in [8, 11, 15, 16, 27, 100]
4+
x = ones(Float64, N)
5+
y = rfft(x)
6+
y_ref = 0*y
7+
y_ref[1] = N
8+
@test y ≈ y_ref atol=1e-12
119
end
1210

13-
@testset verbose = true "against naive implementation. Size: $n" for n in 1:64
11+
@testset "More forward tests. Size: $n" for n in 1:64
1412
x = randn(n)
15-
y = rfft(x)
16-
@test naive_1d_fourier_transform(x, FFTA.FFT_FORWARD)[1:(n ÷ 2 + 1)] ≈ y
1713

18-
@testset "temporarily test real dft separately until used by rfft" begin
19-
y_dft = similar(y)
20-
FFTA.fft_dft!(y_dft, x, n, 1, 1, 1, 1, cispi(-2/n))
21-
@test y ≈ y_dft
14+
@testset "against naive implementation" begin
15+
y = rfft(x)
16+
@test naive_1d_fourier_transform(x, FFTA.FFT_FORWARD)[1:(n ÷ 2 + 1)] ≈ y
17+
18+
@testset "temporarily test real dft separately until used by rfft" begin
19+
y_dft = similar(y)
20+
FFTA.fft_dft!(y_dft, x, n, 1, 1, 1, 1, cispi(-2/n))
21+
@test y ≈ y_dft
22+
end
23+
end
24+
25+
@testset "allocation regression" begin
26+
@test (@test_allocations rfft(x)) <= 48
2227
end
2328
end
2429

‎test/runtests.jl‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,13 @@
11
using Test, Random, FFTA
22

3+
macro test_allocations(args)
4+
if Base.VERSION >= v"1.9"
5+
:(@allocations($(esc(args))))
6+
else
7+
:(0)
8+
end
9+
end
10+
311
function naive_1d_fourier_transform(x::Vector, d::FFTA.Direction)
412
n = length(x)
513
y = zeros(Complex{Float64}, n)

‎test/twodim/complex_backward.jl‎

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,24 @@
11
using FFTA, Test
2-
test_nums = [8, 11, 15, 16, 27, 100]
3-
@testset "backward" begin
4-
for N in test_nums
5-
x = ones(ComplexF64, N, N)
6-
y = bfft(x)
7-
y_ref = 0*y
8-
y_ref[1] = length(x)
9-
@test y ≈ y_ref
10-
end
2+
3+
@testset "backward. N=$N" for N in [8, 11, 15, 16, 27, 100]
4+
x = ones(ComplexF64, N, N)
5+
y = bfft(x)
6+
y_ref = 0*y
7+
y_ref[1] = length(x)
8+
@test y ≈ y_ref
119
end
1210

13-
@testset verbose = true "against naive implementation" for n in 1:64
11+
@testset "More backward tests" for n in 1:64
1412
@testset "size: ($m, $n)" for m in n:(n + 1)
1513
X = complex.(randn(m, n), randn(m, n))
16-
Y = similar(X)
17-
@test naive_2d_fourier_transform(X, FFTA.FFT_BACKWARD) ≈ bfft(X)
14+
15+
@testset "against naive implementation" begin
16+
@test naive_2d_fourier_transform(X, FFTA.FFT_BACKWARD) ≈ bfft(X)
17+
end
18+
19+
@testset "allocations" begin
20+
@test (@test_allocations bfft(X)) <= 111
21+
end
1822
end
1923
end
2024

‎test/twodim/complex_forward.jl‎

Lines changed: 16 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,24 @@
11
using FFTA, Test
2-
test_nums = [8, 11, 15, 16, 27, 100]
3-
@testset " forward" begin
4-
for N in test_nums
5-
x = ones(ComplexF64, N, N)
6-
y = fft(x)
7-
y_ref = 0*y
8-
y_ref[1] = length(x)
9-
@test y ≈ y_ref
10-
end
2+
3+
@testset " forward. N=$N" for N in [8, 11, 15, 16, 27, 100]
4+
x = ones(ComplexF64, N, N)
5+
y = fft(x)
6+
y_ref = 0*y
7+
y_ref[1] = length(x)
8+
@test y ≈ y_ref
119
end
1210

13-
@testset verbose = true "against naive implementation" for n in 1:64
11+
@testset "More forward tests" for n in 1:64
1412
@testset "size: ($m, $n)" for m in n:(n + 1)
1513
X = complex.(randn(m, n), randn(m, n))
16-
Y = similar(X)
17-
@test naive_2d_fourier_transform(X, FFTA.FFT_FORWARD) ≈ fft(X)
14+
15+
@testset "against naive implementation" begin
16+
@test naive_2d_fourier_transform(X, FFTA.FFT_FORWARD) ≈ fft(X)
17+
end
18+
19+
@testset "allocations" begin
20+
@test (@test_allocations fft(X)) <= 111
21+
end
1822
end
1923
end
2024

0 commit comments

Comments
 (0)