Fix LU factorization in-place operations (JuliaLang#22774)

* Add tests for in-place operations of A_ldiv_B!(y, LU, x) for LU factorizations

* Do Ax_ldiv_B! for LU factorizations in-place

* Remove incorrect condition in LU test. The transpose should be tested for complex numbers, and might be skipped for real numbers; not the other way around. For clarity just test both unconditionally.

* Get rid of allocations, apply the (inverse) permutation in-place

base/linalg/lu.jl

@@ -241,34 +241,55 @@ function show(io::IO, F::LU)
  print(io, "\nsuccessful: $(issuccess(F))")
 end
 
+_apply_ipiv!(A::LU, B::StridedVecOrMat) = _ipiv!(A, 1 : length(A.ipiv), B)
+_apply_inverse_ipiv!(A::LU, B::StridedVecOrMat) = _ipiv!(A, length(A.ipiv) : -1 : 1, B)
+
+function _ipiv!(A::LU, order::OrdinalRange, B::StridedVecOrMat)
+ for i = order
+ if i != A.ipiv[i]
+ _swap_rows!(B, i, A.ipiv[i])
+ end
+ end
+ B
+end
+
+function _swap_rows!(B::StridedVector, i::Integer, j::Integer)
+ B[i], B[j] = B[j], B[i]
+ B
+end
+
+function _swap_rows!(B::StridedMatrix, i::Integer, j::Integer)
+ for col = 1 : size(B, 2)
+ B[i,col], B[j,col] = B[j,col], B[i,col]
+ end
+ B
+end
+
 A_ldiv_B!(A::LU{T,<:StridedMatrix}, B::StridedVecOrMat{T}) where {T<:BlasFloat} =
  @assertnonsingular LAPACK.getrs!('N', A.factors, A.ipiv, B) A.info
-A_ldiv_B!(A::LU{<:Any,<:StridedMatrix}, b::StridedVector) =
- A_ldiv_B!(UpperTriangular(A.factors),
- A_ldiv_B!(UnitLowerTriangular(A.factors), b[ipiv2perm(A.ipiv, length(b))]))
-A_ldiv_B!(A::LU{<:Any,<:StridedMatrix}, B::StridedMatrix) =
- A_ldiv_B!(UpperTriangular(A.factors),
- A_ldiv_B!(UnitLowerTriangular(A.factors), B[ipiv2perm(A.ipiv, size(B, 1)),:]))
+
+function A_ldiv_B!(A::LU{<:Any,<:StridedMatrix}, B::StridedVecOrMat)
+ _apply_ipiv!(A, B)
+ A_ldiv_B!(UpperTriangular(A.factors), A_ldiv_B!(UnitLowerTriangular(A.factors), B))
+end
 
 At_ldiv_B!(A::LU{T,<:StridedMatrix}, B::StridedVecOrMat{T}) where {T<:BlasFloat} =
  @assertnonsingular LAPACK.getrs!('T', A.factors, A.ipiv, B) A.info
-At_ldiv_B!(A::LU{<:Any,<:StridedMatrix}, b::StridedVector) =
- At_ldiv_B!(UnitLowerTriangular(A.factors),
- At_ldiv_B!(UpperTriangular(A.factors), b))[invperm(ipiv2perm(A.ipiv, length(b)))]
-At_ldiv_B!(A::LU{<:Any,<:StridedMatrix}, B::StridedMatrix) =
- At_ldiv_B!(UnitLowerTriangular(A.factors),
- At_ldiv_B!(UpperTriangular(A.factors), B))[invperm(ipiv2perm(A.ipiv, size(B,1))),:]
+
+function At_ldiv_B!(A::LU{<:Any,<:StridedMatrix}, B::StridedVecOrMat)
+ At_ldiv_B!(UnitLowerTriangular(A.factors), At_ldiv_B!(UpperTriangular(A.factors), B))
+ _apply_inverse_ipiv!(A, B)
+end
 
 Ac_ldiv_B!(F::LU{T,<:StridedMatrix}, B::StridedVecOrMat{T}) where {T<:Real} =
  At_ldiv_B!(F, B)
 Ac_ldiv_B!(A::LU{T,<:StridedMatrix}, B::StridedVecOrMat{T}) where {T<:BlasComplex} =
  @assertnonsingular LAPACK.getrs!('C', A.factors, A.ipiv, B) A.info
-Ac_ldiv_B!(A::LU{<:Any,<:StridedMatrix}, b::StridedVector) =
- Ac_ldiv_B!(UnitLowerTriangular(A.factors),
- Ac_ldiv_B!(UpperTriangular(A.factors), b))[invperm(ipiv2perm(A.ipiv, length(b)))]
-Ac_ldiv_B!(A::LU{<:Any,<:StridedMatrix}, B::StridedMatrix) =
- Ac_ldiv_B!(UnitLowerTriangular(A.factors),
- Ac_ldiv_B!(UpperTriangular(A.factors), B))[invperm(ipiv2perm(A.ipiv, size(B,1))),:]
+
+function Ac_ldiv_B!(A::LU{<:Any,<:StridedMatrix}, B::StridedVecOrMat)
+ Ac_ldiv_B!(UnitLowerTriangular(A.factors), Ac_ldiv_B!(UpperTriangular(A.factors), B))
+ _apply_inverse_ipiv!(A, B)
+end
 
 At_ldiv_Bt(A::LU{T,<:StridedMatrix}, B::StridedVecOrMat{T}) where {T<:BlasFloat} =
  @assertnonsingular LAPACK.getrs!('T', A.factors, A.ipiv, transpose(B)) A.info

test/linalg/lu.jl

@@ -84,11 +84,30 @@ dimg = randn(n)/2
  @test norm(a*(lua\c) - c, 1) < ε*κ*n # c is a vector
  @test norm(a'*(lua'\c) - c, 1) < ε*κ*n # c is a vector
  @test AbstractArray(lua) ≈ a
- if eltya <: Real && eltyb <: Real
- @test norm(a.'*(lua.'\b) - b,1) < ε*κ*n*2 # Two because the right hand side has two columns
- @test norm(a.'*(lua.'\c) - c,1) < ε*κ*n
- end
+ @test norm(a.'*(lua.'\b) - b,1) < ε*κ*n*2 # Two because the right hand side has two columns
+ @test norm(a.'*(lua.'\c) - c,1) < ε*κ*n
  end
+
+ # Test whether Ax_ldiv_B!(y, LU, x) indeed overwrites y
+ resultT = typeof(oneunit(eltyb) / oneunit(eltya))
+
+ b_dest = similar(b, resultT)
+ c_dest = similar(c, resultT)
+
+ A_ldiv_B!(b_dest, lua, b)
+ A_ldiv_B!(c_dest, lua, c)
+ @test norm(b_dest - lua \ b, 1) < ε*κ*2n
+ @test norm(c_dest - lua \ c, 1) < ε*κ*n
+
+ At_ldiv_B!(b_dest, lua, b)
+ At_ldiv_B!(c_dest, lua, c)
+ @test norm(b_dest - lua.' \ b, 1) < ε*κ*2n
+ @test norm(c_dest - lua.' \ c, 1) < ε*κ*n
+
+ Ac_ldiv_B!(b_dest, lua, b)
+ Ac_ldiv_B!(c_dest, lua, c)
+ @test norm(b_dest - lua' \ b, 1) < ε*κ*2n
+ @test norm(c_dest - lua' \ c, 1) < ε*κ*n
  end
  if eltya <: BlasFloat && eltyb <: BlasFloat
  e = rand(eltyb,n,n)