DArray: Implement LU factorization

Author: Rabab53

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))

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

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

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

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"),