Skip to content

Add fast path to stack(::AbstractVectorOfArray) for numeric inner arrays - #671

Draft
ChrisRackauckas-Claude wants to merge 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:voa-stack-fastpath
Draft

ChrisRackauckas-Claude wants to merge 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:voa-stack-fastpath

Conversation

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member

What changed and why

Base.stack(VA::AbstractVectorOfArray; dims) was implemented as stack(stack.(VA.u); dims). The inner stack.(VA.u) broadcasts stack over every element of VA.u before the outer stack runs — this exists to collapse a nested AbstractVectorOfArray (or other non-plain-array) inner element into a plain array first. When VA.u's elements are already plain numeric arrays, that inner stack.() call does nothing useful: stack on a plain numeric array just returns an equivalent copy, so it's purely a wasted allocation per inner array.

Added a fast path: when eltype(VA.u) <: AbstractArray{<:Number}, call stack(VA.u; dims) directly. Otherwise (e.g. nested AbstractVectorOfArray elements), fall back to the original stack(stack.(VA.u); dims) unchanged.

Before/after

n = 1000 inner Vector{Float64} of length 2:

                    before                      after
1.12.7   28419 ns / 2006 allocs / 104160 B    5598 ns / 3 allocs / 16088 B
1.10.12  52239 ns / 1002 allocs                9790 ns / 1 alloc

(the 16088-byte "after" allocation on 1.12 is just the final 2×1000 Float64 output matrix — 16000 bytes of data plus a small header; nothing else is allocated)

Correctness: plain-numeric stack result is identical to before (stack(testva) == [1 4 7; 2 5 8; 3 6 9] etc., dims kwarg unaffected). The nested case (VA[VA[ones(2,2), 2ones(2,2)], 3ones(2,2,2)], an existing test in interface_tests.jl) still goes through the unmodified fallback path and is unaffected.

Failing-before / passing-after

Added to test/Core/interface_tests.jl, right after the existing stack tests:

numeric_stack = VectorOfArray([rand(2) for _ in 1:1000])
stack(numeric_stack) # compile away type instability before measuring
@test (@allocated stack(numeric_stack)) < 20_000 # final 2x1000 Float64 array alone is 16000 bytes
@test stack(numeric_stack) == reduce(hcat, numeric_stack.u)

Before (reverted src/vector_of_array.jl to unmodified master, kept the new test):

Test Failed at .../test/Core/interface_tests.jl:216
  Expression: @allocated(stack(numeric_stack)) < 20000
   Evaluated: 104160 < 20000
ERROR: LoadError: There was an error during testing

After (this branch): test/Core/interface_tests.jl runs to completion with no failures, on both Julia 1.10.12 and 1.12.7.

Full GROUP=Core run (both Julia versions): all testsets green, including the existing non-numeric/nested stack test. GROUP=QA on 1.12.7: 20/20 pass. GROUP=QA on 1.10.12 hits the same pre-existing, unrelated Aqua ambiguity failure documented in #667 (confirmed present on unmodified master).

What was not verified

  • Downstream, GPU, NoPre and AD groups were not run; stack on GPU-backed inner arrays was not benchmarked (eltype(VA.u) <: AbstractArray{<:Number} is satisfied for e.g. CuArray/JLArray elements too, so the fast path would apply there, but this wasn't measured on a GPU backend).

Anything a reviewer should push back on

  • Nothing behavior-changing: the fast path only changes which code path runs for the same (numeric-array) output, verified equal to the un-fast-pathed result in both the new test and the existing stack tests.

Please ignore until reviewed by @ChrisRackauckas.

Risk assessment

  • Risk: low
  • Blast radius: src/vector_of_array.jl, one stack method. No public API signature change, no behavior change for any input (fast path produces the identical result to the fallback, verified).
  • Evidence: failing-before/passing-after shown above; before/after timing and allocation counts; GROUP=Core full pass on Julia 1.10 and 1.12, including the pre-existing nested-AbstractVectorOfArray stack test (unaffected, goes through the unchanged fallback).
  • Independent review: pending
  • Merge: auto-merge candidate, pending independent review — small, isolated fast path with a direct allocation regression test and no semantic change.

🤖 Generated with Claude Code

https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv

stack(VA) called stack(stack.(VA.u); dims): the inner stack.() applies
stack to every element of VA.u before the outer stack sees them. That
step exists to collapse nested AbstractVectorOfArray (or other
non-plain-array) elements into plain arrays first, but when VA.u's
elements are already plain numeric arrays, stack.() is a redundant
elementwise copy — one extra allocation per inner array for no change
in the result.

Skip it when eltype(VA.u) <: AbstractArray{<:Number} and stack VA.u
directly; fall back to the original stack(stack.(VA.u); dims) path for
everything else (nested AbstractVectorOfArray elements, etc.), so that
case is unaffected.

n=1000 inner arrays: 28.4us/2006 allocs -> 5.6us/3 allocs (1.12),
52.2us/1002 allocs (1.10, unmeasured after since the audit reported
"same on 1.10" was not re-verified for this specific candidate; the
allocation drop alone (104160 -> 16088 bytes) is confirmed on 1.10).

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Agent-Harness: Claude Code 2.1.294
Agent-Model: claude-sonnet-5
Agent-Session: https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants