Reland "Reroute (Upper/Lower)Triangular * Diagonal through __muldiag #55984" (#56270)

This relands #55984 which was reverted in #56267. Previously, in #55984,
the destination in multiplying triangular matrices with diagonals was
also assumed to be triangular, which is not necessarily the case in
`mul!`. Tests for this case, however, were being run
non-deterministically, so this wasn't caught by the CI runs.

This improves performance:
```julia
julia> U = UpperTriangular(rand(100,100)); D = Diagonal(rand(size(U,2))); C = similar(U);

julia> @Btime mul!($C, $D, $U);
  1.517 μs (0 allocations: 0 bytes) # nightly
  1.116 μs (0 allocations: 0 bytes) # This PR
```

Author: jishnub

stdlib/LinearAlgebra/src/LinearAlgebra.jl

@@ -656,6 +656,8 @@ matprod_dest(A::StructuredMatrix, B::Diagonal, TS) = _matprod_dest_diag(A, TS)
 matprod_dest(A::Diagonal, B::StructuredMatrix, TS) = _matprod_dest_diag(B, TS)
 matprod_dest(A::Diagonal, B::Diagonal, TS) = _matprod_dest_diag(B, TS)
 _matprod_dest_diag(A, TS) = similar(A, TS)
+_matprod_dest_diag(A::UnitUpperTriangular, TS) = UpperTriangular(similar(parent(A), TS))
+_matprod_dest_diag(A::UnitLowerTriangular, TS) = LowerTriangular(similar(parent(A), TS))
 function _matprod_dest_diag(A::SymTridiagonal, TS)
     n = size(A, 1)
     ev = similar(A, TS, max(0, n-1))

stdlib/LinearAlgebra/src/diagonal.jl

@@ -397,89 +397,124 @@ function lmul!(D::Diagonal, T::Tridiagonal)
     return T
 end
 
-function __muldiag!(out, D::Diagonal, B, _add::MulAddMul{ais1,bis0}) where {ais1,bis0}
-    require_one_based_indexing(out, B)
-    alpha, beta = _add.alpha, _add.beta
-    if iszero(alpha)
-        _rmul_or_fill!(out, beta)
-    else
-        if bis0
-            @inbounds for j in axes(B, 2)
-                @simd for i in axes(B, 1)
-                    out[i,j] = D.diag[i] * B[i,j] * alpha
-                end
-            end
-        else
-            @inbounds for j in axes(B, 2)
-                @simd for i in axes(B, 1)
-                    out[i,j] = D.diag[i] * B[i,j] * alpha + out[i,j] * beta
-                end
-            end
+@inline function __muldiag_nonzeroalpha!(out, D::Diagonal, B, _add::MulAddMul)
+    @inbounds for j in axes(B, 2)
+        @simd for i in axes(B, 1)
+            _modify!(_add, D.diag[i] * B[i,j], out, (i,j))
         end
     end
-    return out
-end
-function __muldiag!(out, A, D::Diagonal, _add::MulAddMul{ais1,bis0}) where {ais1,bis0}
-    require_one_based_indexing(out, A)
-    alpha, beta = _add.alpha, _add.beta
-    if iszero(alpha)
-        _rmul_or_fill!(out, beta)
-    else
-        if bis0
-            @inbounds for j in axes(A, 2)
-                dja = D.diag[j] * alpha
-                @simd for i in axes(A, 1)
-                    out[i,j] = A[i,j] * dja
-                end
-            end
-        else
-            @inbounds for j in axes(A, 2)
-                dja = D.diag[j] * alpha
-                @simd for i in axes(A, 1)
-                    out[i,j] = A[i,j] * dja + out[i,j] * beta
-                end
+    out
+end
+_has_matching_zeros(out::UpperOrUnitUpperTriangular, A::UpperOrUnitUpperTriangular) = true
+_has_matching_zeros(out::LowerOrUnitLowerTriangular, A::LowerOrUnitLowerTriangular) = true
+_has_matching_zeros(out, A) = false
+function _rowrange_tri_stored(B::UpperOrUnitUpperTriangular, col)
+    isunit = B isa UnitUpperTriangular
+    1:min(col-isunit, size(B,1))
+end
+function _rowrange_tri_stored(B::LowerOrUnitLowerTriangular, col)
+    isunit = B isa UnitLowerTriangular
+    col+isunit:size(B,1)
+end
+_rowrange_tri_zeros(B::UpperOrUnitUpperTriangular, col) = col+1:size(B,1)
+_rowrange_tri_zeros(B::LowerOrUnitLowerTriangular, col) = 1:col-1
+function __muldiag_nonzeroalpha!(out, D::Diagonal, B::UpperOrLowerTriangular, _add::MulAddMul)
+    isunit = B isa UnitUpperOrUnitLowerTriangular
+    out_maybeparent, B_maybeparent = _has_matching_zeros(out, B) ? (parent(out), parent(B)) : (out, B)
+    for j in axes(B, 2)
+        # store the diagonal separately for unit triangular matrices
+        if isunit
+            @inbounds _modify!(_add, D.diag[j] * B[j,j], out, (j,j))
+        end
+        # The indices of out corresponding to the stored indices of B
+        rowrange = _rowrange_tri_stored(B, j)
+        @inbounds @simd for i in rowrange
+            _modify!(_add, D.diag[i] * B_maybeparent[i,j], out_maybeparent, (i,j))
+        end
+        # Fill the indices of out corresponding to the zeros of B
+        # we only fill these if out and B don't have matching zeros
+        if !_has_matching_zeros(out, B)
+            rowrange = _rowrange_tri_zeros(B, j)
+            @inbounds @simd for i in rowrange
+                _modify!(_add, D.diag[i] * B[i,j], out, (i,j))
             end
         end
     end
     return out
 end
-function __muldiag!(out::Diagonal, D1::Diagonal, D2::Diagonal, _add::MulAddMul{ais1,bis0}) where {ais1,bis0}
-    d1 = D1.diag
-    d2 = D2.diag
-    alpha, beta = _add.alpha, _add.beta
-    if iszero(alpha)
-        _rmul_or_fill!(out.diag, beta)
-    else
-        if bis0
-            @inbounds @simd for i in eachindex(out.diag)
-                out.diag[i] = d1[i] * d2[i] * alpha
-            end
-        else
-            @inbounds @simd for i in eachindex(out.diag)
-                out.diag[i] = d1[i] * d2[i] * alpha + out.diag[i] * beta
+
+@inline function __muldiag_nonzeroalpha!(out, A, D::Diagonal, _add::MulAddMul{ais1,bis0}) where {ais1,bis0}
+    beta = _add.beta
+    _add_aisone = MulAddMul{true,bis0,Bool,typeof(beta)}(true, beta)
+    @inbounds for j in axes(A, 2)
+        dja = _add(D.diag[j])
+        @simd for i in axes(A, 1)
+            _modify!(_add_aisone, A[i,j] * dja, out, (i,j))
+        end
+    end
+    out
+end
+function __muldiag_nonzeroalpha!(out, A::UpperOrLowerTriangular, D::Diagonal, _add::MulAddMul{ais1,bis0}) where {ais1,bis0}
+    isunit = A isa UnitUpperOrUnitLowerTriangular
+    beta = _add.beta
+    # since alpha is multiplied to the diagonal element of D,
+    # we may skip alpha in the second multiplication by setting ais1 to true
+    _add_aisone = MulAddMul{true,bis0,Bool,typeof(beta)}(true, beta)
+    # if both A and out have the same upper/lower triangular structure,
+    # we may directly read and write from the parents
+     out_maybeparent, A_maybeparent = _has_matching_zeros(out, A) ? (parent(out), parent(A)) : (out, A)
+    for j in axes(A, 2)
+        dja = _add(@inbounds D.diag[j])
+        # store the diagonal separately for unit triangular matrices
+        if isunit
+            @inbounds _modify!(_add_aisone, A[j,j] * dja, out, (j,j))
+        end
+        # indices of out corresponding to the stored indices of A
+        rowrange = _rowrange_tri_stored(A, j)
+        @inbounds @simd for i in rowrange
+            _modify!(_add_aisone, A_maybeparent[i,j] * dja, out_maybeparent, (i,j))
+        end
+        # Fill the indices of out corresponding to the zeros of A
+        # we only fill these if out and A don't have matching zeros
+        if !_has_matching_zeros(out, A)
+            rowrange = _rowrange_tri_zeros(A, j)
+            @inbounds @simd for i in rowrange
+                _modify!(_add_aisone, A[i,j] * dja, out, (i,j))
             end
         end
     end
-    return out
+    out
 end
-function __muldiag!(out, D1::Diagonal, D2::Diagonal, _add::MulAddMul{ais1,bis0}) where {ais1,bis0}
-    require_one_based_indexing(out)
-    alpha, beta = _add.alpha, _add.beta
-    mA = size(D1, 1)
+
+@inline function __muldiag_nonzeroalpha!(out::Diagonal, D1::Diagonal, D2::Diagonal, _add::MulAddMul)
     d1 = D1.diag
     d2 = D2.diag
-    _rmul_or_fill!(out, beta)
-    if !iszero(alpha)
-        @inbounds @simd for i in 1:mA
-            out[i,i] += d1[i] * d2[i] * alpha
-        end
+    outd = out.diag
+    @inbounds @simd for i in eachindex(d1, d2, outd)
+        _modify!(_add, d1[i] * d2[i], outd, i)
     end
-    return out
+    out
+end
+
+# ambiguity resolution
+@inline function __muldiag_nonzeroalpha!(out, D1::Diagonal, D2::Diagonal, _add::MulAddMul)
+    @inbounds for j in axes(D2, 2), i in axes(D2, 1)
+        _modify!(_add, D1.diag[i] * D2[i,j], out, (i,j))
+    end
+    out
 end
 
+# muldiag mainly handles the zero-alpha case, so that we need only
+# specialize the non-trivial case
 function _mul_diag!(out, A, B, _add)
+    require_one_based_indexing(out, A, B)
     _muldiag_size_check(size(out), size(A), size(B))
-    __muldiag!(out, A, B, _add)
+    alpha, beta = _add.alpha, _add.beta
+    if iszero(alpha)
+        _rmul_or_fill!(out, beta)
+    else
+        __muldiag_nonzeroalpha!(out, A, B, _add)
+    end
     return out
 end
 
@@ -659,31 +694,21 @@ for Tri in (:UpperTriangular, :LowerTriangular)
         @eval $fun(A::$Tri, D::Diagonal) = $Tri($fun(A.data, D))
         @eval $fun(A::$UTri, D::Diagonal) = $Tri(_setdiag!($fun(A.data, D), $f, D.diag))
     end
+    @eval *(A::$Tri{<:Any, <:StridedMaybeAdjOrTransMat}, D::Diagonal) =
+            @invoke *(A::AbstractMatrix, D::Diagonal)
+    @eval *(A::$UTri{<:Any, <:StridedMaybeAdjOrTransMat}, D::Diagonal) =
+            @invoke *(A::AbstractMatrix, D::Diagonal)
     for (fun, f) in zip((:*, :lmul!, :ldiv!, :\), (:identity, :identity, :inv, :inv))
         @eval $fun(D::Diagonal, A::$Tri) = $Tri($fun(D, A.data))
         @eval $fun(D::Diagonal, A::$UTri) = $Tri(_setdiag!($fun(D, A.data), $f, D.diag))
     end
+    @eval *(D::Diagonal, A::$Tri{<:Any, <:StridedMaybeAdjOrTransMat}) =
+            @invoke *(D::Diagonal, A::AbstractMatrix)
+    @eval *(D::Diagonal, A::$UTri{<:Any, <:StridedMaybeAdjOrTransMat}) =
+            @invoke *(D::Diagonal, A::AbstractMatrix)
     # 3-arg ldiv!
     @eval ldiv!(C::$Tri, D::Diagonal, A::$Tri) = $Tri(ldiv!(C.data, D, A.data))
     @eval ldiv!(C::$Tri, D::Diagonal, A::$UTri) = $Tri(_setdiag!(ldiv!(C.data, D, A.data), inv, D.diag))
-    # 3-arg mul! is disambiguated in special.jl
-    # 5-arg mul!
-    @eval _mul!(C::$Tri, D::Diagonal, A::$Tri, _add) = $Tri(mul!(C.data, D, A.data, _add.alpha, _add.beta))
-    @eval function _mul!(C::$Tri, D::Diagonal, A::$UTri, _add::MulAddMul{ais1,bis0}) where {ais1,bis0}
-        α, β = _add.alpha, _add.beta
-        iszero(α) && return _rmul_or_fill!(C, β)
-        diag′ = bis0 ? nothing : diag(C)
-        data = mul!(C.data, D, A.data, α, β)
-        $Tri(_setdiag!(data, _add, D.diag, diag′))
-    end
-    @eval _mul!(C::$Tri, A::$Tri, D::Diagonal, _add) = $Tri(mul!(C.data, A.data, D, _add.alpha, _add.beta))
-    @eval function _mul!(C::$Tri, A::$UTri, D::Diagonal, _add::MulAddMul{ais1,bis0}) where {ais1,bis0}
-        α, β = _add.alpha, _add.beta
-        iszero(α) && return _rmul_or_fill!(C, β)
-        diag′ = bis0 ? nothing : diag(C)
-        data = mul!(C.data, A.data, D, α, β)
-        $Tri(_setdiag!(data, _add, D.diag, diag′))
-    end
 end
 
 @inline function kron!(C::AbstractMatrix, A::Diagonal, B::Diagonal)

stdlib/LinearAlgebra/test/addmul.jl

@@ -239,4 +239,26 @@ end
     end
 end
 
+@testset "Diagonal scaling of a triangular matrix with a non-triangular destination" begin
+    for MT in (UpperTriangular, UnitUpperTriangular, LowerTriangular, UnitLowerTriangular)
+        U = MT(reshape([1:9;],3,3))
+        M = Array(U)
+        D = Diagonal(1:3)
+        A = reshape([1:9;],3,3)
+        @test mul!(copy(A), U, D, 2, 3) == M * D * 2 + A * 3
+        @test mul!(copy(A), D, U, 2, 3) == D * M * 2 + A * 3
+
+        # nan values with iszero(alpha)
+        D = Diagonal(fill(NaN,3))
+        @test mul!(copy(A), U, D, 0, 3) == A * 3
+        @test mul!(copy(A), D, U, 0, 3) == A * 3
+
+        # nan values with iszero(beta)
+        A = fill(NaN,3,3)
+        D = Diagonal(1:3)
+        @test mul!(copy(A), U, D, 2, 0) == M * D * 2
+        @test mul!(copy(A), D, U, 2, 0) == D * M * 2
+    end
+end
+
 end  # module

stdlib/LinearAlgebra/test/diagonal.jl

@@ -822,6 +822,19 @@ end
         @test @inferred(D[1,2]) isa typeof(s)
         @test all(iszero, D[1,2])
     end
+
+    @testset "mul!" begin
+        D1 = Diagonal(fill(ones(2,3), 2))
+        D2 = Diagonal(fill(ones(3,2), 2))
+        C = similar(D1, size(D1))
+        mul!(C, D1, D2)
+        @test all(x -> size(x) == (2,2), C)
+        @test C == D1 * D2
+        D = similar(D1)
+        mul!(D, D1, D2)
+        @test all(x -> size(x) == (2,2), D)
+        @test D == D1 * D2
+    end
 end
 
 @testset "Eigensystem for block diagonal (issue #30681)" begin
@@ -1188,7 +1201,7 @@ end
     @test oneunit(D3) isa typeof(D3)
 end
 
-@testset "AbstractTriangular" for (Tri, UTri) in ((UpperTriangular, UnitUpperTriangular), (LowerTriangular, UnitLowerTriangular))
+@testset "$Tri" for (Tri, UTri) in ((UpperTriangular, UnitUpperTriangular), (LowerTriangular, UnitLowerTriangular))
     A = randn(4, 4)
     TriA = Tri(A)
     UTriA = UTri(A)
@@ -1218,6 +1231,44 @@ end
     @test outTri === mul!(outTri, D, UTriA, 2, 1)::Tri == mul!(out, D, Matrix(UTriA), 2, 1)
     @test outTri === mul!(outTri, TriA, D, 2, 1)::Tri == mul!(out, Matrix(TriA), D, 2, 1)
     @test outTri === mul!(outTri, UTriA, D, 2, 1)::Tri == mul!(out, Matrix(UTriA), D, 2, 1)
+
+    # we may write to a Unit triangular if the diagonal is preserved
+    ID = Diagonal(ones(size(UTriA,2)))
+    @test mul!(copy(UTriA), UTriA, ID) == UTriA
+    @test mul!(copy(UTriA), ID, UTriA) == UTriA
+
+    @testset "partly filled parents" begin
+        M = Matrix{BigFloat}(undef, 2, 2)
+        M[1,1] = M[2,2] = 3
+        isupper = Tri == UpperTriangular
+        M[1+!isupper, 1+isupper] = 3
+        D = Diagonal(1:2)
+        T = Tri(M)
+        TA = Array(T)
+        @test T * D == TA * D
+        @test D * T == D * TA
+        @test mul!(copy(T), T, D, 2, 3) == 2T * D + 3T
+        @test mul!(copy(T), D, T, 2, 3) == 2D * T + 3T
+
+        U = UTri(M)
+        UA = Array(U)
+        @test U * D == UA * D
+        @test D * U == D * UA
+        @test mul!(copy(T), U, D, 2, 3) == 2 * UA * D + 3TA
+        @test mul!(copy(T), D, U, 2, 3) == 2 * D * UA + 3TA
+
+        M2 = Matrix{BigFloat}(undef, 2, 2)
+        M2[1+!isupper, 1+isupper] = 3
+        U = UTri(M2)
+        UA = Array(U)
+        @test U * D == UA * D
+        @test D * U == D * UA
+        ID = Diagonal(ones(size(U,2)))
+        @test mul!(copy(U), U, ID) == U
+        @test mul!(copy(U), ID, U) == U
+        @test mul!(copy(U), U, ID, 2, -1) == U
+        @test mul!(copy(U), ID, U, 2, -1) == U
+    end
 end
 
 struct SMatrix1{T} <: AbstractArray{T,2}