From 726bfc3988267c3cd352cc291257bd000fb4d432 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Thu, 8 Oct 2026 06:40:55 -0400 Subject: [PATCH 1/3] Add @inline to unpack_args(i, ::Tuple{Any}) in ArrayPartition broadcast The tuple-recursion base case was missing @inline, forcing a runtime dispatch whenever an ArrayPartition broadcast mixes partition eltypes (e.g. Float32 and Float64 partitions). Homogeneous partitions were unaffected since the compiler could already fold the single-method recursion. Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Agent-Harness: Claude Code 2.1.294 Agent-Model: claude-sonnet-5 Agent-Session: https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv --- src/array_partition.jl | 2 +- test/Core/partitions_test.jl | 11 +++++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/src/array_partition.jl b/src/array_partition.jl index a1220451..9504b7a9 100644 --- a/src/array_partition.jl +++ b/src/array_partition.jl @@ -572,7 +572,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..c5b4debf 100644 --- a/test/Core/partitions_test.jl +++ b/test/Core/partitions_test.jl @@ -416,3 +416,14 @@ 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 From 31c7293273d56e1ba09be2b0ae01d6bd028df4ab Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Thu, 8 Oct 2026 07:22:50 -0400 Subject: [PATCH 2/3] Instantiate per-partition Broadcasted in ArrayPartition copy/copyto!, use type-level homogeneity check unpack() drops a Broadcasted's axes to make recomputing them cheaper, but the per-partition copyto!/copy never recomputed them before handing the slice to the partition's own copyto!. On Julia 1.12 the resulting Broadcasted{Style,Nothing,...} takes a slower, non-SIMD path in the partition's own broadcast machinery. Wrap each per-partition unpack in Broadcast.instantiate before copying/copying into it. Also replace the runtime `all(x isa typeof(first(dest.x)) for x in dest.x)` homogeneity check (which only looked at dest, allocated, and ran on every call) with a type-level check over dest and every ArrayPartition argument of the broadcast, so it constant-folds. Builds on the @inline fix in the previous PR: without it, re-adding Broadcast.instantiate on a mixed-partition-eltype broadcast makes the existing unpack_args dispatch strictly worse (confirmed: AP3mixed `b+c*d` 5203 -> 9390 ns, allocs 21 -> 42, with instantiate alone). Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Agent-Harness: Claude Code 2.1.294 Agent-Model: claude-sonnet-5 Agent-Session: https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv --- src/array_partition.jl | 21 +++++++++++++++----- test/Core/partitions_test.jl | 38 ++++++++++++++++++++++++++++++++++++ test/QA/qa.jl | 2 +- 3 files changed, 55 insertions(+), 6 deletions(-) diff --git a/src/array_partition.jl b/src/array_partition.jl index 9504b7a9..dbfe9f32 100644 --- a/src/array_partition.jl +++ b/src/array_partition.jl @@ -500,25 +500,36 @@ 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 +# (the previous `all(x isa typeof(first(dest.x)) for x in dest.x)` ran, and allocated, at +# runtime, and only looked at `dest`, ignoring the broadcast arguments). +@inline _homog_ap(::ArrayPartition{T, S}) where {T, S} = 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 diff --git a/test/Core/partitions_test.jl b/test/Core/partitions_test.jl index c5b4debf..a73da28d 100644 --- a/test/Core/partitions_test.jl +++ b/test/Core/partitions_test.jl @@ -427,3 +427,41 @@ struct TestIsnanFunctor end 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 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, ), From 91f03875de81d70495a8d42eecc2dee617781731 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Thu, 8 Oct 2026 11:49:14 -0400 Subject: [PATCH 3/3] Fix BoundsError in _homog_ap for a 0-partition ArrayPartition _homog_ap(::ArrayPartition{T,S}) called fieldtype(S, 1) without guarding fieldcount(S) == 0, so an in-place broadcast on ArrayPartition() threw BoundsError(Tuple{}, 1) instead of succeeding as it does on master (the old all(...) check was vacuously true on an empty iterator). Guard the empty case so it still constant-folds to true. Also trims a changelog-style comment per review feedback. Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Agent-Harness: Claude Code 2.1.294 Agent-Model: claude-sonnet-5 Agent-Session: https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv --- src/array_partition.jl | 6 ++---- test/Core/partitions_test.jl | 15 +++++++++++++++ 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/src/array_partition.jl b/src/array_partition.jl index dbfe9f32..eb1f5137 100644 --- a/src/array_partition.jl +++ b/src/array_partition.jl @@ -506,10 +506,8 @@ end 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 -# (the previous `all(x isa typeof(first(dest.x)) for x in dest.x)` ran, and allocated, at -# runtime, and only looked at `dest`, ignoring the broadcast arguments). -@inline _homog_ap(::ArrayPartition{T, S}) where {T, S} = S <: NTuple{fieldcount(S), fieldtype(S, 1)} +# 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)) diff --git a/test/Core/partitions_test.jl b/test/Core/partitions_test.jl index a73da28d..25c3df32 100644 --- a/test/Core/partitions_test.jl +++ b/test/Core/partitions_test.jl @@ -465,3 +465,18 @@ end 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