diff --git a/src/array_partition.jl b/src/array_partition.jl index a1220451..eb1f5137 100644 --- a/src/array_partition.jl +++ b/src/array_partition.jl @@ -500,25 +500,34 @@ end } N = npartitions(bc) @inline function f(i) - return copy(unpack(bc, i)) + return copy(Broadcast.instantiate(unpack(bc, i))) end return ArrayPartition(f, N) end +# Whether every partition of an ArrayPartition, or every ArrayPartition argument of a +# Broadcasted tree, has the same partition type. A type-level check so it constant-folds. +@inline _homog_ap(::ArrayPartition{T, S}) where {T, S} = fieldcount(S) == 0 || S <: NTuple{fieldcount(S), fieldtype(S, 1)} +@inline _homog_ap(bc::Broadcast.Broadcasted) = _homog_ap_args(bc.args) +@inline _homog_ap(x) = true +@inline _homog_ap_args(args::Tuple) = _homog_ap(first(args)) && _homog_ap_args(Base.tail(args)) +@inline _homog_ap_args(::Tuple{}) = true + @inline function Base.copyto!( dest::ArrayPartition, bc::Broadcast.Broadcasted{ArrayPartitionStyle{Style}} ) where {Style} N = npartitions(dest, bc) - # If dest is all the same underlying array type, use for-loop - if all(x isa typeof(first(dest.x)) for x in dest.x) + # If dest and every ArrayPartition argument share a single partition type, use a + # for-loop; otherwise unroll into an ntuple so each partition's type is known. + if _homog_ap(dest) && _homog_ap(bc) @inbounds for i in 1:N - copyto!(dest.x[i], unpack(bc, i)) + copyto!(dest.x[i], Broadcast.instantiate(unpack(bc, i))) end else # Fall back to original implementation for complex broadcasts @inline function f(i) - return copyto!(dest.x[i], unpack(bc, i)) + return copyto!(dest.x[i], Broadcast.instantiate(unpack(bc, i))) end ntuple(f, Val(N)) end @@ -572,7 +581,7 @@ end @inline function unpack_args(i, args::Tuple) return (unpack(args[1], i), unpack_args(i, Base.tail(args))...) end -unpack_args(i, args::Tuple{Any}) = (unpack(args[1], i),) +@inline unpack_args(i, args::Tuple{Any}) = (unpack(args[1], i),) unpack_args(::Any, args::Tuple{}) = () ## utils diff --git a/test/Core/partitions_test.jl b/test/Core/partitions_test.jl index 93a3e7b8..25c3df32 100644 --- a/test/Core/partitions_test.jl +++ b/test/Core/partitions_test.jl @@ -416,3 +416,67 @@ end struct TestIsnanFunctor end (::TestIsnanFunctor)(x) = isnan(x) @test all(TestIsnanFunctor(), AP[[NaN], [NaN]]) + +@testset "Mixed-eltype in-place broadcast does not allocate" begin + # unpack_args(i, ::Tuple{Any}) was missing @inline, forcing a runtime + # dispatch at the tuple-recursion base case whenever the partitions + # have different element types. + f!(a, b, c, d) = (@. a = b + c * d; nothing) + mk() = ArrayPartition(rand(10), rand(Float32, 20), rand(5)) + a, b, c, d = mk(), mk(), mk(), mk() + f!(a, b, c, d) + @test (@allocated f!(a, b, c, d)) == 0 +end + +# copyto!/copy on an ArrayPartition broadcast unpack each partition's slice of the +# Broadcasted tree without recomputing its axes (unpack drops them). Without +# Broadcast.instantiate before handing that slice to the partition's own copyto!, +# the partition sees a non-instantiated Broadcasted (axes === nothing). +mutable struct RATRecordingArray{T} <: AbstractVector{T} + data::Vector{T} + instantiated::Bool +end +RATRecordingArray(v::Vector) = RATRecordingArray(v, false) +Base.size(a::RATRecordingArray) = size(a.data) +Base.getindex(a::RATRecordingArray, i::Int) = a.data[i] +Base.setindex!(a::RATRecordingArray, v, i::Int) = (a.data[i] = v) +function Base.similar(a::RATRecordingArray, ::Type{T}, dims::Dims) where {T} + return RATRecordingArray(similar(a.data, T, dims), false) +end +function Base.similar(::Type{RATRecordingArray}, dims::Tuple{Union{Integer, Base.OneTo}, Vararg{Union{Integer, Base.OneTo}}}) + return RATRecordingArray(Vector{Float64}(undef, map(length, dims))) +end +Base.BroadcastStyle(::Type{<:RATRecordingArray}) = Broadcast.ArrayStyle{RATRecordingArray}() +function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{RATRecordingArray}}, ::Type{ElType}) where {ElType} + return similar(RATRecordingArray, axes(bc)) +end +function Base.copyto!(dest::RATRecordingArray, bc::Broadcast.Broadcasted{<:Broadcast.ArrayStyle{RATRecordingArray}}) + dest.instantiated = bc.axes !== nothing + copyto!(dest.data, Broadcast.instantiate(bc)) + return dest +end + +@testset "ArrayPartition broadcast instantiates partition Broadcasted" begin + a = ArrayPartition(RATRecordingArray([1.0, 2.0]), RATRecordingArray([1.0, 2.0, 3.0])) + b = ArrayPartition(RATRecordingArray([1.0, 1.0]), RATRecordingArray([1.0, 1.0, 1.0])) + a .= a .+ b + @test all(x -> x.instantiated, a.x) + + c = copy(Broadcast.broadcasted(+, a, b)) + @test all(x -> x.instantiated, c.x) +end + +@testset "Broadcast on 0-partition ArrayPartition" begin + # _homog_ap(::ArrayPartition{T,S}) called fieldtype(S, 1) without guarding + # fieldcount(S) == 0, throwing BoundsError(Tuple{}, 1) for ArrayPartition(). + dest = ArrayPartition() + dest .= ArrayPartition() .+ 1 + @test dest == ArrayPartition() + + @test copy(Broadcast.broadcasted(+, ArrayPartition(), 1)) == ArrayPartition() + @test ArrayPartition() .+ 1 == ArrayPartition() + + dest2 = ArrayPartition() + copyto!(dest2, ArrayPartition()) + @test dest2 == ArrayPartition() +end diff --git a/test/QA/qa.jl b/test/QA/qa.jl index 8b49cb62..fc39ef01 100644 --- a/test/QA/qa.jl +++ b/test/QA/qa.jl @@ -20,7 +20,7 @@ run_qa( :QRCompactWY, :SolvedVariables, :StaticArray, :StaticVecOrMat, :SymbolicTypeTrait, :_InitialValue, :_ipiv_rows!, :_maybe_reshape, :_swap_rows!, :_unsafe_getindex, :_unsafe_getindex!, :adapt_structure, - :add_sum, :flatten, :front, :index_shape, :ismutable, :issingular, + :add_sum, :flatten, :front, :index_shape, :instantiate, :ismutable, :issingular, :promote_op, :restructure, :result_style, :similar_type, :tail, :unalias, :zeromatrix, ),