Various fixes for propagating CachedArrayStyle correctly

ext/LazyArraysBandedMatricesExt.jl

@@ -11,7 +11,7 @@ import LazyArrays: sublayout, symmetriclayout, hermitianlayout, applylayout, cac
                    PaddedColumns, CachedArray, CachedMatrix, LazyLayout, BroadcastLayout, ApplyLayout,
                    paddeddata, resizedata!, broadcastlayout, _broadcastarray2broadcasted, _broadcast_sub_arguments,
                    arguments, call, applybroadcaststyle, simplify, simplifiable, islazy_layout, lazymaterialize, _broadcast_mul_mul, _broadcast_mul_simplifiable,
-                   triangularlayout, AbstractCachedMatrix, cache_layout, _mulbanded_copyto!, ApplyBandedLayout, BroadcastBandedLayout
+                   triangularlayout, AbstractCachedMatrix, cache_layout, _mulbanded_copyto!, ApplyBandedLayout, BroadcastBandedLayout, CachedArrayStyle
 import Base: BroadcastStyle, similar, copy, broadcasted, getindex, OneTo, oneto, tail, sign, abs
 import BandedMatrices: bandedbroadcaststyle, bandwidths, isbanded, bandedcolumns, bandeddata, BandedStyle,
                         AbstractBandedLayout, AbstractBandedMatrix, BandedColumns, BandedRows, BandedSubBandedMatrix, 
@@ -25,12 +25,18 @@ hermitianlayout(::Type{<:Complex}, ::AbstractLazyBandedLayout) = HermitianLayout
 
 
 bandedbroadcaststyle(::AbstractLazyArrayStyle) = LazyArrayStyle{2}()
+bandedbroadcaststyle(::CachedArrayStyle) = CachedArrayStyle{2}()
 
 BroadcastStyle(::AbstractLazyArrayStyle{1}, ::BandedStyle) = LazyArrayStyle{2}()
 BroadcastStyle(::BandedStyle, ::AbstractLazyArrayStyle{1}) = LazyArrayStyle{2}()
 BroadcastStyle(::AbstractLazyArrayStyle{2}, ::BandedStyle) = LazyArrayStyle{2}()
 BroadcastStyle(::BandedStyle, ::AbstractLazyArrayStyle{2}) = LazyArrayStyle{2}()
 
+BroadcastStyle(::CachedArrayStyle{1}, ::BandedStyle) = CachedArrayStyle{2}()
+BroadcastStyle(::BandedStyle, ::CachedArrayStyle{1}) = CachedArrayStyle{2}()
+BroadcastStyle(::CachedArrayStyle{2}, ::BandedStyle) = CachedArrayStyle{2}()
+BroadcastStyle(::BandedStyle, ::CachedArrayStyle{2}) = CachedArrayStyle{2}()
+
 bandedcolumns(::AbstractLazyLayout) = BandedColumns{LazyLayout}()
 bandedcolumns(::DualLayout{<:AbstractLazyLayout}) = BandedColumns{LazyLayout}()
 
@@ -428,6 +434,7 @@ isbanded(::AbstractPaddedLayout, A) = true # always treat as banded
 const HcatBandedMatrix{T,N} = Hcat{T,NTuple{N,BandedMatrix{T,Matrix{T},OneTo{Int}}}}
 const VcatBandedMatrix{T,N} = Vcat{T,2,NTuple{N,BandedMatrix{T,Matrix{T},OneTo{Int}}}}
 
+BroadcastStyle(::Type{HcatBandedMatrix{T,0}}) where {T} = LazyArrayStyle{2}() # This method is needed for e.g. Hcat(), which would otherwise get recognised as a BandedStyle, giving different behaviour before-and-after loading BandedMatrices
 BroadcastStyle(::Type{HcatBandedMatrix{T,N}}) where {T,N} = BandedStyle()
 BroadcastStyle(::Type{VcatBandedMatrix{T,N}}) where {T,N} = BandedStyle()
 

src/LazyArrays.jl

@@ -18,7 +18,7 @@ import Base: *, +, -, /, <, ==, >, \, ≤, ≥, (:), @_gc_preserve_begin, @_gc_p
              oneto, add_sum, promote_op
 
 import Base.Broadcast: AbstractArrayStyle, BroadcastStyle, Broadcasted, DefaultArrayStyle, broadcasted, combine_eltypes,
-                       instantiate
+                       instantiate, result_style, Unknown
 
 import LinearAlgebra: AbstractQ, AdjOrTrans, StructuredMatrixStyle, checksquare, det, diag, dot, lmul!, logabsdet,
                       norm1, norm2, normInf, normp, pinv, rmul!, tr, tril, triu

src/cache.jl

@@ -62,6 +62,8 @@ cacheddata(A::Transpose) = transpose(cacheddata(parent(A)))
 
 maybe_cacheddata(A::AbstractCachedArray) = cacheddata(A)
 maybe_cacheddata(A::SubArray{<:Any,N,<:AbstractCachedArray}) where N = cacheddata(A)
+maybe_cacheddata(A::Adjoint) = adjoint(maybe_cacheddata(parent(A)))
+maybe_cacheddata(A::Transpose) = transpose(maybe_cacheddata(parent(A)))
 maybe_cacheddata(A) = A # no-op
 
 convert(::Type{AbstractArray{T}}, S::CachedArray{T}) where T = S
@@ -402,13 +404,13 @@ for op in (:*, :\, :+, :-)
     @eval layout_broadcasted(::ZerosLayout, ::CachedLayout, ::typeof($op), a::AbstractVector, b::AbstractVector) = broadcast(DefaultArrayStyle{1}(), $op, a, b)
 end
 
-function _bc_resizecacheddata!(::AbstractCachedLayout, a) 
+# We use BroadcastStyle dispatch instead of MemoryLayout dispatch since it's e.g. Vcats have an ApplyLayout but they can have a BroadcastStyle of CachedArrayStyle
+function _bc_resizecacheddata!(::CachedArrayStyle, a) 
     resizedata!(a, size(a)...)
     view(cacheddata(a), axes(a)...)
 end
-_bc_resizecacheddata!(::DualLayout{ML}, a) where {ML<:AbstractCachedLayout} =  _bc_resizecacheddata!(ML(), a)
 _bc_resizecacheddata!(_, a) = a
-_bc_resizecacheddata!(a) = _bc_resizecacheddata!(MemoryLayout(a), a)
+_bc_resizecacheddata!(a) = _bc_resizecacheddata!(_BroadcastStyle(typeof(a)), a)
 resize_bcargs!(bc::Broadcasted{<:CachedArrayStyle}) = broadcasted(bc.f, map(_bc_resizecacheddata!, bc.args)...)
 
 similar(bc::Broadcasted{<:CachedArrayStyle}, ::Type{T}) where T = CachedArray(zeros(T, axes(bc)))

src/lazyapplying.jl

@@ -228,7 +228,12 @@ AbstractArray{T,N}(A::ApplyArray{<:Any,N}) where {T,N} = ApplyArray{T,N}(A.f, ma
 
 @inline applied_axes(f, args...) = map(oneto, applied_size(f, args...))
 
+function cacheddata(A::ApplyArray{T, N, F}) where {T, N, F}
+    args = arguments(A)
+    ApplyArray{T, N}(A.f, map(maybe_cacheddata, args)...)
+end
 
+maybe_cacheddata(A::ApplyArray) = cacheddata(A)
 
 # immutable arrays don't need to copy.
 # Some special cases like vcat overload setindex! and therefore
@@ -329,8 +334,13 @@ function show(io::IO, A::Applied)
     print(io, ')')
 end
 
+_BroadcastStyle(x...) = BroadcastStyle(x...)
+_BroadcastStyle(::Type{<:AbstractQ}) = DefaultArrayStyle{2}()
+
 applybroadcaststyle(::Type{<:AbstractArray{<:Any,N}}, _2) where N = DefaultArrayStyle{N}()
 applybroadcaststyle(::Type{<:AbstractArray{<:Any,N}}, ::AbstractLazyLayout) where N = LazyArrayStyle{N}()
+applybroadcaststyle(::Type{<:ApplyArray{<:Any,N,<:Any,Args}}, ::AbstractLazyLayout) where {N,Args<:Tuple} = result_style(LazyArrayStyle{N}(), tuple_type_broadcastlayout(Args))
+applybroadcaststyle(::Type{<:SubArray{<:Any,<:Any,P}}, lay::AbstractLazyLayout) where {P} = applybroadcaststyle(P, lay)
 BroadcastStyle(M::Type{<:ApplyArray}) = applybroadcaststyle(M, MemoryLayout(M))
 BroadcastStyle(M::Type{<:SubArray{<:Any,N,<:ApplyArray}}) where N = applybroadcaststyle(M, MemoryLayout(M))
 

src/lazybroadcasting.jl

@@ -52,6 +52,8 @@ end
 const BroadcastVector{T,F,Args} = BroadcastArray{T,1,F,Args}
 const BroadcastMatrix{T,F,Args} = BroadcastArray{T,2,F,Args}
 
+BroadcastStyle(::Type{<:BroadcastArray{<:Any,N,<:Any,Args}}) where {N,Args<:Tuple} = result_style(LazyArrayStyle{N}(), tuple_type_broadcastlayout(Args))
+
 LazyArray(bc::Broadcasted) = BroadcastArray(bc)
 
 BroadcastArray{T,N,F,Args}(bc::Broadcasted) where {T,N,F,Args} = BroadcastArray{T,N,F,Args}(bc.f,bc.args)
@@ -120,6 +122,13 @@ sub_materialize(::BroadcastLayout, A) = converteltype(eltype(A), sub_materialize
 
 copy(bc::Broadcasted{<:AbstractLazyArrayStyle}) = BroadcastArray(bc)
 
+function cacheddata(A::BroadcastArray{T, N, F}) where {T, N, F}
+    args = arguments(A)
+    BroadcastArray{T, N}(A.f, map(maybe_cacheddata, args)...)
+end
+
+maybe_cacheddata(A::BroadcastArray) = cacheddata(A)
+
 # BroadcastArray are immutable
 copy(bc::BroadcastArray) = bc
 map(::typeof(copy), bc::BroadcastArray) = bc

src/lazyconcat.jl

@@ -256,7 +256,7 @@ function vcat_copyto!(arr::AbstractVector, arrays...)
     i = firstindex(arr)
     for a in arrays
         m = length(a)
-        copyto!(view(arr, range(i, length=m)), a)
+        m > 0 && copyto!(view(arr, range(i, length=m)), a)
         i += m
     end
     arr
@@ -439,8 +439,16 @@ end
 # to take advantage of special implementations of the sub-components
 ######
 
-BroadcastStyle(::Type{<:Vcat{<:Any,N}}) where N = LazyArrayStyle{N}()
-BroadcastStyle(::Type{<:Hcat{<:Any}}) = LazyArrayStyle{2}()
+@inline tuple_type_broadcastlayout(::Type{I}) where {I<:Tuple} = result_style(_BroadcastStyle(Base.tuple_type_head(I)), tuple_type_broadcastlayout(Base.tuple_type_tail(I)))
+@inline tuple_type_broadcastlayout(::Type{Tuple{}}) = Unknown()
+@inline tuple_type_broadcastlayout(::Type{Tuple{A}}) where {A} = _BroadcastStyle(A)
+@inline tuple_type_broadcastlayout(::Type{Tuple{A,B}}) where {A,B} = result_style(_BroadcastStyle(A), _BroadcastStyle(B))
+@inline tuple_type_broadcastlayout(::Type{Tuple{A,B,C}}) where {A,B,C} = result_style(_BroadcastStyle(A), tuple_type_broadcastlayout(Tuple{B,C}))
+@inline tuple_type_broadcastlayout(::Type{Tuple{A,B,C,D}}) where {A,B,C,D} = result_style(_BroadcastStyle(A), tuple_type_broadcastlayout(Tuple{B,C,D}))
+@inline tuple_type_broadcastlayout(::Type{Tuple{A,B,C,D,E}}) where {A,B,C,D,E} = result_style(_BroadcastStyle(A), tuple_type_broadcastlayout(Tuple{B,C,D,E}))
+
+BroadcastStyle(::Type{<:Vcat{<:Any,N,I}}) where {N,I<:Tuple} = result_style(LazyArrayStyle{N}(), tuple_type_broadcastlayout(I)) # the <:Tuple is to avoid ambiguity 
+BroadcastStyle(::Type{<:Hcat{<:Any,I}}) where {I<:Tuple} = result_style(LazyArrayStyle{2}(), tuple_type_broadcastlayout(I))
 
 # This is if we broadcast a function on a mixed concat f.([1; [2,3]])
 # such that f returns a vector, e.g., f(1) == [1,2], we don't want

src/lazyoperations.jl

@@ -438,25 +438,35 @@ copy(Q::Accumulate) = Accumulate(Q.op, copy(Q.data), copy(Q.v), Q.dims, Q.datasi
 copyto!(x::AbstractArray{<:Any,N}, C::Cumsum{<:Any,N}) where N = cumsum!(x, C.v)
 
 # How we populate the data
-function cache_filldata!(K::Accumulate{<:Any,1}, inds)
-    copyto!(view(K.data,inds), view(K.v,inds))
+function _accumulate_cache_filldata_1d!(data, v, op::F, inds) where {F}
+    copyto!(view(data,inds), view(v,inds))
     @inbounds for k in inds
-        K.data[k] = K.op(K.data[k-1], K.data[k])
+        data[k] = op(data[k-1], v[k])
     end
 end
 
-function cache_filldata!(K::Accumulate{<:Any,2}, kr, jr)
-    if K.dims == 1
+function cache_filldata!(K::Accumulate{<:Any,1}, inds)
+    _accumulate_cache_filldata_1d!(K.data, K.v, K.op, inds)
+end
+
+function _accumulate_cache_filldata_2d!(data, v, op::F, kr, jr, dims) where {F}
+    if dims == 1
+        copyto!(view(data, kr, jr), view(v, kr, jr))
         @inbounds for j in jr, k in kr
-            K.data[k,j] = K.op(K.data[k-1,j], K.v[k,j])
+            data[k,j] = op(data[k-1,j], v[k,j])
         end
-    else # K.dims == 2
+    else # dims == 2
+        copyto!(view(data, kr, jr), view(v, kr, jr))
         @inbounds for j in jr, k in kr
-            K.data[k,j] = K.op(K.data[k,j-1], K.v[k,j])
+            data[k,j] = op(data[k,j-1], v[k,j])
         end
     end
 end
 
+function cache_filldata!(K::Accumulate{<:Any,2}, kr, jr)
+    _accumulate_cache_filldata_2d!(K.data, K.v, K.op, kr, jr, K.dims)
+end
+
 
 # keep lazy
 cumsum(a::LazyArray; kwds...) = Cumsum(a; kwds...)

src/padded.jl

@@ -149,8 +149,6 @@ _vcat_resizedata!(_, B, m...) = B # by default we can't resize
 
 resizedata!(B::Vcat, m...) = _vcat_resizedata!(MemoryLayout(B), B, m...)
 
-cacheddata(B::Vcat) = Vcat(map(maybe_cacheddata, arguments(B))...)
-
 function ==(A::CachedVector{<:Any,<:Any,<:Zeros}, B::CachedVector{<:Any,<:Any,<:Zeros})
     length(A) == length(B) || return false
     n = max(A.datasize[1], B.datasize[1])

test/applytests.jl

@@ -1,10 +1,12 @@
 module ApplyTests
 
-using LazyArrays, FillArrays, ArrayLayouts, Test
-import LazyArrays: materialize, broadcasted, DefaultApplyStyle, Applied, arguments,
-            ApplyArray, ApplyMatrix, ApplyVector, LazyArrayApplyStyle, ApplyLayout, call
-import ArrayLayouts: StridedLayout
-using LinearAlgebra
+    using LazyArrays, FillArrays, ArrayLayouts, Test
+    import LazyArrays: materialize, broadcasted, DefaultApplyStyle, Applied, arguments,
+                ApplyArray, ApplyMatrix, ApplyVector, LazyArrayApplyStyle, ApplyLayout, call,
+                CachedArrayStyle, CachedArray, LazyArrayStyle
+    import ArrayLayouts: StridedLayout
+    import Base.Broadcast: BroadcastStyle
+    using LinearAlgebra
 
 @testset "Applying" begin
     @testset "Applied" begin
@@ -121,6 +123,13 @@ using LinearAlgebra
         end
         @test colsupport(M,1) == 1:5
     end
+
+    @testset "BroadcastStyle with a cached argument" begin
+        A = ApplyArray(*, rand(10, 2), rand(2, 2))
+        @test BroadcastStyle(typeof(A)) == BroadcastStyle(typeof(A')) == BroadcastStyle(typeof(view(A, 1:2, 1:2))) == LazyArrayStyle{2}()
+        A = ApplyArray(*, CachedArray(rand(10, 2)), rand(2, 2))
+        @test BroadcastStyle(typeof(A)) == BroadcastStyle(typeof(A')) == BroadcastStyle(typeof(view(A, 1:2, 1:2))) == CachedArrayStyle{2}()
+    end
 end # testset
 
 end # module

test/bandedtests.jl

@@ -1,8 +1,9 @@
 module LazyBandedTests
 using ArrayLayouts, LazyArrays, BandedMatrices, LinearAlgebra, Test
 using BandedMatrices: AbstractBandedLayout, _BandedMatrix, isbanded, BandedStyle, BandedColumns, BandedRows, resize, bandeddata
-using LazyArrays: PaddedLayout, PaddedRows, PaddedColumns, arguments, call, LazyArrayStyle, ApplyLayout, simplifiable, resizedata!, MulStyle, LazyLayout, BroadcastLayout
+using LazyArrays: PaddedLayout, PaddedRows, Accumulate, PaddedColumns, arguments, call, LazyArrayStyle, ApplyLayout, simplifiable, resizedata!, MulStyle, LazyLayout, BroadcastLayout, CachedArrayStyle
 using ArrayLayouts: OnesLayout, StridedLayout
+import Base.Broadcast: BroadcastStyle
 LazyArraysBandedMatricesExt = Base.get_extension(LazyArrays, :LazyArraysBandedMatricesExt)
 BroadcastBandedLayout = LazyArraysBandedMatricesExt.BroadcastBandedLayout
 ApplyBandedLayout = LazyArraysBandedMatricesExt.ApplyBandedLayout
@@ -967,6 +968,19 @@ LinearAlgebra.lmul!(β::Number, A::PseudoBandedMatrix) = (lmul!(β, A.data); A)
         @test A\D ≈ (B*B)\D
         @test D\A ≈ D\(B*B)
     end
+
+    @testset "BroadcastStyle with cached data" begin
+        A = _BandedMatrix(Accumulate(*, 1:10)', 1:10, 0, 0)
+        @test BroadcastStyle(typeof(A)) == CachedArrayStyle{2}()
+        @test BroadcastStyle(typeof(A')) == CachedArrayStyle{2}()
+    end
+
+    @testset "Combining BandedStyle with CachedArrayStyle" begin
+        @test BroadcastStyle(CachedArrayStyle{1}(), BandedStyle()) == CachedArrayStyle{2}()
+        @test BroadcastStyle(BandedStyle(), CachedArrayStyle{1}()) == CachedArrayStyle{2}()
+        @test BroadcastStyle(CachedArrayStyle{2}(), BandedStyle()) == CachedArrayStyle{2}()
+        @test BroadcastStyle(BandedStyle(), CachedArrayStyle{2}()) == CachedArrayStyle{2}()
+    end
 end
 
 end # module

test/broadcasttests.jl

@@ -2,8 +2,9 @@ module BroadcastTests
 
 using LazyArrays, ArrayLayouts, LinearAlgebra, FillArrays, Base64, Test
 using StaticArrays, Tracker
-import LazyArrays: BroadcastLayout, arguments, LazyArrayStyle, sub_materialize, simplifiable
+import LazyArrays: BroadcastLayout, CachedArrayStyle, arguments, LazyArrayStyle, sub_materialize, simplifiable
 import Base: broadcasted
+import Base.Broadcast: BroadcastStyle
 
 using ..InfiniteArrays
 using Infinities
@@ -480,6 +481,31 @@ using Infinities
             @test_throws "MethodError: no method matching _vec_mul_arguments"  LazyArrays._vec_mul_arguments(2, [])
         end
     end
+
+    @testset "BroadcastStyle" begin
+        @test BroadcastStyle(typeof(BroadcastVector(exp, 1:10))) == LazyArrayStyle{1}()
+        @test BroadcastStyle(typeof(BroadcastVector(+, 1:10, cache(1:10)))) == CachedArrayStyle{1}()
+        @test BroadcastStyle(typeof(BroadcastMatrix(*, Accumulate(*, 1:10)', rand(10)'))) == CachedArrayStyle{2}()
+        @test BroadcastStyle(typeof(BroadcastMatrix(*, rand(5, 5)', LazyArrays.CachedArray(rand(5, 5))))) == CachedArrayStyle{2}()
+    end
+
+    #=
+    @testset "Preserving LazyArrayStyle with adjoints/views" begin
+        for x in (BroadcastVector(exp, 1:10), ApplyArray(+, 1:10))
+            for op in (transpose, adjoint)
+                y = op(view(x, 1:5))
+                @test BroadcastStyle(typeof(y)) == LazyArrayStyle{2}()
+            end
+        end
+        x = BroadcastMatrix(*, rand(10, 10), rand(10, 10))
+        for op in (transpose, adjoint)
+            y = op(view(x, 1:5, 1:5))
+            @test BroadcastStyle(typeof(y)) == LazyArrayStyle{2}()
+        end
+        x = view(BroadcastVector(exp, 1:10), 2:7)
+        @test BroadcastStyle(typeof(x)) == LazyArrayStyle{1}()
+    end
+    =#
 end
 
 end #module