ArenaNode for array-backed expressions - #167
Open
MilesCranmerBot wants to merge 67 commits into
Open
Conversation
Contributor
Benchmark Results (Julia v1)Time benchmarks
Memory benchmarks
|
MilesCranmerBot
force-pushed
the
arenanode-prototype
branch
from
February 21, 2026 12:56
e8acc61 to
847c2c6
Compare
MilesCranmerBot
force-pushed
the
arenanode-prototype
branch
from
June 9, 2026 09:23
847c2c6 to
04e763d
Compare
Member
|
@MilesCranmerBot comments from Claude... pls integrate. Maybe benchmark against master too, to see if it works... |
Contributor
Author
|
Integrated in What changed:
Validation: Benchmark ( |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #167 +/- ##
===========================================
+ Coverage 59.43% 74.02% +14.59%
===========================================
Files 30 32 +2
Lines 2682 3188 +506
===========================================
+ Hits 1594 2360 +766
+ Misses 1088 828 -260 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
MilesCranmer
left a comment
MilesCranmer
left a comment
Member
There was a problem hiding this comment.
@MilesCranmerBot please address above comments.
Also, shouldn't we have a tree_mapreduce? Then we wouldn't need to write a custom is_constant and so forth.
Btw, please expand your test suite. You should actually be using the full Node Interface tests. ArenaNode should pass ALL of them.
Contributor
Author
|
Addressed the review in the latest pushes ( Changes:
Local validation:
All addressed review threads are resolved. CI is pending on the latest push. |
MilesCranmer
requested changes
Jun 10, 2026
MilesCranmer
left a comment
MilesCranmer
left a comment
Member
There was a problem hiding this comment.
See new comments.
Also.
Please test with expressions types too. eg Expression with ArenaNode - do the full expression interface tests.
Also maybe try derivatives too?
MilesCranmer
requested changes
Jun 11, 2026
MilesCranmer
left a comment
MilesCranmer
left a comment
Member
There was a problem hiding this comment.
@MilesCranmerBot please see comments above
Also please resolve any comments that have already been addressed.
Contributor
Author
|
Addressed in Changes:
Validation: Resolved the review threads addressed by this push. CI is pending on the latest commit. |
Member
|
@MilesCranmerBot there are still unresolved comments left. Also the CI is failing. Please address everything. |
Contributor
Author
|
Addressed in Changes:
Validation:
|
MilesCranmer
requested changes
Jun 11, 2026
MilesCranmer
left a comment
MilesCranmer
left a comment
Member
There was a problem hiding this comment.
@MilesCranmerBot please address
Contributor
Author
|
Addressed in Changes:
Validation:
Resolved the review threads addressed by this push. CI is pending on the new commit. |
Member
|
@MilesCranmerBot downgrade compat is failing. Could you make a separate PR that fixes the lower bound? Or, if necessary, in this one (if there's a new dependency for example). Also why does it say only 0.54% of diff is actually tested... That is crazy. Should be 100%... |
Member
|
Also please resolve individual comments if already addressed. |
MilesCranmerBot
force-pushed
the
arenanode-prototype
branch
from
June 11, 2026 15:40
d401e65 to
92d4eef
Compare
Contributor
Author
|
Done.
Root cause: downgrade compat could select On the 0.54% Codecov number: that report is stale. The Codecov comment compares base |
Contributor
Author
|
Follow-up correction: the first pushed bound was too narrow for the Julia 1.10 downgrade path. I pushed LoweredCodeUtils = "2, 3.4 - 3.5.99"That keeps Julia 1.10 on the JET 0.9 / LoweredCodeUtils 2.x path, while excluding the failing LoweredCodeUtils 3.0.x path on Julia 1.12. Local fresh resolver/import checks:
|
Member
|
@MilesCranmerBot CI is still failing, pls fix |
Member
|
@MilesCranmerBot no excluding coverage files... You should try to actually improve the coverage as needed. |
MilesCranmer
left a comment
MilesCranmer
left a comment
Member
There was a problem hiding this comment.
@MilesCranmerBot please address comments.
Contributor
Author
|
Fixed and pushed to the PR branch. Changes:
Verified locally with the full package test command plus targeted SymbolicUtils/Optim/LoopVectorization smoke tests. Remote checks are now green: CI matrix, downgrade compat, format, bench, and Codecov patch/project all pass. |
MilesCranmer
left a comment
MilesCranmer
left a comment
Member
There was a problem hiding this comment.
@MilesCranmerBot pls address comments
MilesCranmer
changed the title
MilesCranmer
changed the title
Member
|
@MilesCranmerBot can you please apply the following diff? Details 5 files changed, 367 insertions(+), 51 deletions(-)
diff --git a/src/ArenaNode.jl b/src/ArenaNode.jl
index 07e606e..faeafc4 100644
--- a/src/ArenaNode.jl
+++ b/src/ArenaNode.jl
@@ -9,32 +9,68 @@ import ..NodeModule:
unsafe_get_children,
get_child,
set_child!,
- set_children!
-
-"""Array-backed arena storing the fields of a tree node in a struct-of-arrays form.
+ set_children!,
+ count_nodes,
+ copy_node
+import ..NodeUtilsModule:
+ count_constant_nodes,
+ count_scalar_constants,
+ has_constants,
+ get_scalar_constants,
+ set_scalar_constants!
+import ..NodePreallocationModule: allocate_container, copy_into!
+import ..ValueInterfaceModule: get_number_type
+
+"""All per-node fields packed into a single isbits struct.
+
+Storing nodes as one `Vector{ArenaEntry}` (array-of-structs) makes whole-tree
+operations flat array operations: `copy` is a single `memcpy`, and traversals
+touch one contiguous stream of memory.
Indices are `Int32` and are 1-based. A child index of `0` indicates an empty slot.
+"""
+struct ArenaEntry{T,D}
+ val::T
+ children::NTuple{D,Int32}
+ feature::UInt16
+ degree::UInt8
+ op::UInt8
+ constant::Bool
+end
+
+@inline function _replace(
+ e::ArenaEntry{T,D};
+ val=e.val,
+ children=e.children,
+ feature=e.feature,
+ degree=e.degree,
+ op=e.op,
+ constant=e.constant,
+) where {T,D}
+ return ArenaEntry{T,D}(val, children, feature, degree, op, constant)
+end
+
+"""Array-backed arena storing the nodes of a tree contiguously.
This is an *experimental prototype* intended to provide an arena-backed representation
with a `Node`-like facade (`ArenaNode`) that supports existing tree algorithms that are
written against `AbstractExpressionNode`.
+
+The `compact` flag tracks whether `nodes` is exactly one postfix-ordered tree
+(children stored before parents, root last, no orphaned nodes and no shared
+subtrees). Trees built via `convert`/`copy` are compact; structural mutations
+through the facade may clear the flag, in which case whole-tree operations fall
+back to generic traversals. A structural `copy` re-compacts.
"""
struct Arena{T,D}
- degree::Vector{UInt8}
- constant::Vector{Bool}
- val::Vector{T}
- feature::Vector{UInt16}
- op::Vector{UInt8}
- children::Vector{NTuple{D,Int32}}
+ nodes::Vector{ArenaEntry{T,D}}
+ compact::Base.RefValue{Bool}
function Arena{T,D}(; capacity::Integer=0) where {T,D}
- degree = sizehint!(UInt8[], capacity)
- constant = sizehint!(Bool[], capacity)
- val = sizehint!(T[], capacity)
- feature = sizehint!(UInt16[], capacity)
- op = sizehint!(UInt8[], capacity)
- children = sizehint!(NTuple{D,Int32}[], capacity)
- return new{T,D}(degree, constant, val, feature, op, children)
+ return new{T,D}(sizehint!(ArenaEntry{T,D}[], capacity), Ref(true))
+ end
+ function Arena{T,D}(nodes::Vector{ArenaEntry{T,D}}, compact::Bool) where {T,D}
+ return new{T,D}(nodes, Ref(compact))
end
end
@@ -54,6 +90,13 @@ end
@inline ArenaNode(arena::Arena{T,D}, idx::Int32) where {T,D} = ArenaNode{T,D}(arena, idx)
+"""Whether `tree` is the root of a compact arena, so that the arena contents
+*are* the tree and whole-tree operations can act on the flat array directly."""
+@inline function is_compact_root(tree::ArenaNode)
+ a = getfield(tree, :arena)
+ return a.compact[] && Int(getfield(tree, :idx)) == length(a.nodes)
+end
+
@inline function _zero_children(::Val{D}) where {D}
return ntuple(_ -> Int32(0), Val(D))
end
@@ -67,13 +110,8 @@ end
op::UInt8,
children::NTuple{D,Int32},
) where {T,D}
- push!(arena.degree, degree)
- push!(arena.constant, constant)
- push!(arena.val, val)
- push!(arena.feature, feature)
- push!(arena.op, op)
- push!(arena.children, children)
- return Int32(length(arena.degree))
+ push!(arena.nodes, ArenaEntry{T,D}(val, children, feature, degree, op, constant))
+ return Int32(length(arena.nodes))
end
@inline function push_constant!(arena::Arena{T,D}, value) where {T,D}
@@ -117,15 +155,15 @@ Base.@constprop :aggressive @inline function Base.getproperty(
elseif k === :idx
return getfield(n, :idx)
elseif k === :degree
- return @inbounds getfield(n, :arena).degree[getfield(n, :idx)]
+ return @inbounds getfield(n, :arena).nodes[getfield(n, :idx)].degree
elseif k === :constant
- return @inbounds getfield(n, :arena).constant[getfield(n, :idx)]
+ return @inbounds getfield(n, :arena).nodes[getfield(n, :idx)].constant
elseif k === :val
- return @inbounds getfield(n, :arena).val[getfield(n, :idx)]::T
+ return @inbounds getfield(n, :arena).nodes[getfield(n, :idx)].val::T
elseif k === :feature
- return @inbounds getfield(n, :arena).feature[getfield(n, :idx)]
+ return @inbounds getfield(n, :arena).nodes[getfield(n, :idx)].feature
elseif k === :op
- return @inbounds getfield(n, :arena).op[getfield(n, :idx)]
+ return @inbounds getfield(n, :arena).nodes[getfield(n, :idx)].op
elseif k === :children
return unsafe_get_children(n)
elseif k === :l
@@ -138,21 +176,26 @@ Base.@constprop :aggressive @inline function Base.getproperty(
end
@inline function Base.setproperty!(n::ArenaNode{T,D}, k::Symbol, v) where {T,D}
+ a = n.arena
i = n.idx
+ e = @inbounds a.nodes[i]
if k === :degree
- @inbounds n.arena.degree[i] = UInt8(v)
+ # Changing arity orphans or exposes child slots, so the flat layout can
+ # no longer be assumed to be exactly this tree.
+ UInt8(v) == e.degree || (a.compact[] = false)
+ @inbounds a.nodes[i] = _replace(e; degree=UInt8(v))
return v
elseif k === :constant
- @inbounds n.arena.constant[i] = Bool(v)
+ @inbounds a.nodes[i] = _replace(e; constant=Bool(v))
return v
elseif k === :val
- @inbounds n.arena.val[i] = convert(T, v)
+ @inbounds a.nodes[i] = _replace(e; val=convert(T, v))
return v
elseif k === :feature
- @inbounds n.arena.feature[i] = UInt16(v)
+ @inbounds a.nodes[i] = _replace(e; feature=UInt16(v))
return v
elseif k === :op
- @inbounds n.arena.op[i] = UInt8(v)
+ @inbounds a.nodes[i] = _replace(e; op=UInt8(v))
return v
elseif k === :l
set_child!(n, v, 1)
@@ -179,13 +222,13 @@ accessing them throws an `UndefRefError`.
"""
@generated function unsafe_get_children(n::ArenaNode{T,D}) where {T,D}
quote
- children = @inbounds getfield(n, :arena).children[getfield(n, :idx)]
+ children = @inbounds getfield(n, :arena).nodes[getfield(n, :idx)].children
return Base.Cartesian.@ntuple($D, j -> _nullable_child(n, children[j]))
end
end
@inline function get_child(n::ArenaNode{T,D}, i::Int) where {T,D}
- c = @inbounds n.arena.children[n.idx][i]
+ c = @inbounds n.arena.nodes[n.idx].children[i]
c == 0 && throw(UndefRefError())
return ArenaNode(n.arena, c)
end
@@ -205,9 +248,14 @@ end
_copy_to_arena!(n.arena, child)
end
- old = @inbounds n.arena.children[n.idx]
- @inbounds n.arena.children[n.idx] = Base.setindex(old, idx, i)
- return ArenaNode(n.arena, idx)
+ a = n.arena
+ e = @inbounds a.nodes[n.idx]
+ if @inbounds(e.children[i]) != idx
+ # Relinking orphans the old child subtree (and may introduce sharing).
+ a.compact[] = false
+ @inbounds a.nodes[n.idx] = _replace(e; children=Base.setindex(e.children, idx, i))
+ end
+ return ArenaNode(a, idx)
end
@inline function set_children!(
@@ -236,17 +284,67 @@ end
idxs = Base.setindex(idxs, idx, i)
end
- @inbounds n.arena.children[n.idx] = idxs
+ a = n.arena
+ e = @inbounds a.nodes[n.idx]
+ if e.children != idxs
+ a.compact[] = false
+ @inbounds a.nodes[n.idx] = _replace(e; children=idxs)
+ end
return nothing
end
-"""Copy a tree into a new arena and return the new root node."""
-function Base.copy(tree::ArenaNode{T,D}; break_sharing::Val{BS}=Val(false)) where {T,D,BS}
+"""Copy a tree into a new arena and return the new root node.
+
+When `tree` is the root of a compact arena, this is a single flat copy of the
+node array (child indices are arena-relative, so they remain valid verbatim).
+Otherwise it falls back to a structural copy, which also re-compacts the
+resulting arena.
+
+This overloads `copy_node` (rather than `Base.copy`) since it is the generic
+entry point: `Base.copy(::AbstractExpressionNode)` forwards here, and the
+fallback `copy_node` would otherwise build a fresh arena per copied node via
+`constructorof`.
+"""
+function copy_node(tree::ArenaNode{T,D}; break_sharing::Val{BS}=Val(false)) where {T,D,BS}
+ if is_compact_root(tree)
+ return ArenaNode{T,D}(Arena{T,D}(copy(tree.arena.nodes), true), tree.idx)
+ end
arena = Arena{T,D}(; capacity=length(tree; break_sharing=Val(true)))
idx = _copy_to_arena!(arena, tree)
return ArenaNode{T,D}(arena, idx)
end
+"""Preallocate an arena for [`copy_into!`](@ref), enabling zero-allocation copies."""
+function allocate_container(
+ prototype::ArenaNode{T,D}, n::Union{Nothing,Integer}=nothing
+) where {T,D}
+ return Arena{T,D}(; capacity=@something(n, length(prototype)))
+end
+
+"""Copy `src` into the preallocated arena `dest`, reusing its storage.
+
+This is the steady-state copy path for population-based search: no allocations
+once `dest` has sufficient capacity.
+"""
+function copy_into!(
+ dest::Arena{T,D},
+ src::ArenaNode{T,D};
+ ref::Union{Nothing,Base.RefValue{<:Integer}}=nothing,
+) where {T,D}
+ @assert dest !== src.arena
+ if is_compact_root(src)
+ nodes = src.arena.nodes
+ resize!(dest.nodes, length(nodes))
+ copyto!(dest.nodes, nodes)
+ dest.compact[] = true
+ return ArenaNode{T,D}(dest, src.idx)
+ end
+ empty!(dest.nodes)
+ idx = _copy_to_arena!(dest, src)
+ dest.compact[] = true
+ return ArenaNode{T,D}(dest, idx)
+end
+
function _copy_to_arena!(arena::Arena{T,D}, tree::AbstractExpressionNode{T,D}) where {T,D}
d = tree.degree
if d == 0
@@ -266,7 +364,7 @@ end
"""Convert an existing tree into an arena-backed representation.
-This copies the entire tree into a fresh arena.
+This copies the entire tree into a fresh arena, in postfix (children-first) order.
"""
@inline function Base.convert(
::Type{ArenaNode{T,D}}, tree::AbstractExpressionNode{T,D}
@@ -281,6 +379,94 @@ end
return convert(ArenaNode{T,D}, tree)
end
+################################################################################
+# Flat whole-tree operations
+#
+# For a compact arena the node array *is* the tree, so tree-wide reductions
+# become linear array scans with no pointer chasing. Each of these falls back
+# to the generic traversal-based implementation when the invariant doesn't hold.
+################################################################################
+
+function count_nodes(tree::ArenaNode; break_sharing::Val{BS}=Val(false)) where {BS}
+ if is_compact_root(tree)
+ return length(tree.arena.nodes)
+ end
+ return invoke(count_nodes, Tuple{AbstractNode}, tree; break_sharing=Val(BS))::Int64
+end
+
+function count_constant_nodes(tree::ArenaNode)
+ if is_compact_root(tree)
+ return count(e -> e.degree == 0x00 && e.constant, tree.arena.nodes)
+ end
+ return invoke(count_constant_nodes, Tuple{AbstractExpressionNode}, tree)
+end
+
+function has_constants(tree::ArenaNode)
+ if is_compact_root(tree)
+ return any(e -> e.degree == 0x00 && e.constant, tree.arena.nodes)
+ end
+ return invoke(has_constants, Tuple{AbstractNode}, tree)
+end
+
+function count_scalar_constants(tree::ArenaNode{T}) where {T<:Number}
+ if is_compact_root(tree)
+ # For scalar `T`, each constant leaf stores exactly one scalar.
+ return count(e -> e.degree == 0x00 && e.constant, tree.arena.nodes)
+ end
+ return invoke(count_scalar_constants, Tuple{AbstractExpressionNode{T}}, tree)
+end
+
+"""Used by `NodeSampler` (random node selection in mutations), once per sample."""
+function Base.count(
+ f::F, tree::ArenaNode{T,D}; init=0, break_sharing::Val{BS}=Val(false)
+) where {F<:Function,T,D,BS}
+ if is_compact_root(tree)
+ a = tree.arena
+ c = init
+ @inbounds for i in 1:length(a.nodes)
+ c += f(ArenaNode{T,D}(a, Int32(i))) ? 1 : 0
+ end
+ return c
+ end
+ return invoke(Base.count, Tuple{F,AbstractNode}, f, tree; init, break_sharing=Val(BS))
+end
+
+"""For compact arenas, constants are gathered by a linear scan, and the
+returned `refs` are plain arena indices (which also remain valid in flat
+copies of the tree)."""
+function get_scalar_constants(
+ tree::ArenaNode{T}, ::Type{BT}=get_number_type(T)
+) where {T<:Number,BT}
+ if is_compact_root(tree)
+ nodes = tree.arena.nodes
+ n_constants = count(e -> e.degree == 0x00 && e.constant, nodes)
+ vals = Vector{T}(undef, n_constants)
+ refs = Vector{Int32}(undef, n_constants)
+ j = 0
+ @inbounds for i in eachindex(nodes)
+ e = nodes[i]
+ if e.degree == 0x00 && e.constant
+ j += 1
+ vals[j] = e.val
+ refs[j] = Int32(i)
+ end
+ end
+ return vals, refs
+ end
+ return invoke(get_scalar_constants, Tuple{AbstractExpressionNode{T},Type{BT}}, tree, BT)
+end
+
+function set_scalar_constants!(
+ tree::ArenaNode{T}, constants, refs::AbstractVector{Int32}
+) where {T<:Number}
+ nodes = tree.arena.nodes
+ @inbounds for j in eachindex(refs, constants)
+ i = refs[j]
+ nodes[i] = _replace(nodes[i]; val=constants[j]::T)
+ end
+ return nothing
+end
+
################################################################################
# Cursor + reusable stack (prototype)
################################################################################
@@ -325,11 +511,10 @@ function next!(c::ArenaCursor{T,D})::Nullable{ArenaNode{T,D}} where {T,D}
node = ArenaNode{T,D}(c.arena, idx)
# Push children in reverse order so the leftmost child is visited next.
- d = @inbounds c.arena.degree[idx]
- if d != 0
- child_idxs = @inbounds c.arena.children[idx]
- @inbounds for i in d:-1:1
- child = child_idxs[i]
+ e = @inbounds c.arena.nodes[idx]
+ if e.degree != 0
+ @inbounds for i in e.degree:-1:1
+ child = e.children[i]
child != 0 && push!(c.stack, child)
end
end
diff --git a/src/DynamicExpressions.jl b/src/DynamicExpressions.jl
index b1fa6c6..5e2688f 100644
--- a/src/DynamicExpressions.jl
+++ b/src/DynamicExpressions.jl
@@ -8,9 +8,9 @@ using DispatchDoctor: @stable, @unstable
include("ExtensionInterface.jl")
include("OperatorEnum.jl")
include("Node.jl")
- include("ArenaNode.jl")
include("NodeUtils.jl")
include("NodePreallocation.jl")
+ include("ArenaNode.jl")
include("Strings.jl")
include("Evaluate.jl")
include("EvaluateDerivative.jl")
diff --git a/test/runtests.jl b/test/runtests.jl
index fa16d57..7bad778 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -56,7 +56,7 @@ if "jet" in test_names
)
s_mod = string(mod.mod)
any(report.vst) do vst
- occursin(s_mod, string(JET.linfomod(vst.linfo)))
+ return occursin(s_mod, string(JET.linfomod(vst.linfo)))
end
end
# On JET 0.10, `target_defined_modules` is not available and also
@@ -82,6 +82,10 @@ testitem_suffixes = String[]
if "main" in test_names
push!(testitem_suffixes, joinpath("test", "unittest.jl"))
push!(testitem_suffixes, joinpath("test", "test_optim.jl"))
+ # NOTE: `@testitem`s defined in a file that `unittest.jl` merely `include`s are
+ # attributed to *their own* filename by TestItemRunner, so files with their own
+ # testitems must be listed here explicitly or they will be silently skipped.
+ push!(testitem_suffixes, joinpath("test", "test_arenanode.jl"))
end
if "optim" in test_names
push!(testitem_suffixes, joinpath("test", "test_optim.jl"))
diff --git a/test/test_arenanode.jl b/test/test_arenanode.jl
index 98f0fc7..342e436 100644
--- a/test/test_arenanode.jl
+++ b/test/test_arenanode.jl
@@ -118,7 +118,8 @@ end
@test string_tree(rewritten, operators) == "sin(x1 * 2.0)"
bad_children = (
- DynamicExpressions.Nullable{Node{Float64,2}}(true), Node{Float32}(; val=1.0f0)
+ DynamicExpressions.Nullable(true, Node{Float64}(; val=0.0)),
+ Node{Float32}(; val=1.0f0),
)
@test_throws ArgumentError set_children!(rewritten, bad_children)
@@ -203,7 +204,7 @@ end
@test grad3[1, :] ≈ fill(1.0, 5)
d_ex = gradient(AutoZygote(), expr_const) do ex
- sum(ex(ones(1, 5)))
+ return sum(ex(ones(1, 5)))
end
@test extract_gradient(d_ex, expr_const) ≈ [5.0]
end
@@ -232,3 +233,119 @@ end
end
end
end
+
+@testitem "ArenaNode flat copy and whole-tree fast paths" begin
+ using DynamicExpressions
+ using DynamicExpressions: Node, copy_node
+ using DynamicExpressions.NodePreallocationModule: allocate_container, copy_into!
+
+ const AN = DynamicExpressions.ArenaNodeModule
+
+ operators = OperatorEnum(; binary_operators=[+, -, *, /], unary_operators=[sin, cos])
+ x1 = Node{Float64}(; feature=1)
+ x2 = Node{Float64}(; feature=2)
+ tree = sin(x1 * 3.2 - 0.9) + x2 * (x1 - 0.5)
+ atree = convert(AN.ArenaNode{Float64}, tree)
+
+ @testset "compact flat copy" begin
+ @test AN.is_compact_root(atree)
+ c = copy(atree)
+ @test c.arena !== atree.arena
+ @test convert(Node, c) == tree
+ c.l.l.r.val = 99.0
+ @test convert(Node, atree) == tree
+ @test convert(Node, c) != tree
+ end
+
+ @testset "subtree copy falls back and re-compacts" begin
+ sub = atree.l
+ @test !AN.is_compact_root(sub)
+ csub = copy(sub)
+ @test AN.is_compact_root(csub)
+ @test convert(Node, csub) == tree.l
+ end
+
+ @testset "structural mutation invalidates fast paths" begin
+ mutated = convert(AN.ArenaNode{Float64}, tree)
+ set_child!(mutated, convert(AN.ArenaNode{Float64}, cos(x2)), 2)
+ @test !mutated.arena.compact[]
+ expected = copy(tree)
+ set_child!(expected, cos(x2), 2)
+ @test convert(Node, mutated) == expected
+ @test count_nodes(mutated) == count_nodes(expected)
+ recompacted = copy(mutated)
+ @test AN.is_compact_root(recompacted)
+ @test count_nodes(recompacted) == count_nodes(expected)
+
+ leafed = convert(AN.ArenaNode{Float64}, tree)
+ node = leafed.r
+ node.degree = 0
+ node.constant = true
+ node.val = 1.0
+ @test !leafed.arena.compact[]
+ expected2 = copy(tree)
+ expected2.r = Node{Float64}(; val=1.0)
+ @test convert(Node, leafed) == expected2
+ @test count_nodes(leafed) == count_nodes(expected2)
+ end
+
+ @testset "preallocated copy_into!" begin
+ dest = allocate_container(atree)
+ out = copy_into!(dest, atree)
+ @test out.arena === dest
+ @test convert(Node, out) == tree
+ out2 = copy_into!(dest, atree)
+ @test convert(Node, out2) == tree
+ end
+
+ @testset "copy_node entry point" begin
+ c = copy_node(atree)
+ @test c.arena !== atree.arena
+ @test convert(Node, c) == tree
+ end
+
+ @testset "Expression-level preallocated copy (SR mutation path)" begin
+ ex = Expression(
+ convert(AN.ArenaNode{Float64}, tree);
+ operators=operators,
+ variable_names=["x1", "x2"],
+ )
+ container = allocate_container(ex)
+ ex2 = copy_into!(container, ex)
+ @test convert(Node, DynamicExpressions.get_tree(ex2)) == tree
+ end
+
+ @testset "whole-tree scans match Node" begin
+ @test count_nodes(atree) == count_nodes(tree)
+ @test count(t -> t.degree == 2, atree) == count(t -> t.degree == 2, tree)
+ @test count(t -> t.degree == 0, atree.l) == count(t -> t.degree == 0, tree.l)
+ @test length(atree) == length(tree)
+ @test count_constant_nodes(atree) == count_constant_nodes(tree)
+ @test has_constants(atree) == has_constants(tree)
+ leaf = convert(AN.ArenaNode{Float64}, Node{Float64}(; feature=1))
+ @test !has_constants(leaf)
+ @test count_constant_nodes(leaf) == 0
+ end
+
+ @testset "scalar constants via arena indices" begin
+ fresh = convert(AN.ArenaNode{Float64}, tree)
+ vals, refs = get_scalar_constants(fresh)
+ @test refs isa Vector{Int32}
+ @test DynamicExpressions.count_scalar_constants(fresh) == length(vals)
+ @test vals == first(get_scalar_constants(tree))
+ set_scalar_constants!(fresh, vals .* 2, refs)
+ @test first(get_scalar_constants(fresh)) == vals .* 2
+
+ # Indices remain valid in flat copies of the tree:
+ c = copy(fresh)
+ set_scalar_constants!(c, vals, refs)
+ @test first(get_scalar_constants(c)) == vals
+
+ # Non-compact trees fall back to the generic Ref-based path:
+ sub = fresh.l
+ vsub, rsub = get_scalar_constants(sub)
+ @test vsub == first(get_scalar_constants(convert(Node, sub)))
+ set_scalar_constants!(sub, vsub .+ 1, rsub)
+ @test first(get_scalar_constants(sub)) == vsub .+ 1
+ end
+end
diff --git a/test/test_arenanode_allocations.jl b/test/test_arenanode_allocations.jl
index fc12b93..429de7b 100644
--- a/test/test_arenanode_allocations.jl
+++ b/test/test_arenanode_allocations.jl
@@ -1,5 +1,6 @@
using Test
using DynamicExpressions
+using DynamicExpressions.NodePreallocationModule: allocate_container, copy_into!
const AN = DynamicExpressions.ArenaNodeModule
@@ -26,6 +27,11 @@ function alloc_eval_tree(tree, X, operators)
return nothing
end
+function alloc_copy_into!(dest, tree)
+ copy_into!(dest, tree)
+ return nothing
+end
+
arena_push = AN.Arena{Float64,2}(; capacity=128)
base_tree = sin(x1)
@@ -40,6 +46,7 @@ child = AN.ArenaNode(child_arena, child_idx)
tree_large = sin(x1) + x1 * 3.2 + cos(x1)
atree_large = convert(AN.ArenaNode{Float64}, tree_large)
+copy_dest = allocate_container(atree_large)
arena_large = AN.Arena{Float64,2}(; capacity=128)
X = randn(Float64, 1, 1_000)
@@ -49,12 +56,14 @@ for _ in 1:5
alloc_copy_tree!(arena_large, tree_large)
alloc_eval_tree(tree_large, X, operators)
alloc_eval_tree(atree_large, X, operators)
+ alloc_copy_into!(copy_dest, atree_large)
end
alloc_counts = Dict(
"push_constant" => @allocations(alloc_push_constant!(arena_push)),
"set_child" => @allocations(alloc_set_child!(parent, child)),
"copy_tree" => @allocations(alloc_copy_tree!(arena_large, tree_large)),
+ "copy_into" => @allocations(alloc_copy_into!(copy_dest, atree_large)),
)
alloc_bytes = Dict(
"eval_node" => @allocated(alloc_eval_tree(tree_large, X, operators)),
@@ -64,4 +73,5 @@ alloc_bytes = Dict(
@test alloc_counts["push_constant"] == 0
@test alloc_counts["set_child"] == 0
@test alloc_counts["copy_tree"] == 0
+@test alloc_counts["copy_into"] == 0
@test alloc_bytes["eval_arena"] <= max(1024, ceil(Int, 1.10 * alloc_bytes["eval_node"]))
-- |
e/a/n/c/d/f/v/st and the kernel abbreviations (sp, ks, is, svs, svals, desc, nfree, fbase, doff, offs, isscal, off, rem, B) become entry, arena, node, child_idx, degree, feature, value, state, stack_top, kinds, idxs, scalar_args, scalar_vals, descriptors, num_free, free_base, dest_offset, offsets, is_scalar, offset, remaining, buffer_rows. No behavior change (2614 tests pass; bench ratios 0.55/0.43/0.39, unchanged).
- expose ArenaNode/Arena as DynamicExpressions.ArenaNode etc.; tests use 'using DynamicExpressions: ArenaNode' instead of a module alias - _push_node! takes keyword defaults; push_constant!/push_feature! only override the relevant fields - PlanState/PlanRegisters (no leading underscore on internal types) - constrain ArenaEntry/Arena/ArenaNode to T<:Number - split _arena_eval into _materialize_features! and _write_root_to_output! - readable planner names: feature_mask, scalar_stack, permanent_stack, arity_mask, num_recyclable_args; the feature-slot computation lives in a named _feature_slot helper shared with _push_leaf! - iszero/isone instead of == 0 / != 1 comparisons - shorten testitem names so headers fit on one line 2614 tests pass; bench ratios 0.56/0.47/0.37 vs Node (unchanged).
_exec_op! is now a thin compile-time arity dispatch; _exec_op_arity! pops the operand descriptors and branches once into _fold_constant_args! (the scalar-lane fold, written with guard clauses) or _run_op_kernel! (slot recycling, destination allocation, operand validation, kernel dispatch). Max nesting depth drops from 6 to 2. No behavior or speed change: 2614 tests pass, and an interleaved A/B against the previous commit gives identical benchmark ratios.
set_child! and set_children! both duplicated the validate-then-link-or-copy logic; it now lives in one helper with the cross-arena copy rationale on it.
_plan_scratch previously re-implemented the executor's stack machine as raw bitmask arithmetic — the two had to agree exactly, since the executor trusts the planner's slot counts under @inbounds, but nothing tied them together. The kind and recycling policy now lives once: _leaf_kind, _op_result_kind, and _is_recyclable are called by both sides, and the planner's bitmask pair becomes a KindStack with push/pop/query operations named after the executor's concepts. Drift between the passes is now a compile error or an obviously-local edit instead of a silent buffer-size mismatch. No behavior change: 2614 tests pass, bench ratios within session noise.
…inbounds in set_scalar_constants! get_arena/get_index are now the only getfield call sites: internal code calls them instead of raw getfield (named, greppable) and instead of property access (functions reachable from getproperty, like get_child via the :l/:r branches, must not cycle back through it). Arena has no custom getproperty, so its getfield(arena, :nodes) calls become plain dot access. set_scalar_constants! gets its @inbounds back: refs are produced by get_scalar_constants, and callers passing stale refs are out of contract.
_pack_descriptor/_descriptor_kind/_descriptor_slot replace the raw kind|slot<<2 packing, _feature_bit names the used-feature bitset position, and _slot_offset replaces _slotoff. All @inline; perf-neutral (verified by interleaved A/B against the pre-refactor base 74438bd: with early_exit disabled the two are indistinguishable; the ~5-8% ratio gap with early_exit on is entirely the operand-validation correctness fix, not the refactors).
The parity fix had moved slot validation to consumption, re-summing feature columns at every use (~5-8% relative vs the pre-fix base). Validation now happens at equivalent-but-cheaper points: - features: one check at materialization (a separate pass over the just-written, cache-hot slot; fusing the sum into the copy loop spoiled its memcpy pattern). Skipped for single-leaf trees, where no operator consumes the feature. - intermediates: checked at production, right after the kernel writes them. Every non-root intermediate is consumed exactly once, so this rejects the same trees as consumption checks. - scalar operands: still checked at consumption (O(1)). - root output: never checked, as in the generic evaluator. Also adds mark_compact!/invalidate_compact!/is_compact instead of raw arena.compact[] writes. Interleaved A/B vs 74438bd: ratios now overlap at all sizes (0.69-0.72 / 0.59-0.64 / 0.57-0.61). 2616 tests pass, including a new bare-feature-root NaN parity case.
@enum OperandKind::UInt8 FoldedConstant PinnedSlot ScratchSlot replaces the raw _K_SCALAR/_K_PSLOT/_K_SLOT byte constants. Interleaved A/B shows no measurable cost: packing and comparisons compile to the same integer ops, and the membership check in the unpacking constructor is branch-predicted away. 2616 tests pass.
This reverts commit bd0a2c7.
Only the three types (ArenaEntry, Arena, ArenaNode) keep docstrings; every internal function's docstring becomes a short '#' comment, so Documenter does not pick them up.
All 38 @inline annotations removed (interleaved A/B shows the inliner does the same job unannotated; the zero-allocation copy tests still pass). The one survivor is the call-site @inline on the predicate in _arena_any, which forces the closure to specialize into the recursion. Also: setproperty! branches share a single trailing 'return value'; getproperty loads the entry once instead of in five branches; trivial helpers collapse to one-line definitions; _nullable_child and setproperty! use get_arena/get_index like everything else.
Mutation hot path; inlining guarantees the constant-symbol branch folds at the call site rather than relying on interprocedural constprop.
Documents the self-buffering experiment result: the generic fused unbuffered evaluator benches even with a plan running on a self-allocated pool, while allocating half the bytes, so the plan stays buffer-gated. Also shows buffered Node is slower than unbuffered Node (strided buffer rows), which is the headroom the plan's contiguous slots reclaim.
ArenaNode.jl keeps the types and the AbstractExpressionNode interface (~510 lines); ArenaNodeEval.jl holds the buffered plan evaluator (~530 lines), included from within the module. Also drops a tutorial-grade comment and records why the plan does not self-allocate a pool.
…t compacting copy
CI:
- _scalar_degn binds T via Tuple{T,Vararg{T}} and derives the arity from
the tuple type, fixing Aqua's unbound-args failure
- formatted with JuliaFormatter v1 (the CI pin; local runs had used v2)
Investigation of the slow any(f, ::ArenaNode) found the generic
AbstractNode traversal machinery is already fast on ArenaNode (within 15%
of Node, zero allocations) -- the custom _arena_any/_arena_mapreduce/
is_constant overrides were unnecessary, and the closure-based recursion in
the any override was nondeterministically up to 8x slower (inlining of the
Base.any layers is compilation-order dependent). So:
- the traversal overrides are deleted; ArenaNode rides the generic
machinery, except a 3-line compact-arena flat scan for (1.4-1.7x
faster than Node, and has_constants/is_constant inherit it)
- the unset-child UndefRefError guard moves from the deleted override into
_load_entry at the facade layer, so poison facades throw like Node's
undefined fields for every generic consumer
Copying out of a non-compact arena is now an entry-level compacting write
(_write_subtree! into a pre-sized vector; no facade traversal, no per-push
growth checks), also used by copy_into! and cross-arena attachment.
copy(::ArenaNode): compact 5.7-11x faster than Node at n=15-255;
non-compact crosses over near n=20 (1.1-1.2x above).
ArenaNodeEval.jl is now its own module included from DynamicExpressions.jl
rather than a nested include. The eval benchmark also reports unbuffered
paths.
Property tests (using the existing supposition_utils generators): Node -> ArenaNode -> Node round-trips with all read-only interface results equal; evaluation matches Node (unbuffered exactly; the buffered plan's ok flag is best-effort, values must agree whenever both report ok); random valid mutations applied to both representations keep them equivalent; copy re-compacts and preserves evaluation.
…_nodes hash went through the generic tree_mapreduce machinery (0.64x vs Node); direct recursion over entries through the shared leaf_hash/branch_hash combinators recovers it to 0.92x -- the rest is the hash arithmetic itself, which must be identical across representations since cross-representation == requires equal hashes (now pinned by the supposition round-trip property). count_constant_nodes on a compact arena is a flat count over the entry vector (1.6x vs Node).
The generic traversal skeleton re-reads the entry through the facade for every field access and shuffles Nullable-wrapped children, which taxed all tree_mapreduce clients at once (hash 0.64x, collect 0.70x, count_depth 0.76x, count_constant_nodes 0.73x vs Node). One override of the primitive -- read each entry exactly once, dispatch on the local degree, hand f the facade -- normalizes the whole family (1.0x for collect/count_depth/ count_constant_nodes, 0.94x index_constant_nodes, 0.85x hash) in both compact and non-compact arenas. This replaces the per-function fixes: the previous hash and count_constant_nodes specializations are deleted, and since the generic hash combinators now just run on the fast skeleton, hash equality across representations holds again (pinned in the supposition round-trip property).
At 31-node trees the generic any falls to 0.52x vs Node on non-compact arenas (facade re-reads per field, like the tree_mapreduce skeleton); the entry-level recursion recovers to 0.67x. The remainder is f reading fields through the facade (arena -> nodes -> entry per access vs Node's single deref), inherent to the facade contract; compact arenas bypass it entirely via the flat scan (1.8x ahead of Node at n=31).
A linear postfix scan over compact arenas benches ~20% faster for pure reductions (hash 1.03x vs Node), but f-application order (parent first, siblings left to right) is observable API through collect/filter/foreach: the postfix scan reorders them, applying positional mutations to different nodes than Node would. Caught by the supposition properties with shrunk counterexamples (abs(5.0e-324), one set_val). Recursion stays; the comment on _entry_mapreduce documents the constraint.
The flat scan was the recursive traversal with recursion disabled, so the
two share a single _entry_any whose children loop is gated on a static
Val{recurse} flag: the compact path drives it from a flat loop with the
recursion compiled out, the non-compact path recurses. Performance
unchanged (compact 1.6-2.0x vs Node at n=31, non-compact 0.63-0.68x).
…tant ops One driver for order-free whole-tree visits, with both walks side by side and the unused one compiled out via a static flag: compact arenas walk the entry array (indices valid by construction, so the load is uncheckable and dead when unused); non-compact arenas walk the children pointers with the per-degree dispatch unrolled like the generic any (a children loop benches ~40% slower; a worklist allocates and loses at SR sizes; closure-driven recursion through Base.any inlines nondeterministically -- all measured). any/count_nodes/get_scalar_constants are one-line clients, replacing the recursive _entry_any, the count_nodes invoke fallback, and the filter_map path. Also, on the canonical-form principle (pure or order-free ops on a compact arena may use the entries directly): - == between compact trees is a flat field-compare (equal trees have identical entry vectors); 2.7x vs Node, from 1.1x - count_constant_nodes: flat count when compact (2.0x), entry-level mapreduce otherwise - set_scalar_constants! writes entries directly: val-only writes cannot change structure, so the setindex! degree/children diff is skipped - @inline restored on getproperty (the removal had been validated against an eval-only benchmark that never exercises property reads) n=31 vs Node: compact 1.4-3.4x on ten ops, floor 0.81x (hash); non-compact 0.85-1.75x, floor 0.65x (copy_into!, the re-compaction price that wins back above n~20).
Older Supposition versions resolved by the downgrade-compat CI job lack length(::Data.OneOf); a single multi-arg map selecting the mutation kind by integer generates the same tuples on any version.
MilesCranmerBot
force-pushed
the
arenanode-prototype
branch
from
August 1, 2026 14:28
60d0023 to
9850b0b
Compare
Member
|
@MilesCranmerBot please merge or rebase on latest master |
# Conflicts: # ext/DynamicExpressionsLoopVectorizationExt.jl # test/Project.toml
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Prototype
ArenaNode: an array-backed arena representation for expressions, implementing the fullNodeInterface. Evaluation runs directly on the arena facade and matchesNodeon allocations (scripts/bench_arenanode.jl), with cheapercopyand a smaller memory footprint.