Merge pull request #29 from dannys4/ds/pow3

Creating radix 3 FFT

Author: dannys4

src/algos.jl

@@ -178,3 +178,51 @@ function fft!(out::AbstractVector{T}, in::AbstractVector{U}, start_out::Int, sta
     s_out = root.s_out
     fft_pow4!(out, in, N, start_out, s_out, start_in, s_in, d)
 end
+
+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}
+    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
+        @muladd out[start_out + 2*stride_out] = in[start_in] + in[start_in + stride_in]*minus120 + in[start_in + 2*stride_in]*plus120
+        return
+    end
+
+    # 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)
+
+    w1 = convert(T, cispi(ds*2/N))
+    w2 = convert(T, cispi(ds*4/N))
+    wk1 = wk2 = one(T)
+    for k in 0:Nprime-1
+        @muladd k0 = start_out + k*stride_out
+        @muladd k1 = start_out + (k+Nprime)*stride_out
+        @muladd k2 = start_out + (k+2*Nprime)*stride_out
+        y_k0, y_k1, y_k2 = out[k0], out[k1], out[k2]
+        @muladd out[k0] = y_k0 + y_k1*wk1 + y_k2*wk2
+        @muladd out[k1] = y_k0 + y_k1*wk1*plus120 + y_k2*wk2*minus120
+        @muladd out[k2] = y_k0 + y_k1*wk1*minus120 + y_k2*wk2*plus120
+        wk1 *= w1
+        wk2 *= w2
+    end
+end
+
+function fft!(out::AbstractVector{T}, in::AbstractVector{U}, start_out::Int, start_in::Int, d::Direction, ::Pow3FFT, g::CallGraph{T}, idx::Int) where {T,U}
+    root = g[idx]
+    N = root.sz
+    s_in = root.s_in
+    s_out = root.s_out
+    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)
+    else
+        fft_pow3!(out, in, N, start_out, s_out, start_in, s_in, d, p_120, m_120)
+    end
+end

src/callgraph.jl

@@ -9,6 +9,9 @@ struct CompositeFFT <: AbstractFFTType end
 # Represents a Radix-2 Cooley-Tukey FFT
 struct Pow2FFT <: AbstractFFTType end
 
+# Represents a Radix-3 Cooley-Tukey FFT
+struct Pow3FFT <: AbstractFFTType end
+
 # Represents a Radix-4 Cooley-Tukey FFT
 struct Pow4FFT <: AbstractFFTType end
 
@@ -95,14 +98,21 @@ end
 
 # Recursively instantiate a set of `CallGraphNode`s
 function CallGraphNode!(nodes::Vector{CallGraphNode}, N::Int, workspace::Vector{Vector{T}}, s_in::Int, s_out::Int)::Int where {T}
-    if N % 2 == 0
+    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))
             return 1
         end
     end
+    if N % 3 == 0
+        if _ispow(N, 3)
+            push!(workspace, T[])
+            push!(nodes, CallGraphNode(0, 0, Pow3FFT(), N, s_in, s_out))
+            return 1
+        end
+    end
     if isprime(N)
         push!(workspace, T[])
         push!(nodes, CallGraphNode(0,0, DFT(),N, s_in, s_out))
@@ -111,6 +121,8 @@ function CallGraphNode!(nodes::Vector{CallGraphNode}, N::Int, workspace::Vector{
     Ns = [first(x) for x in collect(factor(N)) for _ in 1:last(x)]
     if Ns[1] == 2
         N1 = prod(Ns[Ns .== 2])
+    elseif Ns[1] == 3
+        N1 = prod(Ns[Ns .== 3])
     else
         # Greedy search for closest factor of N to sqrt(N)
         Nsqrt = sqrt(N)

test/onedim/complex_backward.jl

@@ -1,5 +1,5 @@
 using FFTA, Test
-test_nums = [8, 11, 15, 16, 100]
+test_nums = [8, 11, 15, 16, 27, 100]
 @testset "backward" begin
     for N in test_nums
         x = ones(ComplexF64, N)

test/onedim/complex_forward.jl

@@ -1,5 +1,5 @@
 using FFTA, Test
-test_nums = [8, 11, 15, 16, 100]
+test_nums = [8, 11, 15, 16, 27, 100]
 @testset verbose = true " forward" begin
     for N in test_nums
         x = ones(ComplexF64, N)

test/onedim/real_backward.jl

@@ -1,5 +1,5 @@
 using FFTA, Test
-test_nums = [8, 11, 15, 16, 100]
+test_nums = [8, 11, 15, 16, 27, 100]
 @testset "backward" begin
     for N in test_nums
         x = ones(Float64, N)

test/onedim/real_forward.jl

@@ -1,5 +1,5 @@
 using FFTA, Test
-test_nums = [8, 11, 15, 16, 100]
+test_nums = [8, 11, 15, 16, 27, 100]
 @testset verbose = true " forward" begin 
     for N in test_nums
         x = ones(Float64, N)

test/runtests.jl

@@ -28,7 +28,7 @@ Random.seed!(1)
             @test x ≈ x2 atol=1e-12
         end
     end
-    @testset verbose = true "2D" begin
+    @testset verbose = false "2D" begin
         @testset verbose = true "Complex" begin
             include("twodim/complex_forward.jl")
             include("twodim/complex_backward.jl")

test/twodim/complex_backward.jl

@@ -1,5 +1,5 @@
 using FFTA, Test
-test_nums = [8, 11, 15, 16, 100]
+test_nums = [8, 11, 15, 16, 27, 100]
 @testset "backward" begin
     for N in test_nums
         x = ones(ComplexF64, N, N)

test/twodim/complex_forward.jl

@@ -1,5 +1,5 @@
 using FFTA, Test
-test_nums = [8, 11, 15, 16, 100]
+test_nums = [8, 11, 15, 16, 27, 100]
 @testset " forward" begin
     for N in test_nums
         x = ones(ComplexF64, N, N)

test/twodim/real_backward.jl

@@ -1,5 +1,5 @@
 using FFTA, Test
-test_nums = [8, 11, 15, 16, 100]
+test_nums = [8, 11, 15, 16, 27, 100]
 @testset "backward" begin
     for N in test_nums
         x = ones(Float64, N, N)