Adding inplace multiplication for (unit)triangular matrices (JuliaLan…

…g#36972)

Co-authored-by: Daniel Karrasch <[email protected]>

stdlib/LinearAlgebra/src/triangular.jl

@@ -706,6 +706,29 @@ mul!(C::AbstractVecOrMat, A::AbstractTriangular, B::AbstractVecOrMat) = lmul!(A,
 mul!(C::AbstractVector, A::AbstractTriangular{<:Any,<:Adjoint}, B::Transpose{<:Any,<:AbstractVecOrMat}) = throw(MethodError(mul!, (C, A, B)))
 mul!(C::AbstractVector, A::AbstractTriangular{<:Any,<:Transpose}, B::Transpose{<:Any,<:AbstractVecOrMat}) = throw(MethodError(mul!, (C, A, B)))
 
+# preserve triangular structure in in-place multiplication
+for (cty, aty, bty) in ((:UpperTriangular, :UpperTriangular, :UpperTriangular),
+ (:UpperTriangular, :UpperTriangular, :UnitUpperTriangular),
+ (:UpperTriangular, :UnitUpperTriangular, :UpperTriangular),
+ (:UnitUpperTriangular, :UnitUpperTriangular, :UnitUpperTriangular),
+ (:LowerTriangular, :LowerTriangular, :LowerTriangular),
+ (:LowerTriangular, :LowerTriangular, :UnitLowerTriangular),
+ (:LowerTriangular, :UnitLowerTriangular, :LowerTriangular),
+ (:UnitLowerTriangular, :UnitLowerTriangular, :UnitLowerTriangular))
+ @eval function mul!(C::$cty, A::$aty, B::$bty)
+ lmul!(A, copyto!(parent(C), B))
+ return C
+ end
+
+ @eval @inline function mul!(C::$cty, A::$aty, B::$bty, alpha::Number, beta::Number)
+ if isone(alpha) && iszero(beta)
+ return mul!(C, A, B)
+ else
+ return generic_matmatmul!(C, 'N', 'N', A, B, MulAddMul(alpha, beta))
+ end
+ end
+end
+
 # direct multiplication/division
 for (t, uploc, isunitc) in ((:LowerTriangular, 'L', 'N'),
  (:UnitLowerTriangular, 'L', 'U'),

stdlib/LinearAlgebra/test/triangular.jl

@@ -630,6 +630,47 @@ end
  end
 end
 
+@testset "inplace mul of appropriate types should preserve triagular structure" begin
+ for elty1 in (Float64, ComplexF32), elty2 in (Float64, ComplexF32)
+ T = promote_type(elty1, elty2)
+ M1 = rand(elty1, 5, 5)
+ M2 = rand(elty2, 5, 5)
+ A = UpperTriangular(M1)
+ A2 = UpperTriangular(M2)
+ Au = UnitUpperTriangular(M1)
+ Au2 = UnitUpperTriangular(M2)
+ B = LowerTriangular(M1)
+ B2 = LowerTriangular(M2)
+ Bu = UnitLowerTriangular(M1)
+ Bu2 = UnitLowerTriangular(M2)
+
+ @test mul!(similar(A), A, A)::typeof(A) == A*A
+ @test mul!(similar(A, T), A, A2) ≈ A*A2
+ @test mul!(similar(A, T), A2, A) ≈ A2*A
+ @test mul!(typeof(similar(A, T))(A), A, A2, 2.0, 3.0) ≈ 2.0*A*A2 + 3.0*A
+ @test mul!(typeof(similar(A2, T))(A2), A2, A, 2.0, 3.0) ≈ 2.0*A2*A + 3.0*A2
+
+ @test mul!(similar(A), A, Au)::typeof(A) == A*Au
+ @test mul!(similar(A), Au, A)::typeof(A) == Au*A
+ @test mul!(similar(Au), Au, Au)::typeof(Au) == Au*Au
+ @test mul!(similar(A, T), A, Au2) ≈ A*Au2
+ @test mul!(similar(A, T), Au2, A) ≈ Au2*A
+ @test mul!(similar(Au2), Au2, Au2) == Au2*Au2
+
+ @test mul!(similar(B), B, B)::typeof(B) == B*B
+ @test mul!(similar(B, T), B, B2) ≈ B*B2
+ @test mul!(similar(B, T), B2, B) ≈ B2*B
+ @test mul!(typeof(similar(B, T))(B), B, B2, 2.0, 3.0) ≈ 2.0*B*B2 + 3.0*B
+ @test mul!(typeof(similar(B2, T))(B2), B2, B, 2.0, 3.0) ≈ 2.0*B2*B + 3.0*B2
+
+ @test mul!(similar(B), B, Bu)::typeof(B) == B*Bu
+ @test mul!(similar(B), Bu, B)::typeof(B) == Bu*B
+ @test mul!(similar(Bu), Bu, Bu)::typeof(Bu) == Bu*Bu
+ @test mul!(similar(B, T), B, Bu2) ≈ B*Bu2
+ @test mul!(similar(B, T), Bu2, B) ≈ Bu2*B
+ end
+end
+
 @testset "special printing of Lower/UpperTriangular" begin
  @test occursin(r"3×3 (LinearAlgebra\.)?LowerTriangular{Int64, Matrix{Int64}}:\n 2 ⋅ ⋅\n 2 2 ⋅\n 2 2 2",
  sprint(show, MIME"text/plain"(), LowerTriangular(2ones(Int64,3,3))))