Add indirection to inplace matrix scaling (#52840)

JuliaLang · Jan 10, 2024 · 124ce94 · 124ce94
1 parent bf13a56
commit 124ce94
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/stdlib/LinearAlgebra/src/generic.jl b/stdlib/LinearAlgebra/src/generic.jl
@@ -129,15 +129,21 @@ match the length of the second, $(length(X))."))
  C
 end
 
-@inline function mul!(C::AbstractArray, s::Number, X::AbstractArray, alpha::Number, beta::Number)
+@inline mul!(C::AbstractArray, s::Number, X::AbstractArray, alpha::Number, beta::Number) =
+ _lscale_add!(C, s, X, alpha, beta)
+
+@inline function _lscale_add!(C::AbstractArray, s::Number, X::AbstractArray, alpha::Number, beta::Number)
  if axes(C) == axes(X)
  C .= (s .* X) .*ₛ alpha .+ C .*ₛ beta
  else
  generic_mul!(C, s, X, MulAddMul(alpha, beta))
  end
  return C
 end
-@inline function mul!(C::AbstractArray, X::AbstractArray, s::Number, alpha::Number, beta::Number)
+@inline mul!(C::AbstractArray, X::AbstractArray, s::Number, alpha::Number, beta::Number) =
+ _rscale_add!(C, X, s, alpha, beta)
+
+@inline function _rscale_add!(C::AbstractArray, X::AbstractArray, s::Number, alpha::Number, beta::Number)
  if axes(C) == axes(X)
  C .= (X .* s) .*ₛ alpha .+ C .*ₛ beta
  else

diff --git a/stdlib/LinearAlgebra/src/triangular.jl b/stdlib/LinearAlgebra/src/triangular.jl
@@ -532,11 +532,9 @@ function copyto!(A::T, B::T) where {T<:Union{LowerTriangular,UnitLowerTriangular
  return A
 end
 
-# Define `mul!` for (Unit){Upper,Lower}Triangular matrices times a number.
-# be permissive here and require compatibility later in _triscale!
-@inline mul!(A::AbstractTriangular, B::AbstractTriangular, C::Number, alpha::Number, beta::Number) =
+@inline _rscale_add!(A::AbstractTriangular, B::AbstractTriangular, C::Number, alpha::Number, beta::Number) =
  _triscale!(A, B, C, MulAddMul(alpha, beta))
-@inline mul!(A::AbstractTriangular, B::Number, C::AbstractTriangular, alpha::Number, beta::Number) =
+@inline _lscale_add!(A::AbstractTriangular, B::Number, C::AbstractTriangular, alpha::Number, beta::Number) =
  _triscale!(A, B, C, MulAddMul(alpha, beta))
 
 function checksize1(A, B)