Merge pull request #60 from JuliaMath/an/precomputew

dannys4 · web-flow · commit ac171b512f5a · 2025-03-19T19:52:55.000Z
Precompute w to avoid recomputing it over and over
diff --git a/src/algos.jl b/src/algos.jl
@@ -4,6 +4,8 @@ fft!(::AbstractVector{T}, ::AbstractVector{T}, ::Int, ::Int, ::Direction, ::Abst
     Int(d)
 end
 
+@inline _conj(w::Complex, d::Direction) = ifelse(direction_sign(d) === 1, w, conj(w))
+
 function (g::CallGraph{T})(out::AbstractVector{T}, in::AbstractVector{U}, start_out::Int, start_in::Int, v::Direction, t::FFTEnum, idx::Int) where {T,U}
     fft!(out, in, start_out, start_in, v, t, g, idx)
 end
@@ -52,7 +54,7 @@ function fft!(out::AbstractVector{T}, in::AbstractVector{U}, start_out::Int, sta
     s_in = root.s_in
     s_out = root.s_out
 
-    w1 = convert(T, cispi(direction_sign(d)*2/N))
+    w1 = _conj(root.w, d)
     wj1 = one(T)
     tmp = g.workspace[idx]
     @inbounds for j1 in 0:N1-1
@@ -82,17 +84,17 @@ Discrete Fourier Transform, O(N^2) algorithm, in place.
 `stride_out`: Stride of the output vector
 `start_in`: Index of the first element of the input vector
 `stride_in`: Stride of the input vector
-`d`: Direction of the transform
+`w`: The value `cispi(direction_sign(d) * 2 / N)`
 
 """
-function fft_dft!(out::AbstractVector{T}, in::AbstractVector{T}, N::Int, start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, d::Direction) where {T}
+function fft_dft!(out::AbstractVector{T}, in::AbstractVector{T}, N::Int, start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, w::T) where {T}
     tmp = in[start_in]
     @inbounds for j in 1:N-1
         tmp += in[start_in + j*stride_in]
     end
     out[start_out] = tmp
 
-    wk = wkn = w = convert(T, cispi(direction_sign(d)*2/N))
+    wk = wkn = w
     @inbounds for d in 1:N-1
         tmp = in[start_in]
         @inbounds for k in 1:N-1
@@ -105,7 +107,7 @@ function fft_dft!(out::AbstractVector{T}, in::AbstractVector{T}, N::Int, start_o
     end
 end
 
-function fft_dft!(out::AbstractVector{Complex{T}}, in::AbstractVector{T}, N::Int, start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, d::Direction) where {T<:Real}
+function fft_dft!(out::AbstractVector{Complex{T}}, in::AbstractVector{T}, N::Int, start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, w::Complex{T}) where {T<:Real}
     halfN = N÷2
 
     tmp = Complex{T}(in[start_in])
@@ -114,7 +116,7 @@ function fft_dft!(out::AbstractVector{Complex{T}}, in::AbstractVector{T}, N::Int
     end
     out[start_out] = tmp
 
-    wk = wkn = w = convert(Complex{T}, cispi(direction_sign(d)*2/N))
+    wk = wkn = w
     @inbounds for d in 1:halfN
         tmp = Complex{T}(in[start_in])
         @inbounds for k in 1:N-1
@@ -129,7 +131,7 @@ end
 
 function fft!(out::AbstractVector{T}, in::AbstractVector{U}, start_out::Int, start_in::Int, d::Direction, ::DFT, g::CallGraph{T}, idx::Int) where {T,U}
     root = g[idx]
-    fft_dft!(out, in, root.sz, start_out, root.s_out, start_in, root.s_in, d)
+    fft_dft!(out, in, root.sz, start_out, root.s_out, start_in, root.s_in, _conj(root.w, d))
 end
 
 """
@@ -144,29 +146,28 @@ Power of 2 FFT, in place
 `stride_out`: Stride of the output vector
 `start_in`: Index of the first element of the input vector
 `stride_in`: Stride of the input vector
-`d`: Direction of the transform
+`w`: The value `cispi(direction_sign(d) * 2 / N)`
 
 """
-function fft_pow2!(out::AbstractVector{T}, in::AbstractVector{U}, N::Int, start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, d::Direction) where {T, U}
+function fft_pow2!(out::AbstractVector{T}, in::AbstractVector{U}, N::Int, start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, w::T) where {T, U}
     if N == 2
         out[start_out]              = in[start_in] + in[start_in + stride_in]
         out[start_out + stride_out] = in[start_in] - in[start_in + stride_in]
         return
     end
     m = N ÷ 2
 
-    fft_pow2!(out, in, m, start_out               , stride_out, start_in            , stride_in*2, d)
-    fft_pow2!(out, in, m, start_out + m*stride_out, stride_out, start_in + stride_in, stride_in*2, d)
+    fft_pow2!(out, in, m, start_out               , stride_out, start_in            , stride_in*2, w*w)
+    fft_pow2!(out, in, m, start_out + m*stride_out, stride_out, start_in + stride_in, stride_in*2, w*w)
 
-    w1 = convert(T, cispi(direction_sign(d)*2/N))
     wj = one(T)
     @inbounds for j in 0:m-1
         j1_out = start_out + j*stride_out
         j2_out = start_out + (j+m)*stride_out
         out_j    = out[j1_out]
         out[j1_out] = out_j + wj*out[j2_out]
         out[j2_out] = out_j - wj*out[j2_out]
-        wj *= w1
+        wj *= w
     end
 end
 
@@ -175,7 +176,7 @@ function fft!(out::AbstractVector{T}, in::AbstractVector{U}, start_out::Int, sta
     N = root.sz
     s_in = root.s_in
     s_out = root.s_out
-    fft_pow2!(out, in, N, start_out, s_out, start_in, s_in, d)
+    fft_pow2!(out, in, N, start_out, s_out, start_in, s_in, _conj(root.w, d))
 end
 
 """
@@ -190,13 +191,12 @@ Power of 4 FFT, in place
 `stride_out`: Stride of the output vector
 `start_in`: Index of the first element of the input vector
 `stride_in`: Stride of the input vector
-`d`: Direction of the transform
+`w`: The value `cispi(direction_sign(d) * 2 / N)`
 
 """
-function fft_pow4!(out::AbstractVector{T}, in::AbstractVector{U}, N::Int, start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, d::Direction) where {T, U}
-    ds = direction_sign(d)
-    plusi = ds*1im
-    minusi = ds*-1im
+function fft_pow4!(out::AbstractVector{T}, in::AbstractVector{U}, N::Int, start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, w::T) where {T, U}
+    plusi = sign(imag(w))*im
+    minusi = -sign(imag(w))*im
     if N == 4
         out[start_out + 0]            = in[start_in] + in[start_in + stride_in]        + in[start_in + 2*stride_in] + in[start_in + 3*stride_in]
         out[start_out +   stride_out] = in[start_in] + in[start_in + stride_in]*plusi  - in[start_in + 2*stride_in] + in[start_in + 3*stride_in]*minusi
@@ -206,17 +206,14 @@ function fft_pow4!(out::AbstractVector{T}, in::AbstractVector{U}, N::Int, start_
     end
     m = N ÷ 4
 
-    @muladd fft_pow4!(out, in, m, start_out                 , stride_out, start_in              , stride_in*4, d)
-    @muladd fft_pow4!(out, in, m, start_out +   m*stride_out, stride_out, start_in +   stride_in, stride_in*4, d)
-    @muladd fft_pow4!(out, in, m, start_out + 2*m*stride_out, stride_out, start_in + 2*stride_in, stride_in*4, d)
-    @muladd fft_pow4!(out, in, m, start_out + 3*m*stride_out, stride_out, start_in + 3*stride_in, stride_in*4, d)
+    @muladd fft_pow4!(out, in, m, start_out                 , stride_out, start_in              , stride_in*4, w^4)
+    @muladd fft_pow4!(out, in, m, start_out +   m*stride_out, stride_out, start_in +   stride_in, stride_in*4, w^4)
+    @muladd fft_pow4!(out, in, m, start_out + 2*m*stride_out, stride_out, start_in + 2*stride_in, stride_in*4, w^4)
+    @muladd fft_pow4!(out, in, m, start_out + 3*m*stride_out, stride_out, start_in + 3*stride_in, stride_in*4, w^4)
 
-    w1 = convert(T, cispi(direction_sign(d)*2/N))
-    wj = one(T)
-
-    w1 = convert(T, cispi(ds*2/N))
-    w2 = convert(T, cispi(ds*4/N))
-    w3 = convert(T, cispi(ds*6/N))
+    w1 = w
+    w2 = w*w1
+    w3 = w*w2
     wk1 = wk2 = wk3 = one(T)
 
     @inbounds for k in 0:m-1
@@ -240,7 +237,7 @@ function fft!(out::AbstractVector{T}, in::AbstractVector{U}, start_out::Int, sta
     N = root.sz
     s_in = root.s_in
     s_out = root.s_out
-    fft_pow4!(out, in, N, start_out, s_out, start_in, s_in, d)
+    fft_pow4!(out, in, N, start_out, s_out, start_in, s_in, _conj(root.w, d))
 end
 
 """
@@ -255,12 +252,12 @@ start_out: Index of the first element of the output vector
 stride_out: Stride of the output vector
 start_in: Index of the first element of the input vector
 stride_in: Stride of the input vector
-d: Direction of the transform
+w: The value `cispi(direction_sign(d) * 2 / N)`
 plus120: Depending on direction, perform either ±120° rotation
 minus120: Depending on direction, perform either ∓120° rotation
 
 """
-function fft_pow3!(out::AbstractVector{T}, in::AbstractVector{U}, N::Int, start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, d::Direction, plus120::T, minus120::T) where {T, U}
+function fft_pow3!(out::AbstractVector{T}, in::AbstractVector{U}, N::Int, start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, w::T, plus120::T, minus120::T) where {T, U}
     if N == 3
         @muladd out[start_out + 0]            = in[start_in] + in[start_in + stride_in]          + in[start_in + 2*stride_in]
         @muladd out[start_out +   stride_out] = in[start_in] + in[start_in + stride_in]*plus120  + in[start_in + 2*stride_in]*minus120
@@ -271,15 +268,13 @@ function fft_pow3!(out::AbstractVector{T}, in::AbstractVector{U}, N::Int, start_
     # Size of subproblem
     Nprime = N ÷ 3
 
-    ds = direction_sign(d)
-
     # Dividing into subproblems
-    fft_pow3!(out, in, Nprime, start_out, stride_out, start_in, stride_in*3, d, plus120, minus120)
-    fft_pow3!(out, in, Nprime, start_out + Nprime*stride_out, stride_out, start_in + stride_in, stride_in*3, d, plus120, minus120)
-    fft_pow3!(out, in, Nprime, start_out + 2*Nprime*stride_out, stride_out, start_in + 2*stride_in, stride_in*3, d, plus120, minus120)
+    fft_pow3!(out, in, Nprime, start_out, stride_out, start_in, stride_in*3, w^3, plus120, minus120)
+    fft_pow3!(out, in, Nprime, start_out + Nprime*stride_out, stride_out, start_in + stride_in, stride_in*3, w^3, plus120, minus120)
+    fft_pow3!(out, in, Nprime, start_out + 2*Nprime*stride_out, stride_out, start_in + 2*stride_in, stride_in*3, w^3, plus120, minus120)
 
-    w1 = convert(T, cispi(ds*2/N))
-    w2 = convert(T, cispi(ds*4/N))
+    w1 = w
+    w2 = w*w1
     wk1 = wk2 = one(T)
     for k in 0:Nprime-1
         @muladd k0 = start_out + k*stride_out
@@ -302,8 +297,8 @@ function fft!(out::AbstractVector{T}, in::AbstractVector{U}, start_out::Int, sta
     p_120 = convert(T, cispi(2/3))
     m_120 = convert(T, cispi(4/3))
     if d == FFT_FORWARD
-        fft_pow3!(out, in, N, start_out, s_out, start_in, s_in, d, m_120, p_120)
+        fft_pow3!(out, in, N, start_out, s_out, start_in, s_in, _conj(root.w, d), m_120, p_120)
     else
-        fft_pow3!(out, in, N, start_out, s_out, start_in, s_in, d, p_120, m_120)
+        fft_pow3!(out, in, N, start_out, s_out, start_in, s_in, _conj(root.w, d), p_120, m_120)
     end
 end
diff --git a/src/callgraph.jl b/src/callgraph.jl
@@ -32,13 +32,14 @@ Node of a call graph
 `sz`: Size of this FFT
 
 """
-struct CallGraphNode
+struct CallGraphNode{T}
     left::Int
     right::Int
     type::FFTEnum
     sz::Int
     s_in::Int
     s_out::Int
+    w::T
 end
 
 """
@@ -51,7 +52,7 @@ Object representing a graph of FFT Calls
 
 """
 struct CallGraph{T<:Complex}
-    nodes::Vector{CallGraphNode}
+    nodes::Vector{CallGraphNode{T}}
     workspace::Vector{Vector{T}}
 end
 
@@ -83,28 +84,29 @@ Recursively instantiate a set of `CallGraphNode`s
 `s_out`: The stride of the output
 
 """
-function CallGraphNode!(nodes::Vector{CallGraphNode}, N::Int, workspace::Vector{Vector{T}}, s_in::Int, s_out::Int)::Int where {T}
+function CallGraphNode!(nodes::Vector{CallGraphNode{T}}, N::Int, workspace::Vector{Vector{T}}, s_in::Int, s_out::Int)::Int where {T}
     if N == 0
         throw(DimensionMismatch("array has to be non-empty"))
     end
+    w = convert(T, cispi(2/N))
     if iseven(N)
         pow = _ispow24(N)
         if !isnothing(pow)
             push!(workspace, T[])
-            push!(nodes, CallGraphNode(0, 0, pow == POW2 ? pow2FFT : pow4FFT, N, s_in, s_out))
+            push!(nodes, CallGraphNode(0, 0, pow == POW2 ? pow2FFT : pow4FFT, N, s_in, s_out, w))
             return 1
         end
     end
     if N % 3 == 0
         if nextpow(3, N) == N
             push!(workspace, T[])
-            push!(nodes, CallGraphNode(0, 0, pow3FFT, N, s_in, s_out))
+            push!(nodes, CallGraphNode(0, 0, pow3FFT, N, s_in, s_out, w))
             return 1
         end
     end
     if N == 1 || isprime(N)
         push!(workspace, T[])
-        push!(nodes, CallGraphNode(0, 0, dft, N, s_in, s_out))
+        push!(nodes, CallGraphNode(0, 0, dft, N, s_in, s_out, w))
         return 1
     end
     Ns = [first(x) for x in collect(factor(N)) for _ in 1:last(x)]
@@ -121,12 +123,12 @@ function CallGraphNode!(nodes::Vector{CallGraphNode}, N::Int, workspace::Vector{
         N1 = N_cp[N1_idx]
     end
     N2 = N ÷ N1
-    push!(nodes, CallGraphNode(0, 0, dft, N, s_in, s_out))
+    push!(nodes, CallGraphNode(0, 0, dft, N, s_in, s_out, w))
     sz = length(nodes)
     push!(workspace, Vector{T}(undef, N))
     left_len = CallGraphNode!(nodes, N1, workspace, N2, N2*s_out)
     right_len = CallGraphNode!(nodes, N2, workspace, N1*s_in, 1)
-    nodes[sz] = CallGraphNode(1, 1 + left_len, compositeFFT, N, s_in, s_out)
+    nodes[sz] = CallGraphNode(1, 1 + left_len, compositeFFT, N, s_in, s_out, w)
     return 1 + left_len + right_len
 end
 
@@ -136,7 +138,7 @@ Instantiate a CallGraph from a number `N`
 
 """
 function CallGraph{T}(N::Int) where {T}
-    nodes = CallGraphNode[]
+    nodes = CallGraphNode{T}[]
     workspace = Vector{Vector{T}}()
     CallGraphNode!(nodes, N, workspace, 1, 1)
     CallGraph(nodes, workspace)
diff --git a/test/onedim/real_forward.jl b/test/onedim/real_forward.jl
@@ -17,7 +17,7 @@ end
 
     @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, FFTA.FFT_FORWARD)
+        FFTA.fft_dft!(y_dft, x, n, 1, 1, 1, 1, cispi(-2/n))
         @test y ≈ y_dft
     end
 end
