Introduce a thread-safe butterfly algorithm and multi-threaded fast spherical harmonic transforms

Author: MikaelSlevinsky

src/FastTransforms.jl

@@ -19,6 +19,7 @@ import Base: getindex, setindex!, Factorization, length
 import Base.LinAlg: BlasFloat, BlasInt
 import HierarchicalMatrices: HierarchicalMatrix, unsafe_broadcasttimes!
 import HierarchicalMatrices: A_mul_B!, At_mul_B!, Ac_mul_B!
+import HierarchicalMatrices: ThreadSafeVector, threadsafezeros
 import LowRankApprox: ColPerm
 import AbstractFFTs: Plan
 

src/SphericalHarmonics/Butterfly.jl

@@ -3,10 +3,10 @@ struct Butterfly{T} <: Factorization{T}
     factors::Vector{Vector{IDPackedV{T}}}
     permutations::Vector{Vector{ColumnPermutation}}
     indices::Vector{Vector{Int}}
-    temp1::Vector{T}
-    temp2::Vector{T}
-    temp3::Vector{T}
-    temp4::Vector{T}
+    temp1::ThreadSafeVector{T}
+    temp2::ThreadSafeVector{T}
+    temp3::ThreadSafeVector{T}
+    temp4::ThreadSafeVector{T}
 end
 
 function size(B::Butterfly, dim::Integer)
@@ -106,7 +106,7 @@ function Butterfly{T}(A::AbstractMatrix{T}, L::Int; isorthogonal::Bool = false,
 
     kk = sumkmax(indices)
 
-    Butterfly(columns, factors, permutations, indices, zeros(T, kk), zeros(T, kk), zeros(T, kk), zeros(T, kk))
+    Butterfly(columns, factors, permutations, indices, threadedzeros(T, kk), threadedzeros(T, kk), threadedzeros(T, kk), threadedzeros(T, kk))
 end
 
 function sumkmax(indices::Vector{Vector{Int}})
@@ -119,7 +119,7 @@ end
 
 #### Helper
 
-function rowperm!(fwd::Bool, x::StridedVecOrMat, p::Vector{Int}, jstart::Int)
+function rowperm!(fwd::Bool, x::AbstractVecOrMat, p::Vector{Int}, jstart::Int)
     n = length(p)
     jshift = jstart-1
     scale!(p, -1)
@@ -151,7 +151,7 @@ function rowperm!(fwd::Bool, x::StridedVecOrMat, p::Vector{Int}, jstart::Int)
     x
 end
 
-function rowperm!(fwd::Bool, y::StridedVector, x::StridedVector, p::Vector{Int}, jstart::Int)
+function rowperm!(fwd::Bool, y::AbstractVector, x::AbstractVector, p::Vector{Int}, jstart::Int)
     n = length(p)
     jshift = jstart-1
     @inbounds if (fwd)
@@ -167,13 +167,13 @@ function rowperm!(fwd::Bool, y::StridedVector, x::StridedVector, p::Vector{Int},
 end
 
 ## ColumnPermutation
-A_mul_B!(A::ColPerm, B::StridedVecOrMat, jstart::Int) = rowperm!(false, B, A.p, jstart)
-At_mul_B!(A::ColPerm, B::StridedVecOrMat, jstart::Int) = rowperm!(true, B, A.p, jstart)
-Ac_mul_B!(A::ColPerm, B::StridedVecOrMat, jstart::Int) = At_mul_B!(A, B, jstart)
+A_mul_B!(A::ColPerm, B::AbstractVecOrMat, jstart::Int) = rowperm!(false, B, A.p, jstart)
+At_mul_B!(A::ColPerm, B::AbstractVecOrMat, jstart::Int) = rowperm!(true, B, A.p, jstart)
+Ac_mul_B!(A::ColPerm, B::AbstractVecOrMat, jstart::Int) = At_mul_B!(A, B, jstart)
 
-A_mul_B!(y::StridedVector, A::ColPerm, x::StridedVector, jstart::Int) = rowperm!(false, y, x, A.p, jstart)
-At_mul_B!(y::StridedVector, A::ColPerm, x::StridedVector, jstart::Int) = rowperm!(true, y, x, A.p, jstart)
-Ac_mul_B!(y::StridedVector, A::ColPerm, x::StridedVector, jstart::Int) = At_mul_B!(y, x, A, jstart)
+A_mul_B!(y::AbstractVector, A::ColPerm, x::AbstractVector, jstart::Int) = rowperm!(false, y, x, A.p, jstart)
+At_mul_B!(y::AbstractVector, A::ColPerm, x::AbstractVector, jstart::Int) = rowperm!(true, y, x, A.p, jstart)
+Ac_mul_B!(y::AbstractVector, A::ColPerm, x::AbstractVector, jstart::Int) = At_mul_B!(y, x, A, jstart)
 
 # Fast A_mul_B!, At_mul_B!, and Ac_mul_B! for an ID. These overwrite the output.
 
@@ -339,7 +339,7 @@ for f! in (:At_mul_B!,:Ac_mul_B!)
     end
 end
 
-function addtemp3totemp2!(temp2::Vector, temp3::Vector, i1::Int, i2::Int)
+function addtemp3totemp2!(temp2::AbstractVector, temp3::AbstractVector, i1::Int, i2::Int)
     z = zero(eltype(temp3))
     @inbounds @simd for i = i1:i2
         temp2[i] += temp3[i]

src/SphericalHarmonics/thinplan.jl

@@ -45,7 +45,7 @@ function Base.A_mul_B!(Y::Matrix, TP::ThinSphericalHarmonicPlan, X::Matrix)
     copy!(B, X)
     M, N = size(X)
 
-    for J = 3:2:N÷2
+    @stepthreads for J = 3:2:N÷2
         if checklayer(J-1)
             A_mul_B_col_J!(Y, BF[J-1], B, 2J)
             2J < N && A_mul_B_col_J!(Y, BF[J-1], B, 2J+1)
@@ -62,7 +62,7 @@ function Base.A_mul_B!(Y::Matrix, TP::ThinSphericalHarmonicPlan, X::Matrix)
         end
     end
 
-    for J = 2:2:N÷2
+    @stepthreads for J = 2:2:N÷2
         if checklayer(J)
             A_mul_B_col_J!(Y, BF[J-1], B, 2J)
             2J < N && A_mul_B_col_J!(Y, BF[J-1], B, 2J+1)
@@ -84,11 +84,11 @@ function Base.A_mul_B!(Y::Matrix, TP::ThinSphericalHarmonicPlan, X::Matrix)
     fill!(Y, zero(eltype(Y)))
 
     A_mul_B_col_J!!(Y, p1, B, 1)
-    for J = 2:4:N
+    @stepthreads for J = 2:4:N
         A_mul_B_col_J!!(Y, p2, B, J)
         J < N && A_mul_B_col_J!!(Y, p2, B, J+1)
     end
-    for J = 4:4:N
+    @stepthreads for J = 4:4:N
         A_mul_B_col_J!!(Y, p1, B, J)
         J < N && A_mul_B_col_J!!(Y, p1, B, J+1)
     end
@@ -101,11 +101,11 @@ function Base.At_mul_B!(Y::Matrix, TP::ThinSphericalHarmonicPlan, X::Matrix)
     copy!(B, X)
     M, N = size(X)
     A_mul_B_col_J!!(Y, p1inv, B, 1)
-    for J = 2:4:N
+    @stepthreads for J = 2:4:N
         A_mul_B_col_J!!(Y, p2inv, B, J)
         J < N && A_mul_B_col_J!!(Y, p2inv, B, J+1)
     end
-    for J = 4:4:N
+    @stepthreads for J = 4:4:N
         A_mul_B_col_J!!(Y, p1inv, B, J)
         J < N && A_mul_B_col_J!!(Y, p1inv, B, J+1)
     end
@@ -114,7 +114,7 @@ function Base.At_mul_B!(Y::Matrix, TP::ThinSphericalHarmonicPlan, X::Matrix)
     fill!(Y, zero(eltype(Y)))
     copy!(Y, 1, B, 1, 3M)
 
-    for J = 3:2:N÷2
+    @stepthreads for J = 3:2:N÷2
         if checklayer(J-1)
             At_mul_B_col_J!(Y, BF[J-1], B, 2J)
             2J < N && At_mul_B_col_J!(Y, BF[J-1], B, 2J+1)
@@ -131,7 +131,7 @@ function Base.At_mul_B!(Y::Matrix, TP::ThinSphericalHarmonicPlan, X::Matrix)
         end
     end
 
-    for J = 2:2:N÷2
+    @stepthreads for J = 2:2:N÷2
         if checklayer(J)
             At_mul_B_col_J!(Y, BF[J-1], B, 2J)
             2J < N && At_mul_B_col_J!(Y, BF[J-1], B, 2J+1)

test/basictests.jl

@@ -58,7 +58,7 @@ end
         x = rand(T, n)
         y = zeros(T, k)
 
-        @test FastTransforms.A_mul_B!(y, A, P, x, 1, 1) == A*x
+        @test norm(FastTransforms.A_mul_B!(y, A, P, x, 1, 1) - A*x) < 10eps()*norm(A*x)
 
         x = rand(T, k)
         y = zeros(T, n)