Merge pull request #6 from YingboMa/avx

Use avx micro

Author: YingboMa

.travis.yml

@@ -3,7 +3,7 @@ language: julia
 os:
   - linux
 julia:
-  - 1.0
+  - 1.1
   - nightly
 #matrix:
 #  allow_failures:

Project.toml

@@ -1,13 +1,15 @@
 name = "RecursiveFactorization"
 uuid = "f2c3362d-daeb-58d1-803e-2bc74f2840b4"
 authors = ["Yingbo Ma <[email protected]>"]
-version = "0.1.0"
+version = "0.1.1"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 
 [compat]
-julia = "1"
+LoopVectorization = "0.7"
+julia = "1.1"
 
 [extras]
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

perf/lu.jl

@@ -1,26 +1,55 @@
 using BenchmarkTools
-import LinearAlgebra, RecursiveFactorization
+using LinearAlgebra, RecursiveFactorization
 
-BenchmarkTools.DEFAULT_PARAMETERS.seconds = 0.5
+BenchmarkTools.DEFAULT_PARAMETERS.seconds = 0.08
 
-luflop(m, n) = n^3÷3 - n÷3 + m*n^2
-luflop(n) = luflop(n, n)
+function luflop(m, n=m; innerflop=2)
+    sum(1:min(m, n)) do k
+        invflop = 1
+        scaleflop = isempty(k+1:m) ? 0 : sum(k+1:m)
+        updateflop = isempty(k+1:n) ? 0 : sum(k+1:n) do j
+            isempty(k+1:m) ? 0 : sum(k+1:m) do i
+                innerflop
+            end
+        end
+        invflop + scaleflop + updateflop
+    end
+end
 
 bas_mflops = Float64[]
 rec_mflops = Float64[]
-ns = 50:50:800
+ref_mflops = Float64[]
+ns = 4:8:500
 for n in ns
+    @info "$n × $n"
     A = rand(n, n)
     bt = @belapsed LinearAlgebra.lu!($(copy(A)))
-    rt = @belapsed RecursiveFactorization.lu!($(copy(A)))
     push!(bas_mflops, luflop(n)/bt/1e9)
+
+    rt = @belapsed RecursiveFactorization.lu!($(copy(A)))
     push!(rec_mflops, luflop(n)/rt/1e9)
+
+    ref = @belapsed LinearAlgebra.generic_lufact!($(copy(A)))
+    push!(ref_mflops, luflop(n)/ref/1e9)
 end
 
-using Plots
-plt = plot(ns, bas_mflops, legend=:bottomright, lab="OpenBLAS", title="LU Factorization Benchmark", marker=:auto, dpi=150)
+using DataFrames, VegaLite
+df = DataFrame(Size = ns, RecursiveFactorization = rec_mflops, OpenBLAS = bas_mflops, Reference = ref_mflops)
+df = stack(df, [:RecursiveFactorization, :OpenBLAS, :Reference], variable_name = :Library, value_name = :GFLOPS)
+plt = df |> @vlplot(
+        :line, color = :Library,
+        x = {:Size}, y = {:GFLOPS},
+        width = 2400, height = 600
+    )
+save(joinpath(homedir(), "Pictures", "lu_float64.png"), plt)
+
+#=
+using Plot
+plt = plot(ns, bas_mflops, legend=:bottomright, lab="OpenBLAS", title="LU Factorization Benchmark", marker=:auto, dpi=300)
 plot!(plt, ns, rec_mflops, lab="RecursiveFactorization", marker=:auto)
+plot!(plt, ns, ref_mflops, lab="Reference", marker=:auto)
 xaxis!(plt, "size (N x N)")
 yaxis!(plt, "GFLOPS")
 savefig("lubench.png")
 savefig("lubench.pdf")
+=#

src/lu.jl

@@ -1,40 +1,52 @@
-using LinearAlgebra: BlasInt, BlasFloat, LU, UnitLowerTriangular, ldiv!, BLAS, checknonsingular
+using LoopVectorization
+using LinearAlgebra: BlasInt, BlasFloat, LU, UnitLowerTriangular, ldiv!, checknonsingular, BLAS, LinearAlgebra
 
-function lu(A::AbstractMatrix, pivot::Union{Val{false}, Val{true}} = Val(true);
-            check::Bool = true, blocksize::Integer = 16)
-    lu!(copy(A), pivot; check = check, blocksize = blocksize)
+function lu(A::AbstractMatrix, pivot::Union{Val{false}, Val{true}} = Val(true); kwargs...)
+    return lu!(copy(A), pivot; kwargs...)
 end
 
-function lu!(A, pivot::Union{Val{false}, Val{true}} = Val(true);
-             check::Bool = true, blocksize::Integer = 16)
-    lu!(A, Vector{BlasInt}(undef, min(size(A)...)), pivot;
-        check = check, blocksize = blocksize)
+function lu!(A, pivot::Union{Val{false}, Val{true}} = Val(true); check=true, kwargs...)
+    m, n  = size(A)
+    minmn = min(m, n)
+    F = if minmn < 10 # avx introduces small performance degradation
+        LinearAlgebra.generic_lufact!(A, pivot; check=check)
+    else
+        lu!(A, Vector{BlasInt}(undef, minmn), pivot; check=check, kwargs...)
+    end
+    return F
 end
 
+# Use a function here to make sure it gets optimized away
+# OpenBLAS' TRSM isn't very good, we use a higher threshold for recursion
+pick_threshold() = BLAS.vendor() === :mkl ? 48 : 192
+
 function lu!(A::AbstractMatrix{T}, ipiv::AbstractVector{<:Integer},
              pivot::Union{Val{false}, Val{true}} = Val(true);
-             check::Bool=true, blocksize::Integer=16) where T
-    info = Ref(zero(BlasInt))
+             check::Bool=true,
+             # the performance is not sensitive wrt blocksize, and 16 is a good default
+             blocksize::Integer=16,
+             threshold::Integer=pick_threshold()) where T
+    info = zero(BlasInt)
     m, n = size(A)
     mnmin = min(m, n)
-    if T <: BlasFloat && A isa StridedArray
-        reckernel!(A, pivot, m, mnmin, ipiv, info, blocksize)
+    if A isa StridedArray && mnmin > threshold
+        info = reckernel!(A, pivot, m, mnmin, ipiv, info, blocksize)
         if m < n # fat matrix
             # [AL AR]
             AL = @view A[:, 1:m]
             AR = @view A[:, m+1:n]
             apply_permutation!(ipiv, AR)
             ldiv!(UnitLowerTriangular(AL), AR)
         end
-      else # generic fallback
-        _generic_lufact!(A, pivot, ipiv, info)
+    else # generic fallback
+        info = _generic_lufact!(A, pivot, ipiv, info)
     end
-    check && checknonsingular(info[])
-    LU{T, typeof(A)}(A, ipiv, info[])
+    check && checknonsingular(info)
+    LU{T, typeof(A)}(A, ipiv, info)
 end
 
 function nsplit(::Type{T}, n) where T
-    k = 128 ÷ sizeof(T)
+    k = 512 ÷ (isbitstype(T) ? sizeof(T) : 8)
     k_2 = k ÷ 2
     return n >= k ? ((n + k_2) ÷ k) * k_2 : n ÷ 2
 end
@@ -44,17 +56,19 @@ Base.@propagate_inbounds function apply_permutation!(P, A)
         i′ = P[i]
         i′ == i && continue
         @simd for j in axes(A, 2)
-            A[i, j], A[i′, j] = A[i′, j], A[i, j]
+            tmp = A[i, j]
+            A[i, j] = A[i′, j]
+            A[i′, j] = tmp
         end
     end
     nothing
 end
 
-function reckernel!(A::AbstractMatrix{T}, pivot::Val{Pivot}, m, n, ipiv, info, blocksize)::Nothing where {T,Pivot}
+function reckernel!(A::AbstractMatrix{T}, pivot::Val{Pivot}, m, n, ipiv, info, blocksize)::BlasInt where {T,Pivot}
     @inbounds begin
         if n <= max(blocksize, 1)
-            _generic_lufact!(A, pivot, ipiv, info)
-            return nothing
+            info = _generic_lufact!(A, pivot, ipiv, info)
+            return info
         end
         n1 = nsplit(T, n)
         n2 = n - n1
@@ -88,7 +102,7 @@ function reckernel!(A::AbstractMatrix{T}, pivot::Val{Pivot}, m, n, ipiv, info, b
         #   [ A11 ]   [ L11 ]
         # P [     ] = [     ] U11
         #   [ A21 ]   [ L21 ]
-        reckernel!(AL, pivot, m, n1, P1, info, blocksize)
+        info = reckernel!(AL, pivot, m, n1, P1, info, blocksize)
         # [ A12 ]    [ P1 ] [ A12 ]
         # [     ] <- [    ] [     ]
         # [ A22 ]    [ 0  ] [ A22 ]
@@ -98,22 +112,33 @@ function reckernel!(A::AbstractMatrix{T}, pivot::Val{Pivot}, m, n, ipiv, info, b
         # Schur complement:
         # We have A22 = L21 U12 + A′22, hence
         # A′22 = A22 - L21 U12
-        BLAS.gemm!('N', 'N', -one(T), A21, A12, one(T), A22)
+        #mul!(A22, A21, A12, -one(T), one(T))
+        schur_complement!(A22, A21, A12)
         # record info
-        previnfo = info[]
+        previnfo = info
         # P2 A22 = L22 U22
-        reckernel!(A22, pivot, m2, n2, P2, info, blocksize)
+        info = reckernel!(A22, pivot, m2, n2, P2, info, blocksize)
         # A21 <- P2 A21
         Pivot && apply_permutation!(P2, A21)
 
-        info[] != previnfo && (info[] += n1)
-        @simd for i in 1:n2
+        info != previnfo && (info += n1)
+        @avx for i in 1:n2
             P2[i] += n1
         end
-        return nothing
+        return info
     end # inbounds
 end
 
+function schur_complement!(𝐂, 𝐀, 𝐁)
+    @avx for m ∈ 1:size(𝐀,1), n ∈ 1:size(𝐁,2)
+        𝐂ₘₙ = zero(eltype(𝐂))
+        for k ∈ 1:size(𝐀,2)
+            𝐂ₘₙ -= 𝐀[m,k] * 𝐁[k,n]
+        end
+        𝐂[m,n] = 𝐂ₘₙ + 𝐂[m,n]
+    end
+end
+
 #=
     Modified from https://github.com/JuliaLang/julia/blob/b56a9f07948255dfbe804eef25bdbada06ec2a57/stdlib/LinearAlgebra/src/lu.jl
     License is MIT: https://julialang.org/license
@@ -147,19 +172,19 @@ function _generic_lufact!(A, ::Val{Pivot}, ipiv, info) where Pivot
                 end
                 # Scale first column
                 Akkinv = inv(A[k,k])
-                @simd for i = k+1:m
+                @avx for i = k+1:m
                     A[i,k] *= Akkinv
                 end
-            elseif info[] == 0
-                info[] = k
+            elseif info == 0
+                info = k
             end
             # Update the rest
-            for j = k+1:n
-                @simd for i = k+1:m
+            @avx for j = k+1:n
+                for i = k+1:m
                     A[i,j] -= A[i,k]*A[k,j]
                 end
             end
         end
     end
-    return nothing
+    return info
 end

test/runtests.jl

@@ -11,24 +11,27 @@ const mylu = RecursiveFactorization.lu
 
 function testlu(A, MF, BF)
     @test MF.info == BF.info
-    @test norm(MF.L*MF.U - A[MF.p, :], Inf) < sqrt(eps(real(first(A))))
+    @test norm(MF.L*MF.U - A[MF.p, :], Inf) < 100sqrt(eps(real(one(float(first(A))))))
     nothing
 end
 
 @testset "Test LU factorization" begin
-    for p in (Val(true), Val(false)), T in (Float64, Float32, ComplexF64, ComplexF32, Real)
-        siz = (50, 100)
-        if isconcretetype(T)
-            A = rand(T, siz...)
-        else
-            _A = rand(50, 100)
-            A  = Matrix{T}(undef, siz...)
-            copyto!(A, _A)
-        end
-        MF = mylu(A, p)
-        BF = baselu(A, p)
-        testlu(A, MF, BF)
-        for i in 50:7:100 # test `MF.info`
+    for _p in (true, false), T in (Float64, Float32, ComplexF64, ComplexF32, Real)
+        p = Val(_p)
+        for s in [1:10; 50:80:200; 300]
+            siz = (s, s+2)
+            @info("size: $(siz[1]) × $(siz[2]), T = $T, p = $_p")
+            if isconcretetype(T)
+                A = rand(T, siz...)
+            else
+                _A = rand(siz...)
+                A  = Matrix{T}(undef, siz...)
+                copyto!(A, _A)
+            end
+            MF = mylu(A, p)
+            BF = baselu(A, p)
+            testlu(A, MF, BF)
+            i = rand(1:s) # test `MF.info`
             A[:, i] .= 0
             MF = mylu(A, p, check=false)
             BF = baselu(A, p, check=false)