Merge pull request #565 from JuliaParallel/rabab/lu

LU implementation for DArray

Author: jpsamaroo

Diff for: .buildkite/pipeline.yml

@@ -15,16 +15,6 @@
     arch: x86_64
     num_cpus: 16
 steps:
-  - label: Julia 1.8
-    timeout_in_minutes: 90
-    <<: *test
-    plugins:
-      - JuliaCI/julia#v1:
-          version: "1.8"
-      - JuliaCI/julia-test#v1:
-          julia_args: "--threads=1"
-      - JuliaCI/julia-coverage#v1:
-          codecov: true
   - label: Julia 1.9
     timeout_in_minutes: 90
     <<: *test
@@ -55,7 +45,7 @@ steps:
           julia_args: "--threads=1"
       - JuliaCI/julia-coverage#v1:
           codecov: true
-  - label: Julia 1.8 (macOS)
+  - label: Julia 1.9 (macOS)
     timeout_in_minutes: 90
     <<: *test
     agents:
@@ -64,26 +54,26 @@ steps:
       arch: x86_64
     plugins:
       - JuliaCI/julia#v1:
-          version: "1.8"
+          version: "1.9"
       - JuliaCI/julia-test#v1:
           julia_args: "--threads=1"
       - JuliaCI/julia-coverage#v1:
           codecov: true
-  - label: Julia 1.8 - TimespanLogging
+  - label: Julia 1.9 - TimespanLogging
     timeout_in_minutes: 20
     <<: *test
     plugins:
       - JuliaCI/julia#v1:
-          version: "1.8"
+          version: "1.9"
       - JuliaCI/julia-coverage#v1:
           codecov: true
     command: "julia --project -e 'using Pkg; Pkg.instantiate(); Pkg.develop(;path=\"lib/TimespanLogging\"); Pkg.test(\"TimespanLogging\")'"
-  - label: Julia 1.8 - DaggerWebDash
+  - label: Julia 1.9 - DaggerWebDash
     timeout_in_minutes: 20
     <<: *test
     plugins:
       - JuliaCI/julia#v1:
-          version: "1.8"
+          version: "1.9"
       - JuliaCI/julia-coverage#v1:
           codecov: true
     command: "julia -e 'using Pkg; Pkg.develop(;path=pwd()); Pkg.develop(;path=\"lib/TimespanLogging\"); Pkg.develop(;path=\"lib/DaggerWebDash\"); include(\"lib/DaggerWebDash/test/runtests.jl\")'"
@@ -92,7 +82,7 @@ steps:
     <<: *bench
     plugins:
       - JuliaCI/julia#v1:
-          version: "1.8"
+          version: "1.9"
       - JuliaCI/julia-test#v1:
           run_tests: false
     command: "julia -e 'using Pkg; Pkg.add(\"BenchmarkTools\"); Pkg.develop(;path=pwd())'; JULIA_PROJECT=\"$PWD\" julia --project benchmarks/benchmark.jl"
@@ -107,7 +97,7 @@ steps:
     timeout_in_minutes: 20
     plugins:
       - JuliaCI/julia#v1:
-          version: "1.8"
+          version: "1.9"
     env:
       JULIA_NUM_THREADS: "4"
     agents:

Diff for: Project.toml

@@ -59,7 +59,7 @@ Statistics = "1"
 StatsBase = "0.28, 0.29, 0.30, 0.31, 0.32, 0.33, 0.34"
 TaskLocalValues = "0.1"
 TimespanLogging = "0.1"
-julia = "1.8"
+julia = "1.9"
 
 [extras]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"

Diff for: docs/src/darray.md

@@ -446,3 +446,4 @@ From `LinearAlgebra`:
 - `*` (Out-of-place Matrix-(Matrix/Vector) multiply)
 - `mul!` (In-place Matrix-Matrix multiply)
 - `cholesky`/`cholesky!` (In-place/Out-of-place Cholesky factorization)
+- `lu`/`lu!` (In-place/Out-of-place LU factorization (`NoPivot` only))

Diff for: src/Dagger.jl

@@ -84,6 +84,7 @@ include("array/sort.jl")
 include("array/linalg.jl")
 include("array/mul.jl")
 include("array/cholesky.jl")
+include("array/lu.jl")
 include("array/random.jl")
 
 # Visualization

Diff for: src/array/linalg.jl

@@ -1,7 +1,7 @@
 function LinearAlgebra.norm2(A::DArray{T,2}) where T
     Ac = A.chunks
-    norms = [Dagger.@spawn mapreduce(LinearAlgebra.norm_sqr, +, chunk) for chunk in Ac]
-    return sqrt(sum(map(fetch, norms)))
+    norms = [Dagger.@spawn mapreduce(LinearAlgebra.norm_sqr, +, chunk) for chunk in Ac]::Matrix{DTask}
+    return sqrt(sum(map(norm->fetch(norm)::real(T), norms)))
 end
 function LinearAlgebra.norm2(A::UpperTriangular{T,<:DArray{T,2}}) where T
     Ac = parent(A).chunks

Diff for: src/array/lu.jl

@@ -0,0 +1,31 @@
+function LinearAlgebra.lu(A::DMatrix{T}, ::LinearAlgebra.NoPivot; check::Bool=true) where T
+    A_copy = LinearAlgebra._lucopy(A, LinearAlgebra.lutype(T))
+    return LinearAlgebra.lu!(A_copy, LinearAlgebra.NoPivot(); check=check)
+end
+function LinearAlgebra.lu!(A::DMatrix{T}, ::LinearAlgebra.NoPivot; check::Bool=true) where T
+    zone = one(T)
+    mzone = -one(T)
+    Ac = A.chunks
+    mt, nt = size(Ac)
+    iscomplex = T <: Complex
+    trans = iscomplex ? 'C' : 'T'
+
+    Dagger.spawn_datadeps() do
+        for k in range(1, min(mt, nt))
+            Dagger.@spawn LinearAlgebra.generic_lufact!(InOut(Ac[k, k]), LinearAlgebra.NoPivot(); check)
+            for m in range(k+1, mt)
+                Dagger.@spawn BLAS.trsm!('R', 'U', 'N', 'N', zone, In(Ac[k, k]), InOut(Ac[m, k]))
+            end
+            for n in range(k+1, nt)
+                Dagger.@spawn BLAS.trsm!('L', 'L', 'N', 'U', zone, In(Ac[k, k]), InOut(Ac[k, n]))
+                for m in range(k+1, mt)
+                    Dagger.@spawn BLAS.gemm!('N', 'N', mzone, In(Ac[m, k]), In(Ac[k, n]), zone, InOut(Ac[m, n]))
+                end
+            end
+        end
+    end
+
+    ipiv = DVector([i for i in 1:min(size(A)...)])
+
+    return LinearAlgebra.LU{T,DMatrix{T},DVector{Int}}(A, ipiv, 0)
+end

Diff for: test/array/linalg/cholesky.jl

@@ -8,6 +8,7 @@
     A = rand(T, 128, 128)
     A = A * A'
     A[diagind(A)] .+= size(A, 1)
+    B = copy(A)
     DA = view(A, Blocks(32, 32))
     if !(T <: Complex)
         @test issymmetric(DA)
@@ -20,6 +21,8 @@
     @test chol_DA isa Cholesky
     @test chol_A.L ≈ chol_DA.L
     @test chol_A.U ≈ chol_DA.U
+    # Check that cholesky did not modify A or DA
+    @test A ≈ DA ≈ B
 
     # In-place
     A_copy = copy(A)

Diff for: test/array/linalg/lu.jl

@@ -0,0 +1,33 @@
+@testset "$T" for T in (Float32, Float64, ComplexF32, ComplexF64)
+    A = rand(T, 128, 128)
+    B = copy(A)
+    DA = view(A, Blocks(64, 64))
+
+    # Out-of-place
+    lu_A = lu(A, NoPivot())
+    lu_DA = lu(DA, NoPivot())
+    @test lu_DA isa LU{T,DMatrix{T},DVector{Int}}
+    if !(T in (Float32, ComplexF32)) # FIXME: NoPivot is unstable for FP32
+        @test lu_A.L ≈ lu_DA.L
+        @test lu_A.U ≈ lu_DA.U
+    end
+    @test lu_A.P ≈ lu_DA.P
+    @test lu_A.p ≈ lu_DA.p
+    # Check that lu did not modify A or DA
+    @test A ≈ DA ≈ B
+
+    # In-place
+    A_copy = copy(A)
+    lu_A = lu!(A_copy, NoPivot())
+    lu_DA = lu!(DA, NoPivot())
+    @test lu_DA isa LU{T,DMatrix{T},DVector{Int}}
+    if !(T in (Float32, ComplexF32)) # FIXME: NoPivot is unstable for FP32
+        @test lu_A.L ≈ lu_DA.L
+        @test lu_A.U ≈ lu_DA.U
+    end
+    @test lu_A.P ≈ lu_DA.P
+    @test lu_A.p ≈ lu_DA.p
+    # Check that changes propagated to A
+    @test DA ≈ A
+    @test !(B ≈ A)
+end

Diff for: test/runtests.jl

@@ -18,6 +18,7 @@ tests = [
     ("Array - MapReduce", "array/mapreduce.jl"),
     ("Array - LinearAlgebra - Matmul", "array/linalg/matmul.jl"),
     ("Array - LinearAlgebra - Cholesky", "array/linalg/cholesky.jl"),
+    ("Array - LinearAlgebra - LU", "array/linalg/lu.jl"),
     ("Array - Random", "array/random.jl"),
     ("Caching", "cache.jl"),
     ("Disk Caching", "diskcaching.jl"),