Branch in diag-bidiag mul

Author: jishnub

src/bidiag.jl

@@ -587,22 +587,27 @@ function _diag(A::Bidiagonal, k)
     end
 end
 
+_MulAddMul_nonzeroalpha(_add::MulAddMul) = _add
+function _MulAddMul_nonzeroalpha(_add::MulAddMul{ais1,bis0,Bool}) where {ais1,bis0}
+    (; beta) = _add
+    MulAddMul{true,bis0,Bool,typeof(beta)}(true, beta)
+end
+
 _mul!(C::AbstractMatrix, A::BiTriSym, B::TriSym, _add::MulAddMul) =
     _bibimul!(C, A, B, _add)
 _mul!(C::AbstractMatrix, A::BiTriSym, B::Bidiagonal, _add::MulAddMul) =
     _bibimul!(C, A, B, _add)
-function _bibimul!(C, A, B, _add::MulAddMul{ais1,bis0}) where {ais1,bis0}
+function _bibimul!(C, A, B, _add::MulAddMul)
     require_one_based_indexing(C)
     matmul_size_check(size(C), size(A), size(B))
     n = size(A,1)
     iszero(n) && return C
     # We use `_rmul_or_fill!` instead of `_modify!` here since using
     # `_modify!` in the following loop will not update the
     # off-diagonal elements for non-zero beta.
-    (; alpha, beta) = _add
-    _rmul_or_fill!(C, beta)
-    ais1 || (iszero(alpha) && return C)
-    _add_nonzeroalpha = alpha isa Bool ? MulAddMul{true,bis0,Bool,typeof(beta)}(true, beta) : _add
+    _rmul_or_fill!(C, _add.beta)
+    iszero(_add.alpha) && return C
+    _add_nonzeroalpha = _MulAddMul_nonzeroalpha(_add)
     _bibimul_nonzeroalpha!(C, A, B, _add_nonzeroalpha)
     C
 end
@@ -844,18 +849,15 @@ function __bibimul!(C, A::Bidiagonal, B::Bidiagonal, _add)
     C
 end
 
-function _mul!(C::AbstractMatrix, A::BiTriSym, B::Diagonal, alpha::Number, beta::Number)
+function _mul!(C::AbstractMatrix, A::BiTriSym, B::Diagonal, _add::MulAddMul)
     require_one_based_indexing(C)
     matmul_size_check(size(C), size(A), size(B))
     n = size(A,1)
     iszero(n) && return C
-    _rmul_or_fill!(C, beta)  # see the same use above
-    iszero(alpha) && return C
-    if alpha isa Bool
-        @stable_muladdmul _bidimul!(C, A, B, MulAddMul(true, beta))
-    else
-        @stable_muladdmul _bidimul!(C, A, B, MulAddMul(alpha, beta))
-    end
+    _rmul_or_fill!(C, _add.beta)  # see the same use above
+    iszero(_add.alpha) && return C
+    _add_nonzeroalpha = _MulAddMul_nonzeroalpha(_add)
+    _bidimul!(C, A, B, _add_nonzeroalpha)
     C
 end
 function _bidimul!(C::AbstractMatrix, A::BiTriSym, B::Diagonal, _add::MulAddMul)
@@ -921,16 +923,13 @@ function _bidimul!(C::AbstractMatrix, A::Bidiagonal, B::Diagonal, _add::MulAddMu
     C
 end
 
-function _mul!(C::Bidiagonal, A::Bidiagonal, B::Diagonal, alpha::Number, beta::Number)
+function _mul!(C::Bidiagonal, A::Bidiagonal, B::Diagonal, _add::MulAddMul)
     matmul_size_check(size(C), size(A), size(B))
     n = size(A,1)
     iszero(n) && return C
-    iszero(alpha) && return _rmul_or_fill!(C, beta)
-    if alpha isa Bool
-        @stable_muladdmul _bidimul!(C, A, B, MulAddMul(true, beta))
-    else
-        @stable_muladdmul _bidimul!(C, A, B, MulAddMul(alpha, beta))
-    end
+    iszero(_add.alpha) && return _rmul_or_fill!(C, _add.beta)
+    _add_nonzeroalpha = _MulAddMul_nonzeroalpha(_add)
+    _bidimul!(C, A, B, _add_nonzeroalpha)
     C
 end
 function _bidimul!(C::Bidiagonal, A::Bidiagonal, B::Diagonal, _add::MulAddMul)
@@ -968,7 +967,15 @@ function _mul!(C::AbstractVecOrMat, A::BiTriSym, B::AbstractVecOrMat, _add::MulA
     nA = size(A,1)
     nB = size(B,2)
     (iszero(nA) || iszero(nB)) && return C
-    iszero(_add.alpha) && return _rmul_or_fill!(C, _add.beta)
+    (; alpha, beta) = _add
+    iszero(alpha) && return _rmul_or_fill!(C, beta)
+    _add_nonzeroalpha = _MulAddMul_nonzeroalpha(_add)
+    _mul_bitrisym_left!(C, A, B, _add_nonzeroalpha)
+    return C
+end
+function _mul_bitrisym_left!(C::AbstractVecOrMat, A::Bidiagonal, B::AbstractVecOrMat, _add::MulAddMul)
+    nA = size(A,1)
+    nB = size(B,2)
     if nA <= 3
         # naive multiplication
         for I in CartesianIndices(C)
@@ -978,6 +985,7 @@ function _mul!(C::AbstractVecOrMat, A::BiTriSym, B::AbstractVecOrMat, _add::MulA
         return C
     end
     _mul_bitrisym!(C, A, B, _add)
+    return C
 end
 function _mul_bitrisym!(C::AbstractVecOrMat, A::Bidiagonal, B::AbstractVecOrMat, _add::MulAddMul)
     nA = size(A,1)
@@ -1037,7 +1045,15 @@ function _mul!(C::AbstractMatrix, A::AbstractMatrix, B::TriSym, _add::MulAddMul)
     matmul_size_check(size(C), size(A), size(B))
     n = size(A,1)
     m = size(B,2)
-    (iszero(_add.alpha) || iszero(m)) && return _rmul_or_fill!(C, _add.beta)
+    (; alpha, beta) = _add
+    (iszero(alpha) || iszero(m)) && return _rmul_or_fill!(C, beta)
+    _add_nonzeroalpha = _MulAddMul_nonzeroalpha(_add)
+    _mul_bitrisym_right!(C, A, B, _add_nonzeroalpha)
+    C
+end
+function _mul_bitrisym_right!(C::AbstractMatrix, A::AbstractMatrix, B::TriSym, _add::MulAddMul)
+    n = size(A,1)
+    m = size(B,2)
     if m == 1
         B11 = B[1,1]
         return mul!(C, A, B11, _add.alpha, _add.beta)
@@ -1073,7 +1089,14 @@ function _mul!(C::AbstractMatrix, A::AbstractMatrix, B::Bidiagonal, _add::MulAdd
     matmul_size_check(size(C), size(A), size(B))
     m, n = size(A)
     (iszero(m) || iszero(n)) && return C
-    iszero(_add.alpha) && return _rmul_or_fill!(C, _add.beta)
+    (; alpha, beta) = _add
+    iszero(alpha) && return _rmul_or_fill!(C, beta)
+    _add_nonzeroalpha = _MulAddMul_nonzeroalpha(_add)
+    _mul_bitrisym_right!(C, A, B, _add_nonzeroalpha)
+    C
+end
+function _mul_bitrisym_right!(C::AbstractMatrix, A::AbstractMatrix, B::Bidiagonal, _add::MulAddMul)
+    m, n = size(A)
     @inbounds if B.uplo == 'U'
         for j in n:-1:2, i in 1:m
             _modify!(_add, A[i,j] * B.dv[j] + A[i,j-1] * B.ev[j-1], C, (i, j))
@@ -1106,6 +1129,12 @@ function _dibimul!(C, A, B, _add)
     # ensure that we fill off-band elements in the destination
     _rmul_or_fill!(C, _add.beta)
     iszero(_add.alpha) && return C
+    _add_nonzeroalpha = _MulAddMul_nonzeroalpha(_add)
+    _dibimul_nonzeroalpha!(C, A, B, _add_nonzeroalpha)
+    C
+end
+function _dibimul_nonzeroalpha!(C, A, B, _add)
+    n = size(A,1)
     if n <= 3
         # For simplicity, use a naive multiplication for small matrices
         # that loops over all elements.
@@ -1142,14 +1171,8 @@ function _dibimul!(C, A, B, _add)
     end # inbounds
     C
 end
-function _dibimul!(C::AbstractMatrix, A::Diagonal, B::Bidiagonal, _add)
-    require_one_based_indexing(C)
-    matmul_size_check(size(C), size(A), size(B))
+function _dibimul_nonzeroalpha!(C::AbstractMatrix, A::Diagonal, B::Bidiagonal, _add)
     n = size(A,1)
-    iszero(n) && return C
-    # ensure that we fill off-band elements in the destination
-    _rmul_or_fill!(C, _add.beta)
-    iszero(_add.alpha) && return C
     Ad = A.diag
     Bdv, Bev = B.dv, B.ev
     rowshift = B.uplo == 'U' ? -1 : 1
@@ -1179,6 +1202,11 @@ function _dibimul!(C::Bidiagonal, A::Diagonal, B::Bidiagonal, _add)
     n = size(A,1)
     n == 0 && return C
     iszero(_add.alpha) && return _rmul_or_fill!(C, _add.beta)
+    _add_nonzeroalpha = _MulAddMul_nonzeroalpha(_add)
+    _dibimul_nonzeroalpha!(C, A, B, _add_nonzeroalpha)
+    C
+end
+function _dibimul_nonzeroalpha!(C::Bidiagonal, A::Diagonal, B::Bidiagonal, _add)
     Ad = A.diag
     Bdv, Bev = B.dv, B.ev
     Cdv, Cev = C.dv, C.ev