Five arg mul! for UniformScaling and improvement in exp! (#40731)

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

stdlib/LinearAlgebra/src/dense.jl

@@ -636,8 +636,8 @@ function exp!(A::StridedMatrix{T}) where T<:BlasFloat
         # Compute U and V: Even/odd terms in Padé numerator & denom
         # Expansion of k=1 in for loop
         P = A2
-        U = C[2]*I + C[4]*P
-        V = C[1]*I + C[3]*P
+        U = mul!(C[4]*P, true, C[2]*I, true, true) #U = C[2]*I + C[4]*P
+        V = mul!(C[3]*P, true, C[1]*I, true, true) #V = C[1]*I + C[3]*P
         for k in 2:(div(size(C, 1), 2) - 1)
             k2 = 2 * k
             P *= A2
@@ -663,8 +663,7 @@ function exp!(A::StridedMatrix{T}) where T<:BlasFloat
         A2 = A * A
         A4 = A2 * A2
         A6 = A2 * A4
-        Ut = CC[4]*A2
-        Ut[diagind(Ut)] .+=  CC[2]
+        Ut = mul!(CC[4]*A2, true,CC[2]*I, true, true); # Ut = CC[4]*A2+CC[2]*I
         # Allocation economical version of:
         #U  = A * (A6 * (CC[14].*A6 .+ CC[12].*A4 .+ CC[10].*A2) .+
         #          CC[8].*A6 .+ CC[6].*A4 .+ Ut)
@@ -676,7 +675,7 @@ function exp!(A::StridedMatrix{T}) where T<:BlasFloat
 
         # Allocation economical version of: Vt = CC[3]*A2 (recycle Ut)
         Vt = mul!(Ut, CC[3], A2, true, false)
-        Vt[diagind(Vt)] .+=  CC[1]
+        mul!(Vt, true, CC[1]*I, true, true); # Vt += CC[1]*I
         # Allocation economical version of:
         #V  = A6 * (CC[13].*A6 .+ CC[11].*A4 .+ CC[9].*A2) .+
         #           CC[7].*A6 .+ CC[5].*A4 .+ Vt

stdlib/LinearAlgebra/src/uniformscaling.jl

@@ -284,6 +284,24 @@ end
     mul!(C, A, J.λ, alpha, beta)
 @inline mul!(C::AbstractVecOrMat, J::UniformScaling, B::AbstractVecOrMat, alpha::Number, beta::Number) =
     mul!(C, J.λ, B, alpha, beta)
+
+function mul!(out::AbstractMatrix{T}, a::Number, B::UniformScaling, α::Number, β::Number) where {T}
+    checksquare(out)
+    if iszero(β)  # zero contribution of the out matrix
+        fill!(out, zero(T))
+    elseif !isone(β)
+        rmul!(out, β)
+    end
+    s = convert(T, a*B.λ*α)
+    if !iszero(s)
+        @inbounds for i in diagind(out)
+            out[i] += s
+        end
+    end
+    return out
+end
+@inline mul!(out::AbstractMatrix, A::UniformScaling, b::Number, α::Number, β::Number)=
+    mul!(out, A.λ, UniformScaling(b), α, β)
 rmul!(A::AbstractMatrix, J::UniformScaling) = rmul!(A, J.λ)
 lmul!(J::UniformScaling, B::AbstractVecOrMat) = lmul!(J.λ, B)
 rdiv!(A::AbstractMatrix, J::UniformScaling) = rdiv!(A, J.λ)

stdlib/LinearAlgebra/test/uniformscaling.jl

@@ -454,6 +454,17 @@ end
     target = J * A * alpha + C * beta
     @test mul!(copy(C), J, A, alpha, beta) ≈ target
     @test mul!(copy(C), A, J, alpha, beta) ≈ target
+
+    a = randn()
+    C = randn(3, 3)
+    target_5mul = a*alpha*J + beta*C
+    @test mul!(copy(C), a, J, alpha, beta) ≈ target_5mul
+    @test mul!(copy(C), J, a, alpha, beta) ≈ target_5mul
+    target_5mul = beta*C # alpha = 0
+    @test mul!(copy(C), a, J, 0, beta) ≈ target_5mul
+    target_5mul = a*alpha*Matrix(J, 3, 3) # beta = 0
+    @test mul!(copy(C), a, J, alpha, 0) ≈ target_5mul
+
 end
 
 @testset "Construct Diagonal from UniformScaling" begin