From 2588d66c3354660426c313e268d0994b731367f9 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Sat, 13 Apr 2024 01:13:23 +0530 Subject: [PATCH 1/3] LinearAlgebra: Type-stability in broadcasting numbers over Bidiagonal --- .../LinearAlgebra/src/structuredbroadcast.jl | 23 +++++++++++-------- .../LinearAlgebra/test/structuredbroadcast.jl | 18 +++++++++++++++ 2 files changed, 31 insertions(+), 10 deletions(-) diff --git a/stdlib/LinearAlgebra/src/structuredbroadcast.jl b/stdlib/LinearAlgebra/src/structuredbroadcast.jl index 476d9e3e5cd02..97d4b84443112 100644 --- a/stdlib/LinearAlgebra/src/structuredbroadcast.jl +++ b/stdlib/LinearAlgebra/src/structuredbroadcast.jl @@ -66,19 +66,22 @@ structured_broadcast_alloc(bc, ::Type{Diagonal}, ::Type{ElType}, n) where {ElTyp # Bidiagonal is tricky as we need to know if it's upper or lower. The promotion # system will return Tridiagonal when there's more than one Bidiagonal, but when # there's only one, we need to make figure out upper or lower -merge_uplos(::Nothing, ::Nothing) = nothing -merge_uplos(a, ::Nothing) = a -merge_uplos(::Nothing, b) = b -merge_uplos(a, b) = a == b ? a : 'T' +# the Val flag checks if only one Bidiagonal is encountered in the broadcast expression, +# in which case we may preserve the type of the array +merge_uplos(::Tuple{Nothing, Val{A}}, ::Tuple{Nothing, Val{B}}) where {A,B} = (nothing, Val(A & B)) +merge_uplos(a::Tuple{Any,Val{A}}, ::Tuple{Nothing, Val{B}}) where {A,B} = (first(a), Val(A & B)) +merge_uplos(::Tuple{Nothing, Val{A}}, b::Tuple{Any,Val{B}}) where {A,B} = (first(b), Val(A & B)) +merge_uplos(a::Tuple{Any,Val}, b::Tuple{Any,Val}) = (first(a) == first(b) ? first(a) : 'T', Val(false)) -find_uplo(a::Bidiagonal) = a.uplo -find_uplo(a) = nothing -find_uplo(bc::Broadcasted) = mapfoldl(find_uplo, merge_uplos, Broadcast.cat_nested(bc), init=nothing) +find_uplo(a::Bidiagonal) = (a.uplo, Val(true)) +find_uplo(a) = (nothing, Val(true)) +find_uplo(bc::Broadcasted) = mapfoldl(find_uplo, merge_uplos, Broadcast.cat_nested(bc), init=(nothing, Val(true))) function structured_broadcast_alloc(bc, ::Type{Bidiagonal}, ::Type{ElType}, n) where {ElType} - uplo = n > 0 ? find_uplo(bc) : 'U' + uplo, val = find_uplo(bc) + uplo = n > 0 ? uplo : 'U' n1 = max(n - 1, 0) - if uplo == 'T' + if val isa Val{false} && uplo == 'T' return Tridiagonal(Array{ElType}(undef, n1), Array{ElType}(undef, n), Array{ElType}(undef, n1)) end return Bidiagonal(Array{ElType}(undef, n),Array{ElType}(undef, n1), uplo) @@ -170,7 +173,7 @@ isvalidstructbc(dest, bc::Broadcasted{T}) where {T<:StructuredMatrixStyle} = (isstructurepreserving(bc) || fzeropreserving(bc)) isvalidstructbc(dest::Bidiagonal, bc::Broadcasted{StructuredMatrixStyle{Bidiagonal}}) = - (size(dest, 1) < 2 || find_uplo(bc) == dest.uplo) && + (size(dest, 1) < 2 || first(find_uplo(bc)) == dest.uplo) && (isstructurepreserving(bc) || fzeropreserving(bc)) function copyto!(dest::Diagonal, bc::Broadcasted{<:StructuredMatrixStyle}) diff --git a/stdlib/LinearAlgebra/test/structuredbroadcast.jl b/stdlib/LinearAlgebra/test/structuredbroadcast.jl index 3767fc10055f2..e4e78ad94102c 100644 --- a/stdlib/LinearAlgebra/test/structuredbroadcast.jl +++ b/stdlib/LinearAlgebra/test/structuredbroadcast.jl @@ -96,6 +96,24 @@ using Test, LinearAlgebra @test broadcast!(*, Z, X, Y) == broadcast(*, fX, fY) end end + + @testset "type-stability in Bidiagonal" begin + B2 = @inferred (B -> .- B)(B) + @test B2 isa Bidiagonal + @test B2 == -1 * B + B2 = @inferred (B -> B .* 2)(B) + @test B2 isa Bidiagonal + @test B2 == B + B + B2 = @inferred (B -> 2 .* B)(B) + @test B2 isa Bidiagonal + @test B2 == B + B + B2 = @inferred (B -> B ./ 1)(B) + @test B2 isa Bidiagonal + @test B2 == B + B2 = @inferred (B -> 1 .\ B)(B) + @test B2 isa Bidiagonal + @test B2 == B + end end @testset "broadcast! where the destination is a structured matrix" begin From 3aaca7cd6e7609428461d7fd27911d03a0201e64 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Wed, 1 May 2024 23:15:25 +0530 Subject: [PATCH 2/3] count_structedmatrix instead of find_uplo flag --- .../LinearAlgebra/src/structuredbroadcast.jl | 23 +++++++++---------- 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/stdlib/LinearAlgebra/src/structuredbroadcast.jl b/stdlib/LinearAlgebra/src/structuredbroadcast.jl index 97d4b84443112..5ea858caf538b 100644 --- a/stdlib/LinearAlgebra/src/structuredbroadcast.jl +++ b/stdlib/LinearAlgebra/src/structuredbroadcast.jl @@ -66,22 +66,19 @@ structured_broadcast_alloc(bc, ::Type{Diagonal}, ::Type{ElType}, n) where {ElTyp # Bidiagonal is tricky as we need to know if it's upper or lower. The promotion # system will return Tridiagonal when there's more than one Bidiagonal, but when # there's only one, we need to make figure out upper or lower -# the Val flag checks if only one Bidiagonal is encountered in the broadcast expression, -# in which case we may preserve the type of the array -merge_uplos(::Tuple{Nothing, Val{A}}, ::Tuple{Nothing, Val{B}}) where {A,B} = (nothing, Val(A & B)) -merge_uplos(a::Tuple{Any,Val{A}}, ::Tuple{Nothing, Val{B}}) where {A,B} = (first(a), Val(A & B)) -merge_uplos(::Tuple{Nothing, Val{A}}, b::Tuple{Any,Val{B}}) where {A,B} = (first(b), Val(A & B)) -merge_uplos(a::Tuple{Any,Val}, b::Tuple{Any,Val}) = (first(a) == first(b) ? first(a) : 'T', Val(false)) +merge_uplos(::Nothing, ::Nothing) = nothing +merge_uplos(a, ::Nothing) = a +merge_uplos(::Nothing, b) = b +merge_uplos(a, b) = a == b ? a : 'T' -find_uplo(a::Bidiagonal) = (a.uplo, Val(true)) -find_uplo(a) = (nothing, Val(true)) -find_uplo(bc::Broadcasted) = mapfoldl(find_uplo, merge_uplos, Broadcast.cat_nested(bc), init=(nothing, Val(true))) +find_uplo(a::Bidiagonal) = a.uplo +find_uplo(a) = nothing +find_uplo(bc::Broadcasted) = mapfoldl(find_uplo, merge_uplos, Broadcast.cat_nested(bc), init=nothing) function structured_broadcast_alloc(bc, ::Type{Bidiagonal}, ::Type{ElType}, n) where {ElType} - uplo, val = find_uplo(bc) - uplo = n > 0 ? uplo : 'U' + uplo = n > 0 ? find_uplo(bc) : 'U' n1 = max(n - 1, 0) - if val isa Val{false} && uplo == 'T' + if count_structedmatrix(Bidiagonal, bc) > 1 && uplo == 'T' return Tridiagonal(Array{ElType}(undef, n1), Array{ElType}(undef, n), Array{ElType}(undef, n1)) end return Bidiagonal(Array{ElType}(undef, n),Array{ElType}(undef, n1), uplo) @@ -138,6 +135,8 @@ iszerodefined(::Type{<:Number}) = true iszerodefined(::Type{<:AbstractArray{T}}) where T = iszerodefined(T) iszerodefined(::Type{<:UniformScaling{T}}) where T = iszerodefined(T) +count_structedmatrix(T, bc::Broadcasted) = sum(Base.Fix2(isa, T), Broadcast.cat_nested(bc); init = 0) + fzeropreserving(bc) = (v = fzero(bc); !ismissing(v) && (iszerodefined(typeof(v)) ? iszero(v) : v == 0)) # Like sparse matrices, we assume that the zero-preservation property of a broadcasted # expression is stable. We can test the zero-preservability by applying the function From 18943f3c54e3645e27b44342a87858ceb62c4b0f Mon Sep 17 00:00:00 2001 From: N5N3 <2642243996@qq.com> Date: Thu, 2 May 2024 08:31:27 +0800 Subject: [PATCH 3/3] Update stdlib/LinearAlgebra/src/structuredbroadcast.jl --- stdlib/LinearAlgebra/src/structuredbroadcast.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/stdlib/LinearAlgebra/src/structuredbroadcast.jl b/stdlib/LinearAlgebra/src/structuredbroadcast.jl index 5ea858caf538b..c9c058c867ba0 100644 --- a/stdlib/LinearAlgebra/src/structuredbroadcast.jl +++ b/stdlib/LinearAlgebra/src/structuredbroadcast.jl @@ -172,7 +172,7 @@ isvalidstructbc(dest, bc::Broadcasted{T}) where {T<:StructuredMatrixStyle} = (isstructurepreserving(bc) || fzeropreserving(bc)) isvalidstructbc(dest::Bidiagonal, bc::Broadcasted{StructuredMatrixStyle{Bidiagonal}}) = - (size(dest, 1) < 2 || first(find_uplo(bc)) == dest.uplo) && + (size(dest, 1) < 2 || find_uplo(bc) == dest.uplo) && (isstructurepreserving(bc) || fzeropreserving(bc)) function copyto!(dest::Diagonal, bc::Broadcasted{<:StructuredMatrixStyle})