Instantiate per-partition Broadcasted in ArrayPartition copy/copyto!, use type-level homogeneity check

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

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

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,
             ),