diff --git a/src/blockconcat.jl b/src/blockconcat.jl index 817e1c9..81e0261 100644 --- a/src/blockconcat.jl +++ b/src/blockconcat.jl @@ -285,6 +285,7 @@ BlockBroadcastArray{T}(::typeof(Diagonal), args...) where T = BlockBroadcastMatr _block_vcat_axes(ax...) = BlockArrays._BlockedUnitRange(1,+(map(blocklasts,ax)...)) +_block_vcat_axes(ax::OneTo...) = OneTo(+(map(length,ax)...)) _block_interlace_axes(::Int, ax::Tuple{BlockedUnitRange{OneTo{Int}}}...) = _block_vcat_axes(ax...) @@ -460,10 +461,6 @@ blockinterlacelayout(_...) = LazyLayout() blockinterlacelayout(::Union{ZerosLayout,PaddedLayout,AbstractBandedLayout}...) = BlockBandedInterlaceLayout() MemoryLayout(::Type{<:BlockBroadcastMatrix{<:Any,typeof(hvcat),Arrays}}) where Arrays = blockinterlacelayout(Base.tail(LazyArrays.tuple_type_memorylayouts(Arrays))...) - -# temporary hack, need to think of how to flag as lazy for infinite case. -MemoryLayout(::Type{<:BlockBroadcastMatrix{<:Any,typeof(hcat),Arrays}}) where Arrays = LazyLayout() - MemoryLayout(::Type{<:BlockBroadcastMatrix{<:Any,typeof(Diagonal),Arrays}}) where Arrays = LazyBandedBlockBandedLayout() ## @@ -510,7 +507,48 @@ blockrowsupport(A::BlockBroadcastMatrix{<:Any,typeof(hvcat)}, k) = Block.(convex blockcolsupport(A::BlockBroadcastVector{<:Any,typeof(vcat)}, j) = Block.(convexunion(colsupport.(tail(A.args), Ref(Int.(j)))...)) -blockbroadcastlayout(FF, args...) = UnknownLayout() + +struct BlockBroadcastLayout{FF} <: AbstractLazyLayout end + +blockbroadcastlayout(::Type{FF}, args...) where FF = BlockBroadcastLayout{FF}() +sublayout(::BlockBroadcastLayout{FF}, ::Type{<:Tuple{Vararg{BlockSlice{<:BlockRange1}}}}) where FF = BlockBroadcastLayout{FF}() + +arguments(::BlockBroadcastLayout, A::BlockBroadcastArray) = A.args + +_viewsubindices(::Tuple{}, inds) = () +_viewsubindices(args::Tuple, inds) = (view(first(args), inds...), _viewsubindices(tail(args), inds)...) + +_deblockslice(bl::BlockSlice) = bl.block + +function arguments(lay::BlockBroadcastLayout, A::SubArray) + args = arguments(lay, parent(A)) + _viewsubindices(args, map(_deblockslice,parentindices(A))) +end + +# function _copyto!(::DenseColumnMajor, lay::BlockBroadcastLayout{typeof(vcat)}, dest::AbstractVector, src::AbstractVector) +# args = arguments(lay, src) +# error("todo") +# end + +function _copyto!(_, lay::BlockBroadcastLayout{typeof(hcat)}, dest::AbstractMatrix, src::AbstractMatrix) + args = arguments(lay, src) + if axes(src,2) isa OneTo # one block + j = 0 + for a in args + if a isa AbstractVector + copyto!(view(dest, :, j+1), a) + else + copyto!(view(dest, :, j .+ axes(a,2)), a) + end + j += size(a,2) + end + else + error("implement") + end +end + + + blockbroadcastlayout(::Type{typeof(vcat)}, ::PaddedLayout...) = PaddedLayout{UnknownLayout}() function paddeddata(B::BlockBroadcastVector{T,typeof(vcat)}) where T diff --git a/test/test_blockconcat.jl b/test/test_blockconcat.jl index 14f21d8..6bab1f9 100644 --- a/test/test_blockconcat.jl +++ b/test/test_blockconcat.jl @@ -292,25 +292,27 @@ end @testset "Interlace" begin @testset "vcat" begin - N = 1000 - a = 1:N - b = 11:10+N - a, b = PseudoBlockArray(a,Ones{Int}(length(a))), PseudoBlockArray(b,Ones{Int}(length(b))) - A = BlockBroadcastArray(vcat, a, b) - if VERSION < v"1.7-" - @test axes(A,1) isa BlockedUnitRange{StepRange{Int,Int}} - else - @test axes(A,1) isa BlockedUnitRange{StepRangeLen{Int,Int,Int,Int}} + @testset "basics" begin + N = 1000 + a = 1:N + b = 11:10+N + a, b = PseudoBlockArray(a,Ones{Int}(length(a))), PseudoBlockArray(b,Ones{Int}(length(b))) + A = BlockBroadcastArray(vcat, a, b) + if VERSION < v"1.7-" + @test axes(A,1) isa BlockedUnitRange{StepRange{Int,Int}} + else + @test axes(A,1) isa BlockedUnitRange{StepRangeLen{Int,Int,Int,Int}} + end + @test @allocated(axes(A)) ≤ 50 + @test A[Block(1)] == PseudoBlockArray(A)[Block(1)] == [A[1],A[2]] == [1,11] + @test A[Block(N)] == PseudoBlockArray(A)[Block(N)] == [1000,1010] + @test convert(AbstractArray{Int},A) ≡ convert(AbstractVector{Int},A) ≡ A + @test copy(A) == AbstractArray{Float64}(A) == AbstractVector{Float64}(A) == convert(AbstractArray{Float64},A) == convert(AbstractVector{Float64},A) == A + @test copy(A') == A' + + @test A .+ A isa BroadcastArray + @test A' .+ A isa BroadcastArray end - @test @allocated(axes(A)) ≤ 50 - @test A[Block(1)] == PseudoBlockArray(A)[Block(1)] == [A[1],A[2]] == [1,11] - @test A[Block(N)] == PseudoBlockArray(A)[Block(N)] == [1000,1010] - @test convert(AbstractArray{Int},A) ≡ convert(AbstractVector{Int},A) ≡ A - @test copy(A) == AbstractArray{Float64}(A) == AbstractVector{Float64}(A) == convert(AbstractArray{Float64},A) == convert(AbstractVector{Float64},A) == A - @test copy(A') == A' - - @test A .+ A isa BroadcastArray - @test A' .+ A isa BroadcastArray @testset "padded" begin a = Vcat(randn(3), Zeros(10)) @@ -332,6 +334,18 @@ end @test_throws BoundsError A[7] @test !isassigned(A,7) end + + @testset "fast materialize" begin + N = 500_000 + a = 1:N + b = 11:10+N + a, b = PseudoBlockArray(a,Ones{Int}(length(a))), PseudoBlockArray(b,Ones{Int}(length(b))) + A = BlockBroadcastArray(vcat, a, b) + + ret = PseudoBlockArray(similar(A)) + @time copyto!(ret, A); + @time copyto!(@view(ret.blocks[1:2:end]), a.blocks) + end end @testset "hcat" begin N = 1000 @@ -354,6 +368,20 @@ end v = BlockVector(randn(3), 1:2) H = BlockBroadcastArray(hcat, v, v) @test H[Block(2,1)] == [v[Block(2)] v[Block(2)]] + + @testset "fast materialize" begin + N = 500_000 + a = PseudoBlockArray(1:N,Ones{Int}(N)) + M = BlockBroadcastArray(hcat, a, Zeros{Int}(axes(a))) + @time ret = PseudoBlockArray(M); + + @time copyto!(view(ret.blocks,:,1), 1:N); + @time fill!(view(ret.blocks,:,2), 0); + + @ent MemoryLayout(M) + args = M.args + @time copyto!(view(ret,:,1), args[1]); + end end @testset "hvcat" begin a = unitblocks(randn(2,3))