From ea792d2e01388e39e077e3a26080028ca193717e Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Sat, 6 Jul 2024 17:52:39 +0530 Subject: [PATCH 1/7] Fix (l/r)mul with Diagonal --- stdlib/LinearAlgebra/src/diagonal.jl | 12 ++++++++++-- stdlib/LinearAlgebra/test/diagonal.jl | 13 +++++++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/stdlib/LinearAlgebra/src/diagonal.jl b/stdlib/LinearAlgebra/src/diagonal.jl index 5b28432fdb520..a8fc12a0e1abb 100644 --- a/stdlib/LinearAlgebra/src/diagonal.jl +++ b/stdlib/LinearAlgebra/src/diagonal.jl @@ -327,8 +327,16 @@ function (*)(D::Diagonal, V::AbstractVector) return D.diag .* V end -rmul!(A::AbstractMatrix, D::Diagonal) = @inline mul!(A, A, D) -lmul!(D::Diagonal, B::AbstractVecOrMat) = @inline mul!(B, D, B) +function rmul!(A::AbstractMatrix, D::Diagonal) + _muldiag_size_check(A, D) + A .*= permutedims(D.diag) + return A +end +function lmul!(D::Diagonal, B::AbstractVecOrMat) + _muldiag_size_check(D, B) + B .= D.diag .* B + return B +end function __muldiag!(out, D::Diagonal, B, _add::MulAddMul{ais1,bis0}) where {ais1,bis0} require_one_based_indexing(out, B) diff --git a/stdlib/LinearAlgebra/test/diagonal.jl b/stdlib/LinearAlgebra/test/diagonal.jl index 4009f841a355c..1a3b8d4fd0ea7 100644 --- a/stdlib/LinearAlgebra/test/diagonal.jl +++ b/stdlib/LinearAlgebra/test/diagonal.jl @@ -1322,4 +1322,17 @@ end @test M == D end +@testset "rmul!/lmul! with banded matrices" begin + @testset "$(nameof(typeof(B)))" for B in ( + Bidiagonal(rand(4), rand(3), :L), + Tridiagonal(rand(3), rand(4), rand(3)) + ) + BA = Array(B) + D = Diagonal(rand(size(B,1))) + DA = Array(D) + @test rmul!(copy(B), D) ≈ B * D ≈ BA * DA + @test lmul!(D, copy(B)) ≈ D * B ≈ DA * BA + end +end + end # module TestDiagonal From b7adcaf63a706f8586951673bdae30fd909176f4 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Sun, 7 Jul 2024 00:33:08 +0530 Subject: [PATCH 2/7] rmul!/lmul! for Bidiagonal --- stdlib/LinearAlgebra/src/bidiag.jl | 38 +++++++++++++++++++++++++++-- stdlib/LinearAlgebra/test/bidiag.jl | 11 +++++++++ 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/stdlib/LinearAlgebra/src/bidiag.jl b/stdlib/LinearAlgebra/src/bidiag.jl index 5de286a2c335b..6891970cd6d08 100644 --- a/stdlib/LinearAlgebra/src/bidiag.jl +++ b/stdlib/LinearAlgebra/src/bidiag.jl @@ -469,8 +469,42 @@ const BiTri = Union{Bidiagonal,Tridiagonal} @inline _mul!(C::AbstractMatrix, A::BandedMatrix, B::BandedMatrix, alpha::Number, beta::Number) = @stable_muladdmul _mul!(C, A, B, MulAddMul(alpha, beta)) -lmul!(A::Bidiagonal, B::AbstractVecOrMat) = @inline _mul!(B, A, B, MulAddMul()) -rmul!(B::AbstractMatrix, A::Bidiagonal) = @inline _mul!(B, B, A, MulAddMul()) +function lmul!(A::Bidiagonal, B::AbstractVecOrMat) + _muldiag_size_check(A, B) + (; dv, ev) = A + if A.uplo == 'U' + for k in axes(B,2) + for i in axes(ev,1) + B[i,k] = dv[i] * B[i,k] + ev[i] * B[i+1,k] + end + B[end,k] = dv[end] * B[end,k] + end + else + for k in axes(B,2) + for i in reverse(axes(dv,1)[2:end]) + B[i,k] = dv[i] * B[i,k] + ev[i-1] * B[i-1,k] + end + B[1,k] = dv[1] * B[1,k] + end + end + return B +end +function rmul!(B::AbstractMatrix, A::Bidiagonal) + _muldiag_size_check(A, B) + (; dv, ev) = A + if A.uplo == 'U' + for k in reverse(axes(dv,1)[2:end]) + @views @. B[:,k] = B[:,k] * dv[k] + B[:,k-1] * ev[k-1] + end + @views B[:,1] .*= dv[1] + else + for k in axes(ev,1) + @views @. B[:,k] = B[:,k] * dv[k] + B[:,k+1] * ev[k] + end + @views B[:,end] .*= dv[end] + end + return B +end function check_A_mul_B!_sizes(C, A, B) mA, nA = size(A) diff --git a/stdlib/LinearAlgebra/test/bidiag.jl b/stdlib/LinearAlgebra/test/bidiag.jl index 416587332c46c..44ed185270e7a 100644 --- a/stdlib/LinearAlgebra/test/bidiag.jl +++ b/stdlib/LinearAlgebra/test/bidiag.jl @@ -933,4 +933,15 @@ end @test B[1,2] == B[Int8(1),UInt16(2)] == B[big(1), Int16(2)] end +@testset "rmul!/lmul! with banded matrices" begin + A = Bidiagonal(rand(4), rand(3), :U) + @testset "$(nameof(typeof(B)))" for B in ( + Bidiagonal(rand(4), rand(3), :L), + Diagonal(rand(4)) + ) + @test_throws ArgumentError rmul!(B, A) + @test_throws ArgumentError lmul!(A, B) + end +end + end # module TestBidiagonal From f1bd47c48e42f3063cb83d988f7aad9cf0431a0a Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Sun, 7 Jul 2024 16:02:32 +0530 Subject: [PATCH 3/7] Replace broadcast by loops --- stdlib/LinearAlgebra/src/bidiag.jl | 16 ++++++++++++---- stdlib/LinearAlgebra/src/diagonal.jl | 10 ++++++++-- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/stdlib/LinearAlgebra/src/bidiag.jl b/stdlib/LinearAlgebra/src/bidiag.jl index 6891970cd6d08..a3b1369f80613 100644 --- a/stdlib/LinearAlgebra/src/bidiag.jl +++ b/stdlib/LinearAlgebra/src/bidiag.jl @@ -494,14 +494,22 @@ function rmul!(B::AbstractMatrix, A::Bidiagonal) (; dv, ev) = A if A.uplo == 'U' for k in reverse(axes(dv,1)[2:end]) - @views @. B[:,k] = B[:,k] * dv[k] + B[:,k-1] * ev[k-1] + for i in axes(B,1) + B[i,k] = B[i,k] * dv[k] + B[i,k-1] * ev[k-1] + end + end + for i in axes(B,1) + B[i,1] *= dv[1] end - @views B[:,1] .*= dv[1] else for k in axes(ev,1) - @views @. B[:,k] = B[:,k] * dv[k] + B[:,k+1] * ev[k] + for i in axes(B,1) + B[i,k] = B[i,k] * dv[k] + B[i,k+1] * ev[k] + end + end + for i in axes(B,1) + B[i,end] *= dv[end] end - @views B[:,end] .*= dv[end] end return B end diff --git a/stdlib/LinearAlgebra/src/diagonal.jl b/stdlib/LinearAlgebra/src/diagonal.jl index a8fc12a0e1abb..3c02119d056d8 100644 --- a/stdlib/LinearAlgebra/src/diagonal.jl +++ b/stdlib/LinearAlgebra/src/diagonal.jl @@ -329,12 +329,18 @@ end function rmul!(A::AbstractMatrix, D::Diagonal) _muldiag_size_check(A, D) - A .*= permutedims(D.diag) + for I in CartesianIndices(A) + row, col = Tuple(I) + @inbounds A[row, col] *= D.diag[col] + end return A end function lmul!(D::Diagonal, B::AbstractVecOrMat) _muldiag_size_check(D, B) - B .= D.diag .* B + for I in CartesianIndices(B) + row = I[1] + @inbounds B[I] = D.diag[row] * B[I] + end return B end From 5b67c41e05a4a2b4bf8e9a6cac51a8f549c99346 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Mon, 8 Jul 2024 16:01:03 +0530 Subject: [PATCH 4/7] lmul!/rmul! for Bidiagonal and Diagonal --- stdlib/LinearAlgebra/src/bidiag.jl | 26 +++++++++++++++++++++++++ stdlib/LinearAlgebra/test/bidiag.jl | 30 ++++++++++++++++++++++------- 2 files changed, 49 insertions(+), 7 deletions(-) diff --git a/stdlib/LinearAlgebra/src/bidiag.jl b/stdlib/LinearAlgebra/src/bidiag.jl index a3b1369f80613..0c8238a42d8e6 100644 --- a/stdlib/LinearAlgebra/src/bidiag.jl +++ b/stdlib/LinearAlgebra/src/bidiag.jl @@ -469,6 +469,7 @@ const BiTri = Union{Bidiagonal,Tridiagonal} @inline _mul!(C::AbstractMatrix, A::BandedMatrix, B::BandedMatrix, alpha::Number, beta::Number) = @stable_muladdmul _mul!(C, A, B, MulAddMul(alpha, beta)) +# B .= A * B function lmul!(A::Bidiagonal, B::AbstractVecOrMat) _muldiag_size_check(A, B) (; dv, ev) = A @@ -489,6 +490,19 @@ function lmul!(A::Bidiagonal, B::AbstractVecOrMat) end return B end +# B .= D * B +function lmul!(D::Diagonal, B::Bidiagonal) + _muldiag_size_check(D, B) + (; dv, ev) = B + isL = B.uplo == 'L' + dv[1] = D.diag[1] * dv[1] + for i in axes(ev,1) + ev[i] = D.diag[i + isL] * ev[i] + dv[i+1] = D.diag[i+1] * dv[i+1] + end + return B +end +# B .= B * A function rmul!(B::AbstractMatrix, A::Bidiagonal) _muldiag_size_check(A, B) (; dv, ev) = A @@ -513,6 +527,18 @@ function rmul!(B::AbstractMatrix, A::Bidiagonal) end return B end +# B .= B * D +function rmul!(B::Bidiagonal, D::Diagonal) + _muldiag_size_check(B, D) + (; dv, ev) = B + isU = B.uplo == 'U' + dv[1] *= D.diag[1] + for i in axes(ev,1) + ev[i] *= D.diag[i + isU] + dv[i+1] *= D.diag[i+1] + end + return B +end function check_A_mul_B!_sizes(C, A, B) mA, nA = size(A) diff --git a/stdlib/LinearAlgebra/test/bidiag.jl b/stdlib/LinearAlgebra/test/bidiag.jl index 44ed185270e7a..245fbfa23f5a6 100644 --- a/stdlib/LinearAlgebra/test/bidiag.jl +++ b/stdlib/LinearAlgebra/test/bidiag.jl @@ -934,13 +934,29 @@ end end @testset "rmul!/lmul! with banded matrices" begin - A = Bidiagonal(rand(4), rand(3), :U) - @testset "$(nameof(typeof(B)))" for B in ( - Bidiagonal(rand(4), rand(3), :L), - Diagonal(rand(4)) - ) - @test_throws ArgumentError rmul!(B, A) - @test_throws ArgumentError lmul!(A, B) + dv, ev = rand(4), rand(3) + for A in (Bidiagonal(dv, ev, :U), Bidiagonal(dv, ev, :L)) + @testset "$(nameof(typeof(B)))" for B in ( + Bidiagonal(dv, ev, :U), + Bidiagonal(dv, ev, :L), + Diagonal(dv) + ) + @test_throws ArgumentError rmul!(B, A) + @test_throws ArgumentError lmul!(A, B) + end + D = Diagonal(dv) + @test rmul!(copy(A), D) ≈ A * D + @test lmul!(D, copy(A)) ≈ D * A + end + @testset "non-commutative" begin + S32 = SizedArrays.SizedArray{(3,2)}(rand(3,2)) + S33 = SizedArrays.SizedArray{(3,3)}(rand(3,3)) + S22 = SizedArrays.SizedArray{(2,2)}(rand(2,2)) + B = Bidiagonal(fill(S32, 4), fill(S32, 3), :U) + D = Diagonal(fill(S22, size(B,2))) + @test rmul!(copy(B), D) ≈ B * D + D = Diagonal(fill(S33, size(B,1))) + @test lmul!(D, copy(B)) ≈ D * B end end From 5a5cb86de950a0e1f6d39f931a5ec75950e331b8 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Mon, 8 Jul 2024 16:15:43 +0530 Subject: [PATCH 5/7] Add more non-commutative tests --- stdlib/LinearAlgebra/test/bidiag.jl | 18 +++++++++++++----- test/testhelpers/SizedArrays.jl | 3 +++ 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/stdlib/LinearAlgebra/test/bidiag.jl b/stdlib/LinearAlgebra/test/bidiag.jl index 245fbfa23f5a6..164a277fd5216 100644 --- a/stdlib/LinearAlgebra/test/bidiag.jl +++ b/stdlib/LinearAlgebra/test/bidiag.jl @@ -952,11 +952,19 @@ end S32 = SizedArrays.SizedArray{(3,2)}(rand(3,2)) S33 = SizedArrays.SizedArray{(3,3)}(rand(3,3)) S22 = SizedArrays.SizedArray{(2,2)}(rand(2,2)) - B = Bidiagonal(fill(S32, 4), fill(S32, 3), :U) - D = Diagonal(fill(S22, size(B,2))) - @test rmul!(copy(B), D) ≈ B * D - D = Diagonal(fill(S33, size(B,1))) - @test lmul!(D, copy(B)) ≈ D * B + for uplo in (:L, :U) + B = Bidiagonal(fill(S32, 4), fill(S32, 3), uplo) + D = Diagonal(fill(S22, size(B,2))) + @test rmul!(copy(B), D) ≈ B * D + D = Diagonal(fill(S33, size(B,1))) + @test lmul!(D, copy(B)) ≈ D * B + end + + B = Bidiagonal(fill(S33, 4), fill(S33, 3), :U) + D = Diagonal(fill(S32, 4)) + @test lmul!(B, Array(D)) ≈ B * D + B = Bidiagonal(fill(S22, 4), fill(S22, 3), :U) + @test rmul!(Array(D), B) ≈ D * B end end diff --git a/test/testhelpers/SizedArrays.jl b/test/testhelpers/SizedArrays.jl index 2d37cead61a08..bc02fb5cbbd20 100644 --- a/test/testhelpers/SizedArrays.jl +++ b/test/testhelpers/SizedArrays.jl @@ -23,6 +23,9 @@ Base.first(::SOneTo) = 1 Base.last(r::SOneTo) = length(r) Base.show(io::IO, r::SOneTo) = print(io, "SOneTo(", length(r), ")") +Broadcast.axistype(a::Base.OneTo, s::SOneTo) = s +Broadcast.axistype(s::SOneTo, a::Base.OneTo) = s + struct SizedArray{SZ,T,N,A<:AbstractArray} <: AbstractArray{T,N} data::A function SizedArray{SZ}(data::AbstractArray{T,N}) where {SZ,T,N} From fd081610d0e6b17c0cb2497ef3e218283ee8bf77 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Mon, 8 Jul 2024 19:43:03 +0530 Subject: [PATCH 6/7] Specialize Tridiagonal-Diagonal multiplications --- stdlib/LinearAlgebra/src/diagonal.jl | 27 +++++++++++++++++++++++++++ stdlib/LinearAlgebra/test/tridiag.jl | 19 +++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/stdlib/LinearAlgebra/src/diagonal.jl b/stdlib/LinearAlgebra/src/diagonal.jl index 3c02119d056d8..bbcda9a14ebff 100644 --- a/stdlib/LinearAlgebra/src/diagonal.jl +++ b/stdlib/LinearAlgebra/src/diagonal.jl @@ -335,6 +335,19 @@ function rmul!(A::AbstractMatrix, D::Diagonal) end return A end +# T .= T * D +function rmul!(T::Tridiagonal, D::Diagonal) + _muldiag_size_check(T, D) + (; dl, d, du) = T + d[1] *= D.diag[1] + for i in axes(dl,1) + dl[i] *= D.diag[i] + du[i] *= D.diag[i+1] + d[i+1] *= D.diag[i+1] + end + return T +end + function lmul!(D::Diagonal, B::AbstractVecOrMat) _muldiag_size_check(D, B) for I in CartesianIndices(B) @@ -344,6 +357,20 @@ function lmul!(D::Diagonal, B::AbstractVecOrMat) return B end +# in-place multiplication with a diagonal +# T .= D * T +function lmul!(D::Diagonal, T::Tridiagonal) + _muldiag_size_check(D, T) + (; dl, d, du) = T + d[1] = D.diag[1] * d[1] + for i in axes(dl,1) + dl[i] = D.diag[i+1] * dl[i] + du[i] = D.diag[i] * du[i] + d[i+1] = D.diag[i+1] * d[i+1] + end + 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 diff --git a/stdlib/LinearAlgebra/test/tridiag.jl b/stdlib/LinearAlgebra/test/tridiag.jl index a067c18d7665d..fae708c4c8db4 100644 --- a/stdlib/LinearAlgebra/test/tridiag.jl +++ b/stdlib/LinearAlgebra/test/tridiag.jl @@ -892,4 +892,23 @@ end end end +@testset "rmul!/lmul! with banded matrices" begin + dl, d, du = rand(3), rand(4), rand(3) + A = Tridiagonal(dl, d, du) + D = Diagonal(d) + @test rmul!(copy(A), D) ≈ A * D + @test lmul!(D, copy(A)) ≈ D * A + + @testset "non-commutative" begin + S32 = SizedArrays.SizedArray{(3,2)}(rand(3,2)) + S33 = SizedArrays.SizedArray{(3,3)}(rand(3,3)) + S22 = SizedArrays.SizedArray{(2,2)}(rand(2,2)) + T = Tridiagonal(fill(S32,3), fill(S32, 4), fill(S32, 3)) + D = Diagonal(fill(S22, size(T,2))) + @test rmul!(copy(T), D) ≈ T * D + D = Diagonal(fill(S33, size(T,1))) + @test lmul!(D, copy(T)) ≈ D * T + end +end + end # module TestTridiagonal From 5f05464e55154b02a18bda63a0ae2a85c9fe634f Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 9 Jul 2024 12:02:44 +0530 Subject: [PATCH 7/7] Trim whitespace --- stdlib/LinearAlgebra/test/bidiag.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/stdlib/LinearAlgebra/test/bidiag.jl b/stdlib/LinearAlgebra/test/bidiag.jl index f61d7c45095ee..2ff3e9b423702 100644 --- a/stdlib/LinearAlgebra/test/bidiag.jl +++ b/stdlib/LinearAlgebra/test/bidiag.jl @@ -967,7 +967,7 @@ end @test rmul!(Array(D), B) ≈ D * B end end - + @testset "conversion to Tridiagonal for immutable bands" begin n = 4 dv = FillArrays.Fill(3, n)