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

[ChrisRackauckas-Claude]

Changes since round 1 review

A Cursor review (round 1) found a blocking bug: _homog_ap(::ArrayPartition{T,S})
called fieldtype(S, 1) without guarding fieldcount(S) == 0, so an in-place
broadcast on the 0-partition ArrayPartition() threw BoundsError(Tuple{}, 1)
where master succeeds. Fixed by guarding the empty case
(fieldcount(S) == 0 || S <: NTuple{...}), confirmed it still constant-folds
to return true for the empty case via code_typed, added a Core test
("Broadcast on 0-partition ArrayPartition") covering in-place broadcast,
copy, OOP broadcast, and copyto!(::ArrayPartition, ::ArrayPartition) on
ArrayPartition(), and confirmed it fails on the pre-fix code
(BoundsError) and passes after. Also corrected the AP2 row below, which
had a copy-paste error (both columns showed the same number; the real
"before" is PR1-only ~92.5 ns, "after" is ~41.0 ns on this run — not
"unchanged" as originally written, per the review's non-blocking finding),
and trimmed a changelog-style comment on _homog_ap to a plain "why" note.
Numbers, bitwise identity, and the earlier regression test are unchanged and
re-verified below. Re-ran Runic, typos, Core and QA on Julia 1.12.4 and
1.10.12: QA on 1.10.12 fails with 1 Aqua method-ambiguity
(similar(::AbstractVectorOfArray, ...) vs. Base.similar), confirmed
pre-existing on unmodified origin/master under Julia 1.10 (same failure,
same ambiguity, via Test.detect_ambiguities on a clean worktree) and
unrelated to vector_of_array.jl, which this PR does not touch; 1.12.4 QA
is unaffected (20 Pass) since this ambiguity apparently doesn't trigger
under 1.12's similar signature.

Stacked on #666

This PR's branch (perf/ap-broadcast-instantiate-homog) is branched off
perf/ap-unpack-args-inline, the branch behind
#666, so the diff below
includes both commits: #666's @inline fix and this PR's own commit. GitHub
does not let a cross-fork PR's base be a branch that only exists on the fork,
so this is opened against master rather than against #666's branch directly
-- please review only the second commit (Instantiate per-partition Broadcasted...); the first is #666.

This PR needs #666: applying only this PR's change (instantiate + type-level
homogeneity check) without the @inline fix makes the mixed-partition-eltype
case strictly worse than master, because Broadcast.instantiate adds work
on a path that still runtime-dispatches. Measured directly (Julia 1.12.4, 3000
samples), instantiate-only vs. master:

case	master	instantiate-only (no PR1)
AP3mixed b+c*d	5203 ns / 21 allocs	9390 ns / 42 allocs
AP5mixed b+c*d	8683 ns / 35 allocs	15380 ns / 70 allocs

With PR1 applied first (this branch), the same cases improve (see below).

What and why

unpack(bc, i) drops a Broadcasted's axes "because it is easier to
recompute" (src/array_partition.jl:558 comment), but the per-partition
copyto!/copy in ArrayPartition's own broadcast machinery
(src/array_partition.jl:503,516,521 before this change) never recomputed
them — it handed each partition a Broadcasted{Style, Nothing, ...}. On
Julia 1.12 that forces the partition's own broadcast implementation onto a
slower, non-SIMD path. This PR wraps each per-partition unpack in
Broadcast.instantiate before copy/copyto!.

It also replaces the runtime homogeneity check
all(x isa typeof(first(dest.x)) for x in dest.x) — which only inspects
dest (ignoring the broadcast's other ArrayPartition arguments), and
allocates an iterator on every call — with a type-level check
(S <: NTuple{fieldcount(S), fieldtype(S,1)}) over dest and every
ArrayPartition argument of the Broadcasted tree, so it constant-folds
instead of running at runtime.

Before/after (reused from the audit's scratch/a1_bcfix2.jl and scratch/a1_ode.jl)

Julia 1.12.4, @benchmark 3000 samples, in-place broadcast on this branch
(PR1 + PR2) vs. PR1 alone:

case	before (PR1 only)	after (PR1+PR2)
AP2 (homog)	92.5 ns / 0 allocs	41.0 ns / 0 allocs
AP2big (homog, 1000+2000) b+c*d	6940.0 ns / 0	1264.0 ns / 0
AP2big 4arg	9700.0 ns / 0	1272.0 ns / 0
AP3mixed b+c*d	101.3 ns / 0	58.4 ns / 0
AP5mixed b+c*d	198.2 ns / 0	116.7 ns / 0

Julia 1.10.12, same workload:

case	before (PR1 only)	after (PR1+PR2)
AP2big b+c*d	1249.9 ns	1154.0 ns
AP3mixed b+c*d	57.2 ns	48.8 ns

Tsit5 ArrayPartition IIP ODE solve (scratch/a1_ode.jl, OrdinaryDiffEqTsit5),
Julia 1.12.4:

workload	before (PR1 only)	after (PR1+PR2)
AP F64/F64 step! (warmed)	5452.0 ns / 80 B	2533.0 ns / 80 B
AP F64/F64/F32 step! (warmed)	10304.0 ns / 1528 B	3886.9 ns / 64 B
AP F64/F64 IIP solve	178.12 us	108.31 us
AP F64/F64/F32 IIP solve	308.03 us	145.28 us

Julia 1.10.12:

workload	before (PR1 only)	after (PR1+PR2)
AP F64/F64 step!	2497.0 ns / 192 B	2378.0 ns / 192 B
AP F64/F64/F32 step!	7595.0 ns / 792 B	4656.0 ns / 64 B

Results are bitwise identical before/after. Checked with a seeded script
(correctness2.jl) comparing collect(a) for every broadcast case above and
(sol.t, collect(sol.u[end])) for both Tsit5 ODE workloads, serialized from
the unpatched tree and compared against this branch:

AP2 f! identical: true
AP2 g! identical: true
AP3mixed f! identical: true
AP3mixed g! identical: true
AP5mixed f! identical: true
AP5mixed g! identical: true
AP2big f! identical: true
AP2big g! identical: true
AP F64/F64 ode identical: true
AP F64/F64/F32 ode identical: true
ALL IDENTICAL

Failing-before / passing-after

Neither @allocated nor JET.@report_opt discriminates this fix: allocations
and runtime-dispatch-report counts for the homogeneous large-ArrayPartition
case are 0 both before and after (verified directly) — the regression this PR
fixes is a lost-SIMD timing effect, not an allocation or dispatch effect, so
those techniques used for PR1 don't apply here.

Instead, added a deterministic structural test,
test/Core/partitions_test.jl::"ArrayPartition broadcast instantiates partition Broadcasted", using a custom partition array type that records whether the
Broadcasted it receives in its own copyto! has axes !== nothing (i.e.
was instantiated). This directly pins the mechanism (not a timing threshold,
so it isn't flaky in CI):

Before this fix (git stash the src/array_partition.jl change):

instantiated: false

After:

instantiated: true

0-partition ArrayPartition (the bug found in round-1 review): added
test/Core/partitions_test.jl::"Broadcast on 0-partition ArrayPartition",
covering in-place broadcast, copy, OOP broadcast, and
copyto!(::ArrayPartition, ::ArrayPartition) on ArrayPartition().

Before this fix (git stash the _homog_ap one-line change):

ERROR: BoundsError(Tuple{}, 1)

After:

(all four pass; dest == ArrayPartition())

Also confirmed with code_typed that _homog_ap(ArrayPartition()) still
constant-folds to return true (not a runtime branch), and spot-checked
_homog_ap on a 1-partition and a mixed-eltype ArrayPartition (true and
false respectively, as expected).

Test results

GROUP=Core: all 12 Core testsets pass (Partitions Tests: 135 Pass, 0 Fail,
includes this PR's two tests and PR1's).
GROUP=QA: QA: 20 Pass, 0 Fail — required adding :instantiate to
test/QA/qa.jl's all_qualified_accesses_are_public ignore list next to the
other legitimately-qualified Base.Broadcast names already there
(:Broadcasted, :flatten, :result_style, etc.), since Broadcast.instantiate
is not itself declared public by Base.

Both re-run on Julia 1.12.4 and 1.10.12.

What was not verified

• Downstream, GPU, AD, and NoPre (JET) test groups were not run.
• AllocCheck.jl-level static allocation checking was not used; @allocated
  and JET.@report_opt were, and neither showed a difference (expected, see
  above).
• copy (vs. copyto!) on an ArrayPartition broadcast with mixed partition
  types was exercised by the existing test suite and the correctness script,
  but no dedicated microbenchmark of mixed-type copy was run (only the
  OOP rows in the in-place-broadcast benchmark table above, which do exercise
  copy).

What a reviewer should push back on

The 0-partition case is now covered (see above); worth independently
re-checking that no other empty-ArrayPartition path regressed beyond the
four covered by the new test (copyto!(dest::AbstractArray, A::ArrayPartition)
and getindex/setindex! on an empty ArrayPartition don't go through
_homog_ap at all, so they were never at risk, but a second pair of eyes on
that claim is welcome).

Please ignore until reviewed by @ChrisRackauckas.

Risk assessment

• Risk: low
• Blast radius: src/array_partition.jl, Base.copy/Base.copyto! for
  ArrayPartition broadcasts and the internal homogeneity check; no signature
  or public API change. Affects every ArrayPartition broadcast (homogeneous
  and mixed).
• Evidence: before/after timing on Julia 1.12.4 and 1.10.12; bitwise-identical
  correctness check across 10 workloads including a Tsit5 ODE solve; Core and
  QA test groups pass; a deterministic structural regression test.
• Independent review: round 1 (Cursor) returned CHANGES for the 0-partition BoundsError above, now fixed; round 2 pending
• Merge: needs human review — depends on Add @inline to unpack_args(i, ::Tuple{Any}) in ArrayPartition broadcast #666 and touches the core
  ArrayPartition broadcast dispatch, so it should not auto-merge even if CI
  is green.

🤖 Generated with Claude Code

https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv

[ChrisRackauckas]

🤖 Automated comment from an AI agent running as @ChrisRackauckas — not written or reviewed by Chris.

Independent review (Devin CLI 3000.11.3, model fusion-claude-opus-5-5-high-sidekick-swe-2-medium): MERGE, risk medium.

Full review

VERDICT: MERGE
RISK: medium

Every ArrayPartition broadcast (copy and copyto!) goes through the changed code, and OrdinaryDiffEq uses it on every step. No public API changes, and the probes below gave identical results before and after. Risk is medium rather than low because of how widely the hot path is used, and because the GPU, Downstream and AD groups were not run.

Blocking findings

None.

Non-blocking findings

1. The branch still contains the Add @inline to unpack_args(i, ::Tuple{Any}) in ArrayPartition broadcast #666 commit, which is now merged (confirmed via gh). Add @inline to unpack_args(i, ::Tuple{Any}) in ArrayPartition broadcast #666 was squash-merged on 2026-10-10 as 1826933. This branch still carries the original 726bfc3, so its diff against my stale local origin/master shows the @inline hunk and the "Mixed-eltype ... does not allocate" testset as well. The hunks are identical, so I expect a clean merge, but I did not check mergeability against the new master. Rebase onto master so the PR contains only 31c7293 and 91f0387.
2. A test comment describes history (from reading), test/Core/partitions_test.jl:470-471: "_homog_ap(::ArrayPartition{T,S}) called fieldtype(S, 1) without guarding fieldcount(S) == 0, throwing BoundsError...". This narrates a bug in an intermediate commit of this same PR, which the CLAUDE.md "describe what exists" rule forbids. Delete it, or replace it with a present-tense line such as "0-partition ArrayPartitions take the homogeneous path".
3. The 0-partition testset only guards against this PR's own intermediate bug (confirmed by running). It passes on the pre-PR baseline (4/4, log_pre_newtests_112.txt). On 31c7293's src it fails: y = ArrayPartition(); y .= ArrayPartition() .+ 1 gave BoundsError: attempt to access DataType at index [1]. So the PR body's failing-before claim holds for that commit, not for master. That is fine as a regression test; I'm noting it so nobody reads it as a fix for a bug on master.
4. NamedArrayPartition is not updated (from reading; not benchmarked). src/named_array_partition.jl:207-229 has the same copy(unpack(bc, i)) / copyto!(..., unpack(bc, i)) pattern without Broadcast.instantiate. If the 1.12 SIMD loss is real for ArrayPartition, NAP broadcasts presumably have it too. This could be a follow-up PR; it doesn't need to go in this one.
5. First calls of some mixed broadcasts compile more slowly (confirmed by running). Case: a homogeneous dest with a heterogeneous ArrayPartition argument. Under the old dest-only check this used the for loop; it now unrolls through ntuple(f, Val(N)).
   • N=40 partitions: first call went from 1.03 s to 1.65 s.
   • Steady-state allocations dropped from 22096 B to 6160 B (and from 320 B to 0 for a small 2-partition case).
   • This is the right trade, but it should be mentioned for users with many partitions.
6. Broadcast.instantiate is not public in Base (confirmed with Base.ispublic(Base.Broadcast, :instantiate) == false on 1.12). It is the documented extension point in the manual's broadcast interface section, and test/QA/qa.jl already ignores sibling names that are just as non-public (Broadcasted, result_style, flatten). Adding it to the ignore list is acceptable.
7. instantiate uses the arguments' axes, not axes(dest.x[i]) (from reading, behaviour confirmed by the probes). Base's materialize! uses the destination's axes. Because of this, the in-place cases "extrude len1", "vec into matrix part" and "transpose arg" still throw DimensionMismatch, exactly as before the PR. Behaviour is unchanged, so this is not a regression, but switching to dest axes would be a separate semantic change if it is ever wanted.

What I ran

All runs used TMPDIR set to the review's tmp/. The scratch trees are git archive extracts:

• scratch/head = PR head.
• scratch/pre = 726bfc3, i.e. master plus Add @inline to unpack_args(i, ::Tuple{Any}) in ArrayPartition broadcast #666, the PR's baseline.
• scratch/mid = src from 31c7293.
• scratch/pre_newtests = baseline src with the PR's tests.

Run	Result
GROUP=Core Pkg.test on head, Julia 1.12.7	All 12 Core testsets pass, Partitions 135/135 (scratch/log_head_core_112.txt)
GROUP=Core Pkg.test on head, Julia 1.10.12	All pass, Partitions 135/135 (scratch/log_head_core_110.txt)
GROUP=QA on head, 1.12	20/20 (scratch/log_head_qa_112.txt)
Baseline src + PR tests, 1.12	Partitions 133 Pass / 2 Fail. Both failures are in "ArrayPartition broadcast instantiates partition Broadcasted" (lines 463 and 466), so that test does fail without the fix. The 0-partition test passes there (see finding 3). Log: scratch/log_pre_newtests_112.txt
scratch/mid (31c7293) empty in-place broadcast	BoundsError, confirming the round-1 bug and that 91f0387 fixes it
Runic --check --diff on the 3 changed files	Clean
typos	Clean

Differential probe (scratch/probe.jl, run on pre and head, 1.12; outputs in scratch/probe_{pre,head}.txt). The printed results and the error types and messages are identical for every case. The only lines that differ are timing and allocation counts. Cases covered:

• mixed-dimension partitions (Vector + Matrix), OOP and in-place;
• scalar and 0-d partitions;
• SVector, SVector+Vector, MVector and MVector+Vector partitions;
• nested ArrayPartitions with mixed eltypes;
• ForwardDiff Dual partitions;
• a homogeneous dest with a heterogeneous argument;
• shape mismatch, length-1 extrusion and partition-count mismatch errors;
• a Ref argument, range partitions, Bool partitions, view partitions and an adjoint argument;
• abstract S (ArrayPartition{Float64,Tuple{AbstractVector{Float64},AbstractVector{Float64}}});
• empty and 1-partition ArrayPartitions;
• NamedArrayPartition OOP and in-place, and an AP dest with a NAP argument.

Mixed-eltype in-place broadcasts still allocate 0 bytes, and OOP return types still infer concretely.

Benchmarks (scratch/bench.jl, BenchmarkTools, 3000 samples, minimum time in ns, pre → head; 1.12 was run twice and the runs matched):

Case	1.12	1.10.12
AP2 f!	45.8 → 25.0	20.8 → 20.8
AP2 g!	125 → 25	not listed
AP2big f!	3171 → 712	729 → 733
AP2big g!	11412 → 946	not listed
AP2big oop	3362 → 946	not listed
AP3mixed f!	54 → 25	not listed
AP5mixed f!	79 → 46	not listed

• On 1.10 every case is unchanged or slightly faster; OOP is 5–13% faster.
• The 1.12 speedups go the same way as the PR's tables. My absolute numbers differ from the author's because this is a different machine.
• So the PR's performance claim reproduces, and there is no regression on 1.10.

What I did not verify

• The GPU, Downstream (OrdinaryDiffEq), AD and NoPre/JET test groups. In particular, I did not exercise instantiate under CUDA/KernelAbstractions broadcast styles.
• The PR's Tsit5 ODE timings and its "bitwise identical" correctness2.jl output. I did not rerun the author's scripts, because they are not in the repo. My own differential probe covers bitwise equality of broadcast results.
• Mergeability against the current master after the Add @inline to unpack_args(i, ::Tuple{Any}) in ArrayPartition broadcast #666 squash-merge (finding 1).
• Whether NamedArrayPartition broadcasts actually show the same slowdown on 1.12 (finding 4 is inferred from reading the code).
• The PR links no issue, and there is no tracking issue, so I did not check for interface conflicts with other open PRs beyond Add @inline to unpack_args(i, ::Tuple{Any}) in ArrayPartition broadcast #666.

---

🤖 Posted by an AI agent — harness: Devin CLI 3000.11.3 (review), Claude Code 2.1.285 (posting) · model: fusion-claude-opus-5-5-high-sidekick-swe-2-medium
Conversation: local Claude session 3c311569-dda6-4687-bc9f-65febb212c33 (fleet master)