diff --git a/Project.toml b/Project.toml index 0da0de6a..bcca1513 100644 --- a/Project.toml +++ b/Project.toml @@ -41,7 +41,7 @@ Optim = "1, 2" NLSolversBase = "7, 8" PrecompileTools = "1.2.1" Reexport = "1.2.2" -SymbolicUtils = "4" +SymbolicUtils = "4.35" Zygote = "0.7" julia = "1.10" Random = "1" diff --git a/ext/DynamicExpressionsLoopVectorizationExt.jl b/ext/DynamicExpressionsLoopVectorizationExt.jl index 7b41f767..521d2e7d 100644 --- a/ext/DynamicExpressionsLoopVectorizationExt.jl +++ b/ext/DynamicExpressionsLoopVectorizationExt.jl @@ -211,12 +211,16 @@ function deg2_r0_eval( end # Interface with Bumper.jl -function bumper_kern!( +@generated function bumper_kern!( op::F, cumulators::Tuple{Vararg{Any,degree}}, ::EvalOptions{true,true,early_exit} ) where {F,degree,early_exit} - cumulator_1 = first(cumulators) - @turbo @. cumulator_1 = op(cumulators...) - return cumulator_1 + quote + Base.Cartesian.@nexprs($degree, i -> cumulator_i = cumulators[i]) + @turbo for j in eachindex(cumulator_1) + cumulator_1[j] = Base.Cartesian.@ncall($degree, op, i -> cumulator_i[j]) + end + return cumulator_1 + end end end diff --git a/scripts/bench_arenanode_eval.jl b/scripts/bench_arenanode_eval.jl new file mode 100644 index 00000000..08fdbb6a --- /dev/null +++ b/scripts/bench_arenanode_eval.jl @@ -0,0 +1,89 @@ +using DynamicExpressions +using DynamicExpressions: EvalOptions, ArrayBuffer +using Random: MersenneTwister + +include(joinpath(@__DIR__, "..", "test", "tree_gen_utils.jl")) + +const AN = DynamicExpressions.ArenaNodeModule + +operators = OperatorEnum(1 => (cos, exp), 2 => (+, -, *, /)) +const T = Float64 +nfeat = 5 +n = 1_000 +rng = MersenneTwister(0) +X = randn(rng, T, nfeat, n) +buf = zeros(T, 64, n) + +function bench(trees, X, operators, buffer; reps=300) + best = Inf + for _ in 1:reps + t0 = time_ns() + for tree in trees + buffer.index[] = 0 + eval_tree_array(tree, X, operators; eval_options=EvalOptions(; buffer)) + end + best = min(best, (time_ns() - t0) / length(trees)) + end + return best +end + +function allocs_per_eval(tree, X, operators, buffer) + @allocated(eval_tree_array(tree, X, operators; eval_options=EvalOptions(; buffer))) +end + +for treesize in (7, 15, 31) + trees = [gen_random_tree_fixed_size(treesize, operators, nfeat, T) for _ in 1:50] + atrees = [convert(AN.ArenaNode{T,2}, t) for t in trees] + buffer = ArrayBuffer(buf, Ref(0)) + + # correctness sanity: both paths must agree + for (t, a) in zip(trees, atrees) + buffer.index[] = 0 + yn, okn = eval_tree_array(t, X, operators; eval_options=EvalOptions(; buffer)) + buffer.index[] = 0 + ya, oka = eval_tree_array(a, X, operators; eval_options=EvalOptions(; buffer)) + okn == oka || error("ok mismatch at size $treesize") + okn && (yn ≈ ya || error("value mismatch at size $treesize")) + end + + # warmup + bench(trees, X, operators, buffer; reps=3) + bench(atrees, X, operators, buffer; reps=3) + + t_node = bench(trees, X, operators, buffer) + t_arena = bench(atrees, X, operators, buffer) + buffer.index[] = 0 + a_node = allocs_per_eval(trees[1], X, operators, buffer) + buffer.index[] = 0 + a_arena = allocs_per_eval(atrees[1], X, operators, buffer) + println( + "n=$treesize Node: $(round(t_node/1e3; digits=2))us " * + "ArenaNode: $(round(t_arena/1e3; digits=2))us " * + "ratio=$(round(t_arena/t_node; digits=3)) " * + "allocs/eval Node=$a_node Arena=$a_arena", + ) + + bench_nobuf(trees) = begin + best = Inf + for _ in 1:300 + t0 = time_ns() + for tree in trees + eval_tree_array(tree, X, operators) + end + best = min(best, (time_ns() - t0) / length(trees)) + end + best + end + bench_nobuf(trees[1:2]) + bench_nobuf(atrees[1:2]) # warmup + tn_nb = bench_nobuf(trees) + ta_nb = bench_nobuf(atrees) + an_nb = @allocated(eval_tree_array(trees[1], X, operators)) + aa_nb = @allocated(eval_tree_array(atrees[1], X, operators)) + println( + "n=$treesize unbuffered: Node: $(round(tn_nb/1e3; digits=2))us " * + "ArenaNode: $(round(ta_nb/1e3; digits=2))us " * + "ratio=$(round(ta_nb/tn_nb; digits=3)) " * + "allocs/eval Node=$an_nb Arena=$aa_nb", + ) +end diff --git a/src/ArenaNode.jl b/src/ArenaNode.jl new file mode 100644 index 00000000..65febb22 --- /dev/null +++ b/src/ArenaNode.jl @@ -0,0 +1,645 @@ +module ArenaNodeModule + +using ..UtilsModule: Nullable, Undefined + +import ..NodeModule: + AbstractNode, + AbstractExpressionNode, + Node, + unsafe_get_children, + get_child, + set_child!, + set_children!, + count_nodes, + copy_node, + tree_mapreduce, + inner_is_equal +import ..NodeUtilsModule: count_constant_nodes, get_scalar_constants, set_scalar_constants! +import ..NodePreallocationModule: allocate_container, copy_into! +import ..ValueInterfaceModule: get_number_type, is_valid, is_valid_array +import ..OperatorEnumModule: OperatorEnum + +"""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<:Number,D} + val::T + children::NTuple{D,Int32} + feature::UInt16 + degree::UInt8 + op::UInt8 + constant::Bool +end + +function _replace( + entry::ArenaEntry{T,D}; + val=entry.val, + children=entry.children, + feature=entry.feature, + degree=entry.degree, + op=entry.op, + constant=entry.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. + +`Arena` implements the (1-based, linear) array interface over its entries, and +this is the only sanctioned way to mutate entries: `setindex!` compares the old +and new entry and automatically clears `compact` whenever the structural fields +(`degree`, `children`) change. Non-structural writes (`val`/`op`/`feature`/ +`constant`) preserve the flag, which keeps constant optimization on the fast +paths. Do not write `arena.nodes` directly outside this file's bulk-copy +internals. +""" +struct Arena{T<:Number,D} <: AbstractVector{ArenaEntry{T,D}} + nodes::Vector{ArenaEntry{T,D}} + compact::Base.RefValue{Bool} + + function Arena{T,D}(; capacity::Integer=0) where {T,D} + 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 + +# Mark/clear the one-postfix-tree invariant (root last, no orphans): +mark_compact!(arena::Arena) = (arena.compact[] = true; arena) +invalidate_compact!(arena::Arena) = (arena.compact[] = false; arena) +is_compact(arena::Arena) = arena.compact[] + +Base.size(arena::Arena) = size(arena.nodes) +Base.IndexStyle(::Type{<:Arena}) = IndexLinear() +Base.@propagate_inbounds Base.getindex(arena::Arena, i::Integer) = arena.nodes[i] +Base.@propagate_inbounds function Base.setindex!( + arena::Arena{T,D}, entry::ArenaEntry{T,D}, i::Integer +) where {T,D} + nodes = arena.nodes + old = nodes[i] + if entry.degree != old.degree || entry.children != old.children + invalidate_compact!(arena) + end + nodes[i] = entry + return arena +end +function Base.push!(arena::Arena{T,D}, entry::ArenaEntry{T,D}) where {T,D} + nodes = arena.nodes + # A single leaf in a fresh arena is a valid tree; any further append breaks + # the one-postfix-tree invariant until a builder re-establishes it. + isempty(nodes) || invalidate_compact!(arena) + push!(nodes, entry) + return arena +end +function Base.sizehint!(arena::Arena, capacity::Integer) + sizehint!(arena.nodes, capacity) + return arena +end + +"""A lightweight facade for a node stored in an [`Arena`](@ref). + +This wrapper is intentionally minimal: it stores an arena reference and an index. +Core fields are accessed and mutated via `getproperty`/`setproperty!`. + +!!! warning + Unlike `Node`, attaching a child from a *different* arena + (`set_child!`/`set_children!`, including keyword construction) copies the + subtree into the parent's arena: the original handle stays attached to its + own arena, so later mutations through it do not affect the new parent. + Same-arena attachments keep reference semantics. +""" +struct ArenaNode{T<:Number,D} <: AbstractExpressionNode{T,D} + arena::Arena{T,D} + idx::Int32 + + function ArenaNode{T,D}(arena::Arena{T,D}, idx::Int32) where {T,D} + return new{T,D}(arena, idx) + end +end + +ArenaNode(arena::Arena{T,D}, idx::Int32) where {T,D} = ArenaNode{T,D}(arena, idx) + +# The only sanctioned `getfield` sites: internal code uses these instead of +# property access so that functions reachable from `getproperty` (e.g. +# `get_child` via the `:l`/`:r` branches) do not cycle back into it. +get_arena(node::ArenaNode) = getfield(node, :arena) +get_index(node::ArenaNode) = getfield(node, :idx) + +# Entry behind a facade. A zero index is a poison facade (an unset child +# slot); accessing it throws like reading an undefined `Node` field. +function _load_entry(arena::Arena, idx::Int32) + iszero(idx) && throw(UndefRefError()) + return @inbounds arena.nodes[idx] +end + +# True when the arena contents *are* `tree`: compact, rooted at the last entry. +function is_compact_root(tree::ArenaNode) + arena = get_arena(tree) + return is_compact(arena) && get_index(tree) == length(arena.nodes) +end + +_zero_children(::Val{D}) where {D} = ntuple(_ -> Int32(0), Val(D)) + +function _push_node!( + arena::Arena{T,D}; + degree::UInt8=UInt8(0), + constant::Bool=false, + val::T=zero(T), + feature::UInt16=UInt16(0), + op::UInt8=UInt8(0), + children::NTuple{D,Int32}=_zero_children(Val(D)), +) where {T,D} + push!(arena, ArenaEntry{T,D}(val, children, feature, degree, op, constant)) + return Int32(length(arena)) +end + +function push_constant!(arena::Arena{T,D}, value) where {T,D} + return _push_node!(arena; constant=true, val=convert(T, value)) +end + +function push_feature!(arena::Arena{T,D}, feature::Integer) where {T,D} + return _push_node!(arena; feature=UInt16(feature)) +end + +# Default node: a zero constant leaf in its own fresh arena. +function ArenaNode{T,D}() where {T,D} + arena = Arena{T,D}() + idx = push_constant!(arena, zero(T)) + return ArenaNode{T,D}(arena, idx) +end + +Base.@constprop :aggressive @inline function Base.getproperty( + node::ArenaNode{T}, property_name::Symbol +) where {T} + if property_name === :arena + return get_arena(node) + elseif property_name === :idx + return get_index(node) + elseif property_name === :children + return unsafe_get_children(node) + elseif property_name === :l + return get_child(node, UInt8(1)) + elseif property_name === :r + return get_child(node, UInt8(2)) + end + entry = _load_entry(get_arena(node), get_index(node)) + if property_name === :degree + return entry.degree + elseif property_name === :constant + return entry.constant + elseif property_name === :val + return entry.val::T + elseif property_name === :feature + return entry.feature + elseif property_name === :op + return entry.op + else + return getfield(node, property_name) + end +end + +@inline function Base.setproperty!( + node::ArenaNode{T,D}, property_name::Symbol, value +) where {T,D} + arena = get_arena(node) + i = get_index(node) + entry = _load_entry(arena, i) + if property_name === :degree + @inbounds arena[i] = _replace(entry; degree=UInt8(value)) + elseif property_name === :constant + @inbounds arena[i] = _replace(entry; constant=Bool(value)) + elseif property_name === :val + @inbounds arena[i] = _replace(entry; val=convert(T, value)) + elseif property_name === :feature + @inbounds arena[i] = _replace(entry; feature=UInt16(value)) + elseif property_name === :op + @inbounds arena[i] = _replace(entry; op=UInt8(value)) + elseif property_name === :l + set_child!(node, value, 1) + elseif property_name === :r + set_child!(node, value, 2) + else + throw(ArgumentError("Unsupported field $property_name for ArenaNode")) + end + return value +end + +function _nullable_child( + node::ArenaNode{T,D}, child_idx::Int32 +)::Nullable{ArenaNode{T,D}} where {T,D} + child = ArenaNode{T,D}(get_arena(node), child_idx) + return Nullable{ArenaNode{T,D}}(iszero(child_idx), child) +end + +# Children as `Nullable` wrappers; unused slots are poison nodes (like `Node`), +# so accessing them throws an `UndefRefError`. +@generated function unsafe_get_children(node::ArenaNode{T,D}) where {T,D} + quote + $(Expr(:meta, :inline)) + children = _load_entry(get_arena(node), get_index(node)).children + return Base.Cartesian.@ntuple($D, j -> _nullable_child(node, children[j])) + end +end + +function get_child(node::ArenaNode{T,D}, i::Integer) where {T,D} + # Avoid routing through getproperty here: the :l/:r property branches call + # get_child, and the resulting inference cycle widens property access. + arena = get_arena(node) + entry = _load_entry(arena, get_index(node)) + child_idx = entry.children[i] # bounds-checked: i > D must throw, not crash + iszero(child_idx) && throw(UndefRefError()) + return ArenaNode{T,D}(arena, child_idx) +end + +# Same-arena children attach by index; anything else is copied into `node`'s +# arena, since arenas cannot link across each other. +function _resolve_child_index!(node::ArenaNode{T,D}, child) where {T,D} + child isa AbstractExpressionNode{T,D} || throw( + ArgumentError( + "ArenaNode children must be AbstractExpressionNode{$T,$D} (got $(typeof(child)))", + ), + ) + if child isa ArenaNode{T,D} && child.arena === node.arena + return child.idx + else + return _copy_to_arena!(node.arena, child) + end +end + +function set_child!(node::ArenaNode{T,D}, child::AbstractNode{D}, i::Int) where {T,D} + arena = node.arena + entry = @inbounds arena[node.idx] + 1 <= i <= D || throw(BoundsError(entry.children, i)) + idx = _resolve_child_index!(node, child) + if @inbounds(entry.children[i]) != idx + @inbounds arena[node.idx] = _replace( + entry; children=Base.setindex(entry.children, idx, i) + ) + end + return ArenaNode{T,D}(arena, idx) +end + +function set_children!( + node::ArenaNode{T,D}, children::Union{Tuple,AbstractVector{<:AbstractNode{D}}} +) where {T,D} + D2 = length(children) + idxs = _zero_children(Val(D)) + @inbounds for i in 1:min(D, D2) + child = children[i] + if child isa Nullable + child.null && continue + child = child[] + end + idxs = Base.setindex(idxs, _resolve_child_index!(node, child), i) + end + + arena = node.arena + entry = @inbounds arena[node.idx] + @inbounds arena[node.idx] = _replace(entry; children=idxs) + return nothing +end + +# Append `idx`'s subtree (children first, root last) to `dest`, returning the +# root's new index. Works directly on entries -- no facade traversal -- so +# copying out of a non-compact arena stays an array operation. `dest` is +# grown to `length(dest) + count(subtree)`; the caller may oversize it first +# to skip per-entry growth (see `_append_subtree!`'s wrapper below). +function _write_subtree!( + dest::Vector{ArenaEntry{T,D}}, src::Vector{ArenaEntry{T,D}}, idx::Int32, cursor::Int +) where {T,D} + iszero(idx) && throw(UndefRefError()) # unset child slot, like Node + entry = @inbounds src[idx] + if !iszero(entry.degree) + children = _zero_children(Val(D)) + @inbounds for j in 1:(entry.degree) + child_idx, cursor = _write_subtree!(dest, src, entry.children[j], cursor) + children = Base.setindex(children, child_idx, j) + end + entry = _replace(entry; children) + end + cursor += 1 + @inbounds dest[cursor] = entry + return Int32(cursor), cursor +end + +# `src` and `dest` may be the same vector: reads are below the original +# length, writes at or above `cursor`, and the subtree size never exceeds +# the source length, so the regions cannot overlap incorrectly. +function _append_subtree!( + dest::Vector{ArenaEntry{T,D}}, src::Vector{ArenaEntry{T,D}}, idx::Int32 +) where {T,D} + cursor = length(dest) + resize!(dest, cursor + length(src)) # upper bound; trimmed below + root_idx, cursor = _write_subtree!(dest, src, idx, cursor) + resize!(dest, cursor) + return root_idx +end + +# Compact arenas copy as one flat array copy (child indices stay valid +# verbatim); otherwise the subtree is appended entry-by-entry into a fresh +# arena, which also re-compacts. Overloads `copy_node` (not `Base.copy`) +# since that is the generic entry point. +function copy_node(tree::ArenaNode{T,D}; break_sharing::Val{BS}=Val(false)) where {T,D,BS} + arena = get_arena(tree) + if is_compact_root(tree) + return ArenaNode{T,D}(Arena{T,D}(copy(arena.nodes), true), get_index(tree)) + end + nodes = sizehint!(ArenaEntry{T,D}[], length(arena.nodes)) + idx = _append_subtree!(nodes, arena.nodes, get_index(tree)) + return ArenaNode{T,D}(Arena{T,D}(nodes, true), idx) +end + +# Preallocated arena for `copy_into!`, enabling zero-allocation copies. +function allocate_container( + prototype::ArenaNode{T,D}, num_nodes::Union{Nothing,Integer}=nothing +) where {T,D} + return Arena{T,D}(; capacity=@something(num_nodes, length(prototype))) +end + +# Steady-state copy for population search: reuse `dest`'s storage, with no +# allocations once it has sufficient capacity. +function copy_into!( + dest::Arena{T,D}, + src::ArenaNode{T,D}; + ref::Union{Nothing,Base.RefValue{<:Integer}}=nothing, +) where {T,D} + set_ref!(n) = ref === nothing ? nothing : (ref[] = n) + if dest === get_arena(src) + # Container reuse: the tree already lives in `dest`. A compact root is + # a no-op; otherwise compact through a temporary copy. + if is_compact_root(src) + set_ref!(length(src)) + return src + end + return copy_into!(dest, copy_node(src); ref) + end + if is_compact_root(src) + nodes = src.arena.nodes + resize!(dest.nodes, length(nodes)) + copyto!(dest.nodes, nodes) + mark_compact!(dest) + set_ref!(length(src)) + return ArenaNode{T,D}(dest, src.idx) + end + empty!(dest.nodes) + idx = _append_subtree!(dest.nodes, get_arena(src).nodes, get_index(src)) + mark_compact!(dest) + set_ref!(length(dest.nodes)) + return ArenaNode{T,D}(dest, idx) +end + +function _copy_to_arena!(arena::Arena{T,D}, tree::ArenaNode{T,D}) where {T,D} + invalidate_compact!(arena) + return _append_subtree!(arena.nodes, get_arena(tree).nodes, get_index(tree)) +end +function _copy_to_arena!( + arena::Arena{T,D}, tree::AbstractExpressionNode{T2,D} +) where {T,T2,D} + degree = tree.degree + if degree == 0 + if tree.constant + return push_constant!(arena, tree.val) + else + return push_feature!(arena, tree.feature) + end + end + + idxs = _zero_children(Val(D)) + @inbounds for i in 1:degree + idxs = Base.setindex(idxs, _copy_to_arena!(arena, get_child(tree, i)), i) + end + return _push_node!(arena; degree=UInt8(degree), op=tree.op, children=idxs) +end + +# Copy the tree into a fresh arena, in postfix (children-first) order. +function Base.convert( + ::Type{ArenaNode{T,D}}, tree::AbstractExpressionNode{T2,D} +) where {T,T2,D} + arena = Arena{T,D}(; capacity=length(tree; break_sharing=Val(true))) + idx = _copy_to_arena!(arena, tree) + mark_compact!(arena) + return ArenaNode{T,D}(arena, idx) +end +function Base.convert( + ::Type{ArenaNode{T}}, tree::AbstractExpressionNode{T2,D} +) where {T,T2,D} + return convert(ArenaNode{T,D}, tree) +end + +# Cross-representation comparisons (`==` promotes its arguments) promote +# toward the arena representation. +function Base.promote_rule(::Type{ArenaNode{T1,D}}, ::Type{Node{T2,D}}) where {T1,T2,D} + return ArenaNode{promote_type(T1, T2),D} +end +function Base.promote_rule(::Type{ArenaNode{T1,D}}, ::Type{ArenaNode{T2,D}}) where {T1,T2,D} + return ArenaNode{promote_type(T1, T2),D} +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. +################################################################################ + +# The single iteration driver for order-free whole-tree visits. +# `visit(node, entry)` receives the facade and the already-loaded entry, and +# returns true to stop early (the driver then returns true). The two walks +# sit side by side with the unused one compiled out: a compact arena *is* +# the tree, so it is walked as an array (storage order, leaves first); a +# non-compact arena is walked through the children pointers (preorder). +function _foreach_node(visit::F, tree::ArenaNode{T,D}) where {F<:Function,T,D} + arena = get_arena(tree) + if is_compact_root(tree) + return _foreach_node(visit, arena, get_index(tree), Val(true)) + else + return _foreach_node(visit, arena, get_index(tree), Val(false)) + end +end +@generated function _foreach_node( + visit::F, arena::Arena{T,D}, root::Int32, ::Val{compact} +) where {F<:Function,T,D,compact} + if compact + # The arena is the tree: walk the array. Indices 1:root are valid by + # construction, so the entry load needs no poison check (and is dead + # code when `visit` ignores it). + return quote + @inbounds for i in 1:root + entry = arena.nodes[i] + @inline(visit(ArenaNode{T,D}(arena, Int32(i)), entry)) && return true + end + return false + end + else + # Walk the children pointers, with the per-degree dispatch unrolled + # like the generic `any` -- this benches at Node parity, where a + # children loop or a worklist do not. + return quote + entry = _load_entry(arena, root) + @inline(visit(ArenaNode{T,D}(arena, root), entry)) && return true + degree = entry.degree + iszero(degree) && return false + children = entry.children + return Base.Cartesian.@nif( + $D, + i -> degree == i, # COV_EXCL_LINE + i -> Base.Cartesian.@nany( + i, j -> _foreach_node(visit, arena, children[j], Val(false)) + ), + ) + end + end +end + +# A compact arena is a canonical form: equal trees have identical entry +# vectors (same layout, same child indices), so equality is a flat +# field-compare. `val` is compared with `==` (not bitwise) to keep the +# -0.0 == 0.0 and NaN != NaN semantics of the structural fallback. +function Base.:(==)(a::ArenaNode{T,D}, b::ArenaNode{T,D})::Bool where {T,D} + if is_compact_root(a) && is_compact_root(b) + nodes_a = get_arena(a).nodes + nodes_b = get_arena(b).nodes + length(nodes_a) == length(nodes_b) || return false + @inbounds for i in eachindex(nodes_a) + ea, eb = nodes_a[i], nodes_b[i] + same = if ea.degree != eb.degree + false + elseif iszero(ea.degree) + ea.constant == eb.constant && + (ea.constant ? ea.val == eb.val : ea.feature == eb.feature) + else + ea.op == eb.op && ea.children == eb.children + end + same || return false + end + return true + end + return inner_is_equal(a, b, nothing) # preserve_sharing is false +end + +function count_nodes(tree::ArenaNode; break_sharing::Val{BS}=Val(false)) where {BS} + is_compact_root(tree) && return length(get_arena(tree).nodes) + counter = Base.RefValue(0) + _foreach_node((_, _) -> (counter[] += 1; false), tree) + return counter[] +end + +function count_constant_nodes(tree::ArenaNode; break_sharing::Val{BS}=Val(false)) where {BS} + arena = get_arena(tree) + if is_compact_root(tree) + return count(entry -> iszero(entry.degree) && entry.constant, arena.nodes) + end + return _entry_mapreduce( + node -> Int(node.degree == 0 && node.constant), + Returns(0), + +, + arena, + get_index(tree), + ) +end + +# The deep traversal primitive. The generic skeleton re-reads the entry +# through the facade for every field access and shuffles Nullable-wrapped +# children; reading each entry exactly once and dispatching on the local +# degree recovers Node-level traversal speed for every tree_mapreduce +# client (hash, collect, count_depth, count_constant_nodes, ...) at once. +# +# Recursion is load-bearing here: the f-application order (parent first, +# siblings left to right) is observable through collect/filter/foreach and +# must match Node's. A linear postfix scan over the compact arena benches +# ~20% faster for pure reductions but visits leaves first, which reorders +# collect and friends (caught by the supposition properties). +# `f_leaf`/`f_branch` still receive facades, so semantics are unchanged. +# Sharing kwargs are accepted and ignored: ArenaNode has +# preserve_sharing == false, for which the generic path never uses them. +function tree_mapreduce( + f_leaf::F1, + f_branch::F2, + op::G, + tree::ArenaNode{T,D}, + result_type::Type{RT}=Undefined; + f_on_shared::H=(result, is_shared) -> result, + break_sharing::Val{BS}=Val(false), +) where {T,D,F1<:Function,F2<:Function,G<:Function,H<:Function,RT,BS} + return _entry_mapreduce(f_leaf, f_branch, op, get_arena(tree), get_index(tree)) +end + +@generated function _entry_mapreduce( + f_leaf::F1, f_branch::F2, op::G, arena::Arena{T,D}, idx::Int32 +) where {F1<:Function,F2<:Function,G<:Function,T,D} + quote + entry = _load_entry(arena, idx) + node = ArenaNode{T,D}(arena, idx) + degree = entry.degree + iszero(degree) && return @inline(f_leaf(node)) + branch = @inline(f_branch(node)) + return Base.Cartesian.@nif( + $D, + i -> degree == i, # COV_EXCL_LINE + i -> @inline( + op( + branch, + Base.Cartesian.@ntuple( + i, + j -> _entry_mapreduce( + f_leaf, f_branch, op, arena, entry.children[j] + ) + )..., + ) + ), + ) + end +end + +function Base.any(f::F, tree::ArenaNode{T,D}) where {F<:Function,T,D} + return _foreach_node((node, _) -> @inline(f(node)), tree) +end + +# Constants as plain Int32 arena indices (also valid in flat copies). Both +# walks visit leaves left to right, so the order matches Node's. +function get_scalar_constants( + tree::ArenaNode{T}, ::Type{BT}=get_number_type(T) +) where {T<:Number,BT} + arena = get_arena(tree) + refs = Int32[] + _foreach_node(tree) do node, entry + if iszero(entry.degree) && entry.constant + push!(refs, get_index(node)) + end + return false + end + vals = T[@inbounds(arena[i].val) for i in refs] + return vals, refs +end + +function set_scalar_constants!( + tree::ArenaNode{T}, constants, refs::AbstractVector{Int32} +) where {T<:Number} + # Val-only writes cannot change structure, so they go to the entries + # directly, skipping `setindex!`'s degree/children diff. + nodes = get_arena(tree).nodes + @inbounds for j in eachindex(refs, constants) + i = refs[j] + nodes[i] = _replace(nodes[i]; val=convert(T, constants[j])) + end + return tree +end + +end diff --git a/src/ArenaNodeEval.jl b/src/ArenaNodeEval.jl new file mode 100644 index 00000000..6a1e7ee7 --- /dev/null +++ b/src/ArenaNodeEval.jl @@ -0,0 +1,544 @@ +module ArenaNodeEvalModule + +using ..UtilsModule: ResultOk +import ..NodeModule: AbstractExpressionNode +using ..ValueInterfaceModule: is_valid +import ..OperatorEnumModule: OperatorEnum +import ..EvaluateModule: _eval_tree_array, EvalOptions, ArrayBuffer, get_nops +import ..ArenaNodeModule: + ArenaNode, Arena, ArenaEntry, get_arena, get_index, is_compact_root + +################################################################################ +# Plan-style buffered evaluation +################################################################################ + +# Plan-style postfix evaluation into the caller-provided `EvalOptions.buffer`, +# treated as a flat pool of contiguous `n_rows` slots (slot 1 is the output; +# the `index` protocol is bypassed). Each used feature materializes once into +# a permanent slot, intermediates are register-allocated with a free list, +# and constant subtrees fold in a scalar lane. Addressing the buffer +# contiguously is the point: the generic evaluator hands out strided buffer +# rows, which vectorize poorly. Without a (sufficient) buffer the generic +# evaluator is used -- its fused unbuffered path benches even with a +# self-allocated plan pool while allocating half the bytes, so we do not +# build a pool ourselves. Other fallbacks: non-compact arena, non-isbits `T` +# (the branchless kernels issue dead loads from unwritten slots), depth or +# feature count over 64, turbo, or `use_fused=Val(false)` (callers may +# overload `deg1_eval` etc., which this path bypasses). +function _eval_tree_array( + tree::ArenaNode{T,D}, + cX::AbstractMatrix{T}, + operators::OperatorEnum, + eval_options::EvalOptions, +)::ResultOk where {T<:Number,D} + buffer = eval_options.buffer + if buffer isa ArrayBuffer{Matrix{T}} && + isbitstype(T) && + cX isa Matrix{T} && + size(buffer.array, 2) == size(cX, 2) && + is_compact_root(tree) && + eval_options.turbo isa Val{false} && + eval_options.use_fused isa Val{true} + ok_plan, num_slots, max_stack, feature_mask = _plan_scratch(get_arena(tree)) + # +1 for the output slot; capacity is the buffer's row count + if ok_plan && num_slots + 1 <= size(buffer.array, 1) + return _arena_eval( + get_arena(tree), + cX, + operators, + eval_options.early_exit, + num_slots, + max_stack, + feature_mask, + buffer.array, + ) + end + end + return invoke( + _eval_tree_array, + Tuple{AbstractExpressionNode{T,D},AbstractMatrix{T},OperatorEnum,EvalOptions}, + tree, + cX, + operators, + eval_options, + ) +end + +# Pool slot of materialized `feature`: slot 1 is the output; features fill +# slots 2, 3, ... in ascending feature order. +function _feature_slot(feature_mask::UInt64, feature::Integer) + return count_ones(feature_mask & (_feature_bit(feature) - 1)) + 2 +end + +# Descriptor kinds for evaluation stack slots: +const _K_SCALAR = 0x00 # folded constant; value lives in the scalar lane +const _K_PSLOT = 0x01 # permanent slot (output or a materialized feature) +const _K_SLOT = 0x02 # recyclable slot (an intermediate) + +# The planner (`_plan_scratch`) and the executor (`_push_leaf!`/`_exec_op!`) +# walk the same postfix program, so they must make identical kind and +# recycling decisions: the executor trusts the planner's slot counts under +# `@inbounds`. These three functions are the single source of that policy. + +# Leaves: constants fold into the scalar lane, features live in permanent slots. +_leaf_kind(entry::ArenaEntry) = entry.constant ? _K_SCALAR : _K_PSLOT + +# Operators: an all-scalar application constant-folds; anything else lands in +# a recyclable intermediate slot. +_op_result_kind(all_args_scalar::Bool) = all_args_scalar ? _K_SCALAR : _K_SLOT + +# Whether consuming an operand of this kind frees its slot for reuse. +_is_recyclable(kind::UInt8) = kind == _K_SLOT + +# A stack descriptor is an Int64 packing a kind (low 2 bits) with a slot +# index; scalar descriptors carry no slot (their value lives in the scalar +# lane). `_feature_bit` is the feature's position in the `feature_mask` +# bitset of used features. +_pack_descriptor(kind::UInt8, slot::Integer=0) = Int64(kind) | (Int64(slot) << 2) +_descriptor_kind(descriptor::Int64) = UInt8(descriptor & 3) +_descriptor_slot(descriptor::Int64) = Int32(descriptor >> 2) +_feature_bit(feature::Integer) = UInt64(1) << (feature - 1) + +# Alloc-free stack of descriptor kinds for the planner: two bitmask lanes +# (bit 1 = top) mark `_K_SCALAR`/`_K_PSLOT`; neither lane set = `_K_SLOT`. +# Capacity is 64 entries. +struct KindStack + scalar::UInt64 + permanent::UInt64 +end + +function _push_kind(kinds::KindStack, kind::UInt8) + return KindStack( + (kinds.scalar << 1) | (kind == _K_SCALAR), + (kinds.permanent << 1) | (kind == _K_PSLOT), + ) +end +function _pop_kinds(kinds::KindStack, count::UInt8) + return KindStack(kinds.scalar >> count, kinds.permanent >> count) +end +function _args_all_scalar(kinds::KindStack, degree::UInt8) + arity_mask = (UInt64(1) << degree) - 1 + return (kinds.scalar & arity_mask) == arity_mask +end +function _count_recyclable_args(kinds::KindStack, degree::UInt8) + arity_mask = (UInt64(1) << degree) - 1 + return count_ones(~kinds.scalar & ~kinds.permanent & arity_mask) +end + +# Alloc-free pre-pass: record used features and simulate the descriptor stack +# to count slots. Makes the same kind/recycling decisions as the executor (the +# shared policy functions above), so the counts are exact. Trees deeper than +# 64 or features beyond 64 report failure and take the generic path. +function _plan_scratch(arena::Arena{T,D}) where {T,D} + nodes = arena.nodes + feature_mask = UInt64(0) + kinds = KindStack(0, 0) + stack_top = 0 + max_stack = 0 + num_live = 0 + num_free = 0 + max_live_intermediates = 0 + @inbounds for i in eachindex(nodes) + entry = nodes[i] + degree = entry.degree + if iszero(degree) + stack_top >= 64 && return (false, 0, 0, UInt64(0)) + stack_top += 1 + max_stack = max(max_stack, stack_top) + kind = _leaf_kind(entry) + if kind == _K_PSLOT + feature = entry.feature + (1 <= feature <= 64) || return (false, 0, 0, UInt64(0)) + feature_mask |= _feature_bit(feature) + end + kinds = _push_kind(kinds, kind) + else + result_kind = _op_result_kind(_args_all_scalar(kinds, degree)) + if _is_recyclable(result_kind) + num_free += _count_recyclable_args(kinds, degree) + if num_free > 0 + num_free -= 1 + else + num_live += 1 + max_live_intermediates = max(max_live_intermediates, num_live) + end + end + kinds = _push_kind(_pop_kinds(kinds, degree), result_kind) + stack_top -= degree - 1 + end + end + num_slots = count_ones(feature_mask) + max_live_intermediates + # A well-formed postfix tree collapses the stack to exactly the root; + # anything else (e.g. orphaned roots) must fail closed. + return (stack_top == 1, num_slots, max_stack, feature_mask) +end + +# `args` is `Tuple{T,Vararg{T}}` (not `NTuple{A,T}`) so that `T` stays bound +# for empty tuples (Aqua's unbound-args check); the arity comes from the +# tuple type itself. +@generated function _scalar_degn( + op_idx::UInt8, args::Tuple{T,Vararg{T}}, operators::O +) where {T,O<:OperatorEnum} + A = length(args.parameters) + nops = get_nops(O, Val(A)) + nops == 0 && return :(throw(ArgumentError("no operators of arity " * string($A)))) + return quote + Base.Cartesian.@nif( + $nops, + i -> i == op_idx, # COV_EXCL_LINE + i -> (Base.Cartesian.@ncall($A, operators.ops[$A][i], k -> args[k]))::T, + ) + end +end + +# Branchless arity-generic kernel: each operand selects per element between +# its scalar value and its pool slot (scalar operands carry offset 0), so an +# arity-A operator needs one kernel rather than 2^A variants. +@generated function _kern_n!( + pool::Matrix{T}, + dest_offset::Int, + op::F, + is_scalar::NTuple{A,Bool}, + scalar_args::NTuple{A,T}, + offsets::NTuple{A,Int}, + num_rows::Int, +) where {T,F,A} + quote + @inbounds @simd for j in 1:num_rows + pool[dest_offset + j] = Base.Cartesian.@ncall( + $A, op, k -> ifelse(is_scalar[k], scalar_args[k], pool[offsets[k] + j]) + ) + end + return nothing + end +end +# `is_valid_array` over a pool slot without constructing a view. +function _valid_slot(pool::Matrix{T}, offset::Int, num_rows::Int) where {T} + total = zero(T) + @inbounds @simd for j in 1:num_rows + total += pool[offset + j] + end + return is_valid(total) +end +_slot_offset(slot::Int32, nrows::Int) = (slot - 1) * nrows + +@generated function _dispatch_degn!( + ::Val{A}, + pool::Matrix{T}, + dest_offset::Int, + op_idx::UInt8, + is_scalar::NTuple{A,Bool}, + scalar_args::NTuple{A,T}, + offsets::NTuple{A,Int}, + nrows::Int, + operators::O, +) where {A,T,O<:OperatorEnum} + nops = get_nops(O, Val(A)) + nops == 0 && return :(throw(ArgumentError("no operators of arity " * string($A)))) + quote + Base.Cartesian.@nif( + $nops, + i -> i == op_idx, # COV_EXCL_LINE + i -> _kern_n!( + pool, + dest_offset, + operators.ops[$A][i], + is_scalar, + scalar_args, + offsets, + nrows, + ), + ) + return nothing + end +end + +# Loop-invariant evaluation state. +struct PlanState{T} + pool::Matrix{T} + descriptors::Vector{Int64} + scalar_vals::Vector{T} + free_base::Int + nrows::Int +end + +# Descriptor stack top, free-list length, and high-water slot, threaded +# through `_push_leaf!`/`_exec_op!`. +struct PlanRegisters + stack_top::Int + num_free::Int + next_slot::Int32 +end + +function _push_leaf!( + state::PlanState{T}, regs::PlanRegisters, entry::ArenaEntry{T}, feature_mask::UInt64 +) where {T} + stack_top = regs.stack_top + 1 + @inbounds if _leaf_kind(entry) == _K_SCALAR + state.descriptors[stack_top] = _pack_descriptor(_K_SCALAR) + state.scalar_vals[stack_top] = entry.val + else + feature_slot = _feature_slot(feature_mask, entry.feature) + state.descriptors[stack_top] = _pack_descriptor(_K_PSLOT, feature_slot) + end + return PlanRegisters(stack_top, regs.num_free, regs.next_slot) +end + +# Dispatch the runtime degree to a compile-time arity, then fold (all-scalar +# operands) or run the array kernel. Returns `(regs, ok)`. +@generated function _exec_op!( + state::PlanState{T}, + regs::PlanRegisters, + op_idx::UInt8, + degree::UInt8, + is_root::Bool, + early_exit::Bool, + operators::O, + ::Val{D}, +) where {T,O<:OperatorEnum,D} + quote + return Base.Cartesian.@nif( + $D, + A -> A == degree, # COV_EXCL_LINE + A -> _exec_op_arity!( + Val(A), state, regs, op_idx, is_root, early_exit, operators + ), + ) + end +end + +# Pop the top `A` operand descriptors; constant-fold if every operand is a +# scalar, otherwise run the kernel. +@generated function _exec_op_arity!( + ::Val{A}, + state::PlanState{T}, + regs::PlanRegisters, + op_idx::UInt8, + is_root::Bool, + early_exit::Bool, + operators::O, +) where {A,T,O<:OperatorEnum} + quote + (; descriptors, scalar_vals) = state + (; stack_top, num_free, next_slot) = regs + @inbounds begin + kinds = Base.Cartesian.@ntuple( + $A, k -> _descriptor_kind(descriptors[stack_top - $A + k]) + ) + idxs = Base.Cartesian.@ntuple( + $A, k -> _descriptor_slot(descriptors[stack_top - $A + k]) + ) + scalar_args = Base.Cartesian.@ntuple( + $A, k -> if kinds[k] == _K_SCALAR + scalar_vals[stack_top - $A + k] + else + zero(T) + end + ) + end + regs = PlanRegisters(stack_top - ($A - 1), num_free, next_slot) + all_args_scalar = Base.Cartesian.@nall($A, k -> kinds[k] == _K_SCALAR) + if _op_result_kind(all_args_scalar) == _K_SCALAR + return _fold_constant_args!(state, regs, op_idx, scalar_args, operators) + end + return _run_op_kernel!( + state, regs, op_idx, kinds, idxs, scalar_args, is_root, early_exit, operators + ) + end +end + +# Constant-fold an all-scalar operator at the (already popped) stack top. +# Like `dispatch_constant_tree`, operand values and the fold result are +# validated unconditionally; folded args are valid by induction, so the arg +# check only screens constant leaves. +function _fold_constant_args!( + state::PlanState{T}, + regs::PlanRegisters, + op_idx::UInt8, + scalar_args::NTuple{A,T}, + operators::OperatorEnum, +) where {A,T} + all(is_valid, scalar_args) || return (regs, false) + value = _scalar_degn(op_idx, scalar_args, operators) + is_valid(value) || return (regs, false) + @inbounds state.descriptors[regs.stack_top] = _pack_descriptor(_K_SCALAR) + @inbounds state.scalar_vals[regs.stack_top] = value + return (regs, true) +end + +# Recycle the freed argument slots, allocate the destination (slot 1 at the +# root), and dispatch the kernel. `early_exit` validation mirrors the generic +# evaluator at lower cost: scalar operands are checked at consumption (O(1)); +# slot operands are covered by checking each kernel output at production +# (every non-root intermediate is consumed exactly once, so this rejects the +# same trees as per-consumption checks); features are validated once at +# materialization; the root output is never checked, as in the generic +# evaluator. +@generated function _run_op_kernel!( + state::PlanState{T}, + regs::PlanRegisters, + op_idx::UInt8, + kinds::NTuple{A,UInt8}, + idxs::NTuple{A,Int32}, + scalar_args::NTuple{A,T}, + is_root::Bool, + early_exit::Bool, + operators::O, +) where {A,T,O<:OperatorEnum} + quote + (; pool, descriptors, free_base, nrows) = state + (; stack_top, num_free, next_slot) = regs + # free recyclable argument slots first; the destination may then reuse + # one (kernels are alias-safe: reads and writes of the same slot are + # at the same element index) + @inbounds Base.Cartesian.@nexprs( + $A, k -> if _is_recyclable(kinds[k]) + num_free += 1 + descriptors[free_base + num_free] = Int64(idxs[k]) + end + ) + if is_root + slot = Int32(1) + elseif num_free > 0 + slot = Int32(@inbounds(descriptors[free_base + num_free])) + num_free -= 1 + else + next_slot += Int32(1) + slot = next_slot + end + @inbounds descriptors[stack_top] = _pack_descriptor(_K_SLOT, slot) + dest_offset = _slot_offset(slot, nrows) + is_scalar = Base.Cartesian.@ntuple($A, k -> kinds[k] == _K_SCALAR) + offsets = Base.Cartesian.@ntuple( + $A, k -> kinds[k] == _K_SCALAR ? 0 : _slot_offset(idxs[k], nrows) + ) + regs = PlanRegisters(stack_top, num_free, next_slot) + scalars_valid = + !early_exit || + Base.Cartesian.@nall($A, k -> !is_scalar[k] || is_valid(scalar_args[k])) + scalars_valid || return (regs, false) + _dispatch_degn!( + Val($A), + pool, + dest_offset, + op_idx, + is_scalar, + scalar_args, + offsets, + nrows, + operators, + ) + if early_exit && !is_root && !_valid_slot(pool, dest_offset, nrows) + return (regs, false) + end + return (regs, true) + end +end + +# Copy each used feature column into its pinned pool slot (layout: 1 = +# output; 2 .. 1+num_features = features, ascending). Validity is checked +# once per feature here, replacing per-consumption checks: every +# materialized feature is consumed by some operator -- except in a +# single-leaf tree, where `check_validity` is passed as false to match +# `deg0_eval`, which never validates a bare leaf. +function _materialize_features!( + pool::Matrix{T}, cX::Matrix{T}, feature_mask::UInt64, nrows::Int, check_validity::Bool +) where {T} + slot = 1 + remaining = feature_mask + while !iszero(remaining) + feature = trailing_zeros(remaining) + 1 + slot += 1 + offset = (slot - 1) * nrows + @inbounds @simd for j in 1:nrows + pool[offset + j] = cX[feature, j] + end + # Separate validity pass over the just-written (cache-hot) slot keeps + # the copy loop a pure memcpy pattern. + check_validity && !_valid_slot(pool, offset, nrows) && return false + remaining &= remaining - 1 + end + return true +end + +# Land the result in pool row 1: copy a scalar or passthrough root into the +# output chunk if needed, then convert the contiguous chunk into the strided +# row the generic buffered evaluator returns (keeps `eval_tree_array` +# type stable). +function _write_root_to_output!( + pool::Matrix{T}, descriptors::Vector{Int64}, scalar_vals::Vector{T}, nrows::Int +) where {T} + # Root never went through a kernel (bare leaf or fully folded scalar), or + # an op-root wrote into a non-output slot via in-place deg1 reuse. A bare + # leaf root is never validity-checked (`deg0_eval` semantics); a folded + # scalar root is already valid by induction. + root_kind = _descriptor_kind(descriptors[1]) + root_slot = _descriptor_slot(descriptors[1]) + if root_kind == _K_SCALAR + value = scalar_vals[1] + @inbounds @simd for j in 1:nrows + pool[j] = value + end + elseif !isone(root_slot) + root_offset = _slot_offset(root_slot, nrows) + @inbounds @simd for j in 1:nrows + pool[j] = pool[root_offset + j] + end + end + # The chunk and row 1 overlap in memory; iterating downward is safe: when + # reading chunk index j, every already-written row position (j''-1)*B+1 + # with j'' > j exceeds j for B = size(pool, 1) >= 2, and for B == 1 the + # chunk and row coincide elementwise. + if size(pool, 1) > 1 + @inbounds for j in nrows:-1:1 + pool[1, j] = pool[j] + end + end + return nothing +end + +function _arena_eval( + arena::Arena{T,D}, + cX::Matrix{T}, + operators::OperatorEnum, + ::Val{early_exit}, + num_slots::Int, + max_stack::Int, + feature_mask::UInt64, + pool::Matrix{T}, +) where {T,D,early_exit} + nodes = arena.nodes + num_nodes = length(nodes) + nrows = size(cX, 2) + num_features = count_ones(feature_mask) + output = @view(pool[1, :]) + + check_features = early_exit && num_nodes > 1 + if !_materialize_features!(pool, cX, feature_mask, nrows, check_features) + return ResultOk(output, false) + end + + # Per-call descriptor state (tiny; the pool itself is caller-owned): + descriptors = Vector{Int64}(undef, max_stack + num_slots) + scalar_vals = Vector{T}(undef, max_stack) + state = PlanState(pool, descriptors, scalar_vals, max_stack, nrows) + regs = PlanRegisters(0, 0, Int32(1 + num_features)) + + @inbounds for i in 1:num_nodes + entry = nodes[i] + if iszero(entry.degree) + regs = _push_leaf!(state, regs, entry, feature_mask) + else + is_root = i == num_nodes + regs, ok = _exec_op!( + state, regs, entry.op, entry.degree, is_root, early_exit, operators, Val(D) + ) + ok || return ResultOk(output, false) + end + end + + _write_root_to_output!(pool, descriptors, scalar_vals, nrows) + return ResultOk(output, true) +end + +end diff --git a/src/DynamicExpressions.jl b/src/DynamicExpressions.jl index 355e7b98..ff6736c7 100644 --- a/src/DynamicExpressions.jl +++ b/src/DynamicExpressions.jl @@ -12,6 +12,8 @@ using DispatchDoctor: @stable, @unstable include("NodePreallocation.jl") include("Strings.jl") include("Evaluate.jl") + include("ArenaNode.jl") + include("ArenaNodeEval.jl") include("EvaluateDerivative.jl") include("ChainRules.jl") include("EvaluationHelpers.jl") @@ -51,6 +53,7 @@ import .ValueInterfaceModule: filter_map, filter_map! import .NodePreallocationModule: allocate_container, copy_into! +import .ArenaNodeModule: ArenaNode, Arena import .NodeModule: constructorof, with_type_parameters, diff --git a/src/Interfaces.jl b/src/Interfaces.jl index 53e5c65d..f1b64b98 100644 --- a/src/Interfaces.jl +++ b/src/Interfaces.jl @@ -32,6 +32,7 @@ using ..NodeModule: set_node!, filter_map, filter_map! +using ..ArenaNodeModule: ArenaNode using ..NodeUtilsModule: NodeIndex, is_node_constant, @@ -450,6 +451,11 @@ ni_description = ( ParametricNode, [Arguments()] ) +@implements( + NodeInterface{all_ni_methods_except(())}, + ArenaNode, + [Arguments()] +) #! format: on diff --git a/src/Node.jl b/src/Node.jl index 63b0a159..1340ee87 100644 --- a/src/Node.jl +++ b/src/Node.jl @@ -210,7 +210,7 @@ end Return the `i`-th child of a node (1-indexed). """ -@inline function get_child(n::AbstractNode{D}, i::Int) where {D} +@inline function get_child(n::AbstractNode{D}, i::Integer) where {D} return unsafe_get_children(n)[i][] end diff --git a/test/Project.toml b/test/Project.toml index e4b346d5..5cf7e62b 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -9,6 +9,8 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" Interfaces = "85a1e053-f937-4924-92a5-1367d23b7b87" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +LoweredCodeUtils = "6f1432cf-f94c-5a45-995e-cdbf5db27b0b" +LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890" Optim = "429524aa-4258-5aef-a3af-852621145aeb" NLSolversBase = "d41bc354-129a-5804-8e4c-c37616107c6c" @@ -26,8 +28,13 @@ TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] -Aqua = "0.8" -JET = "0.9, 0.10" +Aqua = "0.8.16" +JET = "0.9, 0.10, 0.11" +LoweredCodeUtils = "2, 3.4 - 3.5.99" +LineSearches = "7 - 7.4" +LoopVectorization = "0.12" +Supposition = "0.3.5" +SymbolicUtils = "4.35" TestItems = "1" TestItemRunner = "1" diff --git a/test/runtests.jl b/test/runtests.jl index fa16d575..7bad7782 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 new file mode 100644 index 00000000..c87ae2fa --- /dev/null +++ b/test/test_arenanode.jl @@ -0,0 +1,889 @@ +@testmodule ArenaTreeGen begin + using DynamicExpressions + using Random: AbstractRNG + + export random_tree + + function random_leaf(rng, ::Type{T}, ::Val{D}, nfeat, const_p) where {T,D} + if rand(rng) < const_p + return Node{T,D}(; val=randn(rng, T)) + else + return Node{T,D}(; feature=rand(rng, 1:nfeat)) + end + end + + """Grow a random `Node{T,D}` tree to `n` nodes. `arity_cdf[d]` is the + cumulative probability of expanding a leaf to degree `d`; `nops[d]` is the + number of degree-`d` operators in the testitem's `OperatorEnum`.""" + function random_tree( + rng::AbstractRNG, + n; + T=Float64, + D=2, + nfeat=3, + nops=(2, 4), + arity_cdf=(0.3, 1.0), + const_p=0.5, + ) + leaf() = random_leaf(rng, T, Val(D), nfeat, const_p) + tree = leaf() + while count_nodes(tree) < n + node = rand(rng, filter(t -> t.degree == 0, tree)) + r = rand(rng) + d = something(findfirst(>=(r), arity_cdf), D) + node.degree = d + node.op = rand(rng, 1:nops[d]) + set_children!(node, ntuple(_ -> leaf(), d)) + node.constant = false + node.val = zero(T) + end + return tree + end +end + +@testitem "ArenaNode interface and evaluation" begin + using Test + using DynamicExpressions + using DynamicExpressions: NodeInterface + using Interfaces: Interfaces + + using DynamicExpressions: ArenaNode, Arena + using DynamicExpressions.ArenaNodeModule: + ArenaEntry, _replace, is_compact_root, push_constant!, push_feature! + + operators = OperatorEnum(1 => (sin, cos), 2 => (+, *)) + x1 = Node{Float64}(; feature=1) + tree = sin(x1) + x1 * 3.2 + atree = convert(ArenaNode{Float64}, tree) + + @test atree isa ArenaNode{Float64,2} + @test count_nodes(atree) == count_nodes(tree) + @test string_tree(atree, operators) == string_tree(tree, operators) + @test tree_mapreduce(_ -> 1, +, atree, Int) == count_nodes(atree) + + @test Interfaces.test( + NodeInterface, + ArenaNode, + [ + atree, + convert(ArenaNode{Float64}, sin(x1)), + convert(ArenaNode{Float64}, x1), + convert(ArenaNode{Float64}, Node{Float64}(; val=1.0)), + ], + ) + + if atree.degree != 0 + cs = DynamicExpressions.NodeModule.unsafe_get_children(atree) + @test cs isa NTuple{2,DynamicExpressions.Nullable{typeof(atree)}} + @test length(get_children(atree, atree.degree)) == atree.degree + @test get_child(tree, UInt8(1)) == get_child(tree, 1) + @test get_child(atree, UInt8(1)) == get_child(atree, 1) + end + + collected = collect(atree; break_sharing=Val(true)) + @test !isempty(collected) && collected[1].idx == atree.idx + + X = randn(Float64, 1, 50) + y_tree, ok_tree = eval_tree_array(tree, X, operators) + y_atree, ok_atree = eval_tree_array(atree, X, operators) + @test ok_tree + @test ok_atree + @test y_tree ≈ y_atree + + const_nodes = filter(t -> t.degree == 0 && t.constant, atree) + @test !isempty(const_nodes) + const_nodes[1].val = 10.0 + y_mut, ok_mut = eval_tree_array(atree, X, operators) + @test ok_mut + @test !(y_mut ≈ y_tree) + + atree2 = copy(atree) + @test atree2 == atree + const_nodes2 = filter(t -> t.degree == 0 && t.constant, atree2) + const_nodes2[1].val = -5.0 + @test atree2 != atree + + tree2 = convert(Node, atree) + y_tree2, ok_tree2 = eval_tree_array(tree2, X, operators) + @test ok_tree2 + @test y_tree2 ≈ y_mut +end + +@testitem "ArenaNode mutation and simplification" begin + using Test + using DynamicExpressions + + using DynamicExpressions: ArenaNode, Arena + using DynamicExpressions.ArenaNodeModule: + ArenaEntry, _replace, is_compact_root, push_constant!, push_feature! + + operators = OperatorEnum(1 => (sin, cos), 2 => (+, *)) + x1 = Node{Float64}(; feature=1) + tree = sin(x1) + x1 * 3.2 + X = randn(Float64, 1, 50) + + atree_setnode = convert(ArenaNode{Float64}, tree) + atree_setnode2 = copy(atree_setnode) + set_node!(atree_setnode, atree_setnode2) + @test string_tree(atree_setnode, operators) == string_tree(atree_setnode2, operators) + + default = ArenaNode{Float64,2}() + @test default.degree == 0 + @test default.constant + @test default.val == 0.0 + default.val = 2 + default.feature = 1 + default.op = 2 + default.constant = false + default.degree = 0 + @test default.val == 2.0 + @test default.feature == 1 + @test default.op == 2 + @test !default.constant + @test_throws ArgumentError default.foo = 1 + + parent = convert(ArenaNode{Float64}, sin(x1)) + other = convert(ArenaNode{Float64}, x1 * 3.2) + set_child!(parent, other, 1) + @test get_child(parent, 1).arena === parent.arena + other.r.val = 99.0 + y_parent, ok_parent = eval_tree_array(parent, X, operators) + @test ok_parent + @test y_parent ≈ sin.(X[1, :] .* 3.2) + + guarded = convert(ArenaNode{Float64}, sin(x1)) + foreign_child = convert(ArenaNode{Float64}, x1 * 3.2) + n_before = length(guarded.arena.nodes) + @test_throws BoundsError set_child!(guarded, foreign_child, 3) + @test length(guarded.arena.nodes) == n_before + + @test_throws ArgumentError set_child!(parent, Node{Float32}(; val=1.0f0), 1) + @test_throws UndefRefError get_child(convert(ArenaNode{Float64}, x1), 1) + + rewritten = convert(ArenaNode{Float64}, sin(x1)) + set_children!(rewritten, (convert(ArenaNode{Float64}, x1 * 2.0),)) + @test get_child(rewritten, 1).arena === rewritten.arena + @test string_tree(rewritten, operators) == "sin(x1 * 2.0)" + + bad_children = ( + DynamicExpressions.Nullable(true, Node{Float64}(; val=0.0)), + Node{Float32}(; val=1.0f0), + ) + @test_throws ArgumentError set_children!(rewritten, bad_children) + + tree_fold = Node{Float64}(; val=2.0) + Node{Float64}(; val=3.0) + atree_fold = convert(ArenaNode{Float64}, tree_fold) + simplify_tree!(atree_fold, operators) + @test atree_fold.degree == 0 + @test atree_fold.constant + @test atree_fold.val == 5.0 +end + +@testitem "Expression with ArenaNode" begin + using Test + using DynamicExpressions + using DynamicExpressions: ExpressionInterface, get_tree + using Interfaces: test + + using DynamicExpressions: ArenaNode, Arena + using DynamicExpressions.ArenaNodeModule: + ArenaEntry, _replace, is_compact_root, push_constant!, push_feature! + + operators = OperatorEnum(1 => (sin, cos), 2 => (+, *)) + x1 = Node{Float64}(; feature=1) + atree = convert(ArenaNode{Float64}, sin(x1) + x1 * 3.2) + expr = Expression(atree; operators, variable_names=["x"]) + @test get_tree(expr) === atree + @test test(ExpressionInterface, Expression, [expr]) + + simple_expr = Expression( + convert(ArenaNode{Float64}, x1); operators, variable_names=["x"] + ) + @test test(ExpressionInterface, Expression, [simple_expr]) +end + +@testitem "ArenaNode derivatives through Expression" begin + using Test + using DynamicExpressions + using DynamicExpressions: eval_grad_tree_array, extract_gradient + using DifferentiationInterface: AutoZygote, gradient + using Zygote + + using DynamicExpressions: ArenaNode, Arena + using DynamicExpressions.ArenaNodeModule: + ArenaEntry, _replace, is_compact_root, push_constant!, push_feature! + + operators_grad = OperatorEnum(1 => (sin, cos, exp), 2 => (+, -, *, /)) + x1 = Expression( + convert(ArenaNode{Float64}, Node{Float64}(; feature=1)); + operators=operators_grad, + variable_names=[:x1, :x2], + ) + x2 = Expression( + convert(ArenaNode{Float64}, Node{Float64}(; feature=2)); + operators=operators_grad, + variable_names=[:x1, :x2], + ) + expr_grad = sin(2.0 * x1 + exp(x2 + 5.0)) + + Xg = rand(Float64, 2, 10) .+ 1 + expected = @. sin(2.0 * Xg[1, :] + exp(Xg[2, :] + 5.0)) + expected_dy_dx1 = @. 2.0 * cos(2.0 * Xg[1, :] + exp(Xg[2, :] + 5.0)) + + result, ok = eval_tree_array(expr_grad, Xg) + @test ok + @test result ≈ expected + + _, grad2, ok2 = eval_grad_tree_array(expr_grad, Xg; variable=Val(true)) + @test ok2 + @test grad2[1, :] ≈ expected_dy_dx1 + + grad_zygote = expr_grad'(Xg) + @test grad_zygote[1, :] ≈ expected_dy_dx1 + + operators_const = OperatorEnum(2 => (+,)) + x1c = Expression( + convert(ArenaNode{Float64}, Node{Float64}(; feature=1)); + operators=operators_const, + variable_names=["x1"], + ) + expr_const = x1c + 1.5 + _, grad3, ok3 = eval_grad_tree_array(expr_const, ones(1, 5); variable=Val(false)) + @test ok3 + @test grad3[1, :] ≈ fill(1.0, 5) + + d_ex = gradient(AutoZygote(), expr_const) do ex + return sum(ex(ones(1, 5))) + end + @test extract_gradient(d_ex, expr_const) ≈ [5.0] +end + +@testitem "ArenaNode allocations" begin + include(joinpath(@__DIR__, "test_arenanode_allocations.jl")) +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! + + using DynamicExpressions: ArenaNode, Arena + using DynamicExpressions.ArenaNodeModule: + ArenaEntry, _replace, is_compact_root, push_constant!, push_feature! + + 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(ArenaNode{Float64}, tree) + + @testset "compact flat copy" begin + @test 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 !is_compact_root(sub) + csub = copy(sub) + @test is_compact_root(csub) + @test convert(Node, csub) == tree.l + end + + @testset "structural mutation invalidates fast paths" begin + mutated = convert(ArenaNode{Float64}, tree) + set_child!(mutated, convert(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) + @test has_constants(mutated) == has_constants(expected) + recompacted = copy(mutated) + @test is_compact_root(recompacted) + @test count_nodes(recompacted) == count_nodes(expected) + + leafed = convert(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) + @test has_constants(leafed) == has_constants(expected2) + end + + @testset "preallocated copy_into!" begin + dest = allocate_container(atree) + ref = Ref(-1) + out = copy_into!(dest, atree; ref) + @test out.arena === dest + @test convert(Node, out) == tree + @test ref[] == length(atree) + out2 = copy_into!(dest, atree; ref) + @test convert(Node, out2) == tree + @test ref[] == length(atree) + + same_ref = Ref(-1) + same = copy_into!(out2.arena, out2; ref=same_ref) + @test same === out2 + @test same_ref[] == length(out2) + 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(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(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(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)) + @test set_scalar_constants!(fresh, vals .* 2, refs) === fresh + @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 also use arena-index refs: + sub = fresh.l + vsub, rsub = get_scalar_constants(sub) + @test rsub isa Vector{Int32} + @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 + +@testitem "Arena array interface guards compactness" begin + using DynamicExpressions + using DynamicExpressions: Node + + using DynamicExpressions: ArenaNode, Arena + using DynamicExpressions.ArenaNodeModule: + ArenaEntry, _replace, is_compact_root, push_constant!, push_feature! + + operators = OperatorEnum(; binary_operators=[+, *], unary_operators=[sin]) + x1 = Node{Float64}(; feature=1) + tree = sin(x1 * 3.2) + 0.5 + atree = convert(ArenaNode{Float64}, tree) + a = atree.arena + @test a isa AbstractVector{ArenaEntry{Float64,2}} + @test length(a) == count_nodes(tree) + + e = a[1] + a[1] = _replace(e; val=42.0) + @test is_compact_root(atree) + + vals, refs = get_scalar_constants(atree) + set_scalar_constants!(atree, vals .* 2, refs) + @test is_compact_root(atree) + + bi = findfirst(e -> e.degree == 0x02, collect(a)) + e = a[bi] + a[bi] = _replace(e; degree=0x00) + @test !a.compact[] + + b = convert(ArenaNode{Float64}, tree).arena + scrambled = reverse(collect(b)) + copyto!(b, scrambled) + @test !b.compact[] + + c = convert(ArenaNode{Float64}, tree).arena + copyto!(c, collect(c)) + @test c.compact[] +end + +@testitem "ArenaNode fast paths match Node" setup = [ArenaTreeGen] begin + using DynamicExpressions + using DynamicExpressions: Node + using Random + + using DynamicExpressions: ArenaNode, Arena + using DynamicExpressions.ArenaNodeModule: + ArenaEntry, _replace, is_compact_root, push_constant!, push_feature! + + operators = OperatorEnum(; binary_operators=[+, -, *, /], unary_operators=[sin, cos]) + rng = MersenneTwister(42) + + function all_facades(n, acc=ArenaNode{Float64,2}[]) + push!(acc, n) + for i in 1:(n.degree) + all_facades(get_child(n, i), acc) + end + return acc + end + + for _ in 1:20 + root = convert(ArenaNode{Float64}, ArenaTreeGen.random_tree(rng, rand(rng, 5:25))) + for _ in 1:8 + nodes = all_facades(root) + node = rand(rng, nodes) + choice = rand(rng, 1:4) + if choice == 1 && node.degree == 0 + if node.constant + node.val = randn(rng) + else + node.feature = rand(rng, 1:3) + end + elseif choice == 2 && node.degree > 0 + node.op = rand(rng, 1:(node.degree == 1 ? 2 : 4)) + elseif choice == 3 && node.degree > 0 + set_child!( + node, + convert( + ArenaNode{Float64}, ArenaTreeGen.random_tree(rng, rand(rng, 1:5)) + ), + rand(rng, 1:(node.degree)), + ) + else + node.degree = 0 + node.constant = true + node.val = randn(rng) + end + + expected = convert(Node, root) + @test count_nodes(root) == count_nodes(expected) + @test count_constant_nodes(root) == count_constant_nodes(expected) + @test has_constants(root) == has_constants(expected) + @test DynamicExpressions.NodeUtilsModule.is_constant(root) == + DynamicExpressions.NodeUtilsModule.is_constant(expected) + @test count(t -> t.degree == 2, root) == count(t -> t.degree == 2, expected) + @test first(get_scalar_constants(root)) == first(get_scalar_constants(expected)) + c = copy(root) + @test is_compact_root(c) + @test convert(Node, c) == expected + end + end +end + +@testitem "ArenaNode buffered plan eval" setup = [ArenaTreeGen] begin + using DynamicExpressions + using DynamicExpressions: Node, EvalOptions, ArrayBuffer + using Random + + using DynamicExpressions: ArenaNode, Arena + using DynamicExpressions.ArenaNodeModule: + ArenaEntry, _replace, is_compact_root, push_constant!, push_feature! + + operators = OperatorEnum(; binary_operators=[+, *, /, -], unary_operators=[cos, exp]) + rng = MersenneTwister(11) + + for T in (Float32, Float64) + X = randn(rng, T, 5, 37) + buf_a = zeros(T, 40, 37) + buf_n = zeros(T, 40, 37) + for trial in 1:60, early_exit in (true, false) + tree = ArenaTreeGen.random_tree(rng, rand(rng, 1:30); T, nfeat=5, const_p=0.4) + atree = convert(ArenaNode{T}, tree) + opts_a = EvalOptions(; early_exit, buffer=ArrayBuffer(buf_a, Ref(0))) + opts_n = EvalOptions(; early_exit, buffer=ArrayBuffer(buf_n, Ref(0))) + # Unbuffered ground truth; results of buffered evals are views, so + # copy before they can alias each other. + rt = try + ( + eval_tree_array( + tree, X, operators; eval_options=EvalOptions(; early_exit) + ), + false, + ) + catch + (nothing, true) + end + ra = try + (eval_tree_array(atree, X, operators; eval_options=opts_a), false) + catch + (nothing, true) + end + @test rt[2] == ra[2] + (rt[2] || ra[2]) && continue + (yref, okref) = rt[1] + (ya, oka) = ra[1] + @test okref == oka + if okref + @test yref ≈ ya || (any(!isfinite, yref) && any(!isfinite, ya)) + end + # buffered Node evaluation agrees too + (yn, okn) = eval_tree_array(tree, X, operators; eval_options=opts_n) + @test okn == okref + end + + # Stacks deeper than 64 take the generic path (with the same buffer): + deep = Node{T}(; val=T(0.5)) + for _ in 1:70 + deep = Node{T}(; op=1, l=Node{T}(; feature=1), r=deep) + end + adeep = convert(ArenaNode{T}, deep) + big = zeros(T, 80, 37) + o = EvalOptions(; buffer=ArrayBuffer(big, Ref(0))) + y1, ok1 = eval_tree_array(copy(deep), X, operators) + y2, ok2 = eval_tree_array(adeep, X, operators; eval_options=o) + @test ok1 == ok2 + ok1 && @test y1 ≈ y2 + end +end + +@testitem "ArenaNode buffered eval, degree 3" setup = [ArenaTreeGen] begin + using DynamicExpressions + using DynamicExpressions: Node, EvalOptions, ArrayBuffer + using Random + + using DynamicExpressions: ArenaNode, Arena + using DynamicExpressions.ArenaNodeModule: + ArenaEntry, _replace, is_compact_root, push_constant!, push_feature! + + my3(x, y, z) = x * y + z + operators = OperatorEnum(1 => (cos, exp), 2 => (+, *, -, /), 3 => (fma, my3)) + rng = MersenneTwister(5) + T = Float64 + + X = randn(rng, T, 5, 29) + buf = zeros(T, 48, 29) + for trial in 1:40, early_exit in (true, false) + tree = ArenaTreeGen.random_tree( + rng, + rand(rng, 1:25); + D=3, + nfeat=5, + nops=(2, 4, 2), + arity_cdf=(0.25, 0.6, 1.0), + const_p=0.4, + ) + atree = convert(ArenaNode{T,3}, tree) + o = EvalOptions(; early_exit, buffer=ArrayBuffer(buf, Ref(0))) + rt = try + ( + eval_tree_array( + copy(tree), X, operators; eval_options=EvalOptions(; early_exit) + ), + false, + ) + catch + (nothing, true) + end + ra = try + (eval_tree_array(atree, X, operators; eval_options=o), false) + catch + (nothing, true) + end + @test rt[2] == ra[2] + (rt[2] || ra[2]) && continue + (yref, okref) = rt[1] + (ya, oka) = ra[1] + @test okref == oka + if okref + @test yref ≈ ya || (any(!isfinite, yref) && any(!isfinite, ya)) + end + end +end + +@testitem "ArenaNode review regressions" begin + using DynamicExpressions + using DynamicExpressions: Node, EvalOptions, ArrayBuffer + using DynamicExpressions.NodePreallocationModule: allocate_container, copy_into! + + using DynamicExpressions: ArenaNode, Arena + using DynamicExpressions.ArenaNodeModule: + ArenaEntry, _replace, is_compact_root, push_constant!, push_feature! + + operators = OperatorEnum(1 => (cos, exp), 2 => (+, *, -, /)) + T = Float64 + X = T[1.0 2.0 3.0; 0.5 1.5 2.5] + nrows = size(X, 2) + buf() = ArrayBuffer(zeros(T, 16, nrows), Ref(0)) + to_arena(t) = convert(ArenaNode{T,2}, t) + x1 = Node{T}(; feature=1) + x2 = Node{T}(; feature=2) + + @testset "multi-root arenas fail closed" begin + a = Arena{T,2}() + push_constant!(a, 1.0) + i2 = push_feature!(a, 2) + n2 = ArenaNode(a, i2) + @test !is_compact_root(n2) + @test count_nodes(n2) == 1 + y, ok = eval_tree_array(n2, X, operators; eval_options=EvalOptions(; buffer=buf())) + @test ok && y ≈ X[2, :] + end + + @testset "set_scalar_constants! converts the eltype" begin + small = to_arena(Node{T}(; op=1, l=x1, r=Node{T}(; val=9.0))) + _, refs = get_scalar_constants(small) + set_scalar_constants!(small, Float32[2.5], refs) + @test any(n -> n.degree == 0 && n.constant && n.val == 2.5, small) + end + + @testset "out-of-arity get_child throws" begin + leaf = ArenaNode{T,1}() + @test_throws BoundsError leaf.r + end + + @testset "half-built nodes throw on traversal" begin + half = to_arena(x1) + half.degree = 1 + half.op = 1 + @test_throws UndefRefError DynamicExpressions.NodeUtilsModule.is_constant(half) + @test_throws UndefRefError any(_ -> false, half) + end + + function check_parity(tree, ops, Xm; early_exit) + atree = convert(ArenaNode{T,2}, tree) + yn, okn = eval_tree_array( + copy(tree), Xm, ops; eval_options=EvalOptions(; early_exit) + ) + o = EvalOptions(; early_exit, buffer=ArrayBuffer(zeros(T, 16, size(Xm, 2)), Ref(0))) + ya, oka = eval_tree_array(atree, Xm, ops; eval_options=o) + @test okn == oka + if okn + @test yn ≈ ya || (any(!isfinite, yn) && any(!isfinite, ya)) + end + end + + @testset "ok-flag parity with the generic evaluator" begin + nanleaf = Node{T}(; val=NaN) + infdiv = Node{T}(; op=4, l=Node{T}(; val=1.0), r=Node{T}(; val=0.0)) + for early_exit in (true, false) + check_parity(nanleaf, operators, X; early_exit) + check_parity(Node{T}(; op=1, l=x1, r=nanleaf), operators, X; early_exit) + check_parity(Node{T}(; op=1, l=x1, r=infdiv), operators, X; early_exit) + check_parity( + Node{T}(; op=1, l=x1, r=Node{T}(; val=Inf)), operators, X; early_exit + ) + grow = Node{T}(; + op=2, + l=Node{T}(; op=2, l=x1, r=Node{T}(; val=1e300)), + r=Node{T}(; val=1e300), + ) + check_parity(Node{T}(; op=1, l=grow, r=x2), operators, X; early_exit) # Inf intermediate + check_parity(grow, operators, X; early_exit) # Inf at the root + end + + nanclean(x) = ifelse(isnan(x), zero(x), x) + ops2 = OperatorEnum(1 => (nanclean,), 2 => (+, *)) + Xnan = copy(X) + Xnan[1, 2] = NaN + for early_exit in (true, false) + check_parity(copy(x1), ops2, Xnan; early_exit) # bare feature root, NaN input + check_parity(Node{T}(; op=1, l=x1), ops2, Xnan; early_exit) # NaN input absorbed + check_parity(Node{T}(; op=1, l=Node{T}(; val=NaN)), ops2, X; early_exit) # NaN leaf folded away + end + end + + @testset "use_fused=false takes the generic path" begin + t = to_arena(Node{T}(; op=1, l=x1, r=Node{T}(; op=2, l=x2, r=Node{T}(; val=3.0)))) + o_nofuse = EvalOptions(; buffer=buf(), use_fused=Val(false)) + y1, ok1 = eval_tree_array(t, X, operators; eval_options=o_nofuse) + @test ok1 && o_nofuse.buffer.index[] > 0 # generic buffer protocol engaged + o_plan = EvalOptions(; buffer=buf()) + y2, ok2 = eval_tree_array(t, X, operators; eval_options=o_plan) + @test ok2 && o_plan.buffer.index[] == 0 # plan path bypasses the index + @test y1 ≈ y2 + end + + @testset "cross-representation ==" begin + tn = Node{T}(; op=1, l=x1, r=Node{T}(; op=2, l=x2, r=Node{T}(; val=0.5))) + ta = to_arena(tn) + @test ta == tn + @test tn == ta + @test convert(ArenaNode{Float32,2}, tn) == tn + end + + @testset "copy_into! container reuse" begin + t = to_arena(Node{T}(; op=1, l=x1, r=Node{T}(; val=2.0))) + c = allocate_container(t) + t2 = copy_into!(c, t) + t3 = copy_into!(c, t2) + @test string_tree(t3, operators) == string_tree(t, operators) + end + + @testset "non-isbits eltype takes the generic path safely" begin + tb = Node{BigFloat,2}(; + op=1, l=Node{BigFloat,2}(; feature=1), r=Node{BigFloat,2}(; val=big"1.5") + ) + ab = convert(ArenaNode{BigFloat,2}, tb) + Xb = BigFloat.(X) + bufb = ArrayBuffer(Matrix{BigFloat}(undef, 16, nrows), Ref(0)) + yb, okb = eval_tree_array( + ab, Xb, operators; eval_options=EvalOptions(; buffer=bufb) + ) + @test okb && yb ≈ Xb[1, :] .+ big"1.5" + end +end + +@testitem "ArenaNode supposition invariants" begin + using Test + using Supposition + using Supposition: @check, Data + using DynamicExpressions + using DynamicExpressions: Node, EvalOptions, ArrayBuffer, get_tree + + using DynamicExpressions: ArenaNode + + include("supposition_utils.jl") + + const T = Float64 + const N_FEATURES = 5 + const OPERATORS = OperatorEnum(1 => (abs, cos), 2 => (+, -, *, /)) + + expr_gen = make_expression_generator( + T; num_features=N_FEATURES, max_layers=8, operators=OPERATORS + ) + tree_gen = map(get_tree, expr_gen) + input_gen = make_input_matrix_generator(T; n_features=N_FEATURES) + + # Round-trip and read-only properties are representation-independent. + roundtrip = @check function arena_roundtrip(tree=tree_gen) + atree = convert(ArenaNode{T,2}, tree) + back = convert(Node, atree) + return back == tree && + string_tree(atree, OPERATORS) == string_tree(tree, OPERATORS) && + count_nodes(atree) == count_nodes(tree) && + count_constant_nodes(atree) == count_constant_nodes(tree) && + count_depth(atree) == count_depth(tree) && + hash(atree) == hash(tree) && + has_constants(atree) == has_constants(tree) && + has_operators(atree) == has_operators(tree) && + copy(atree) == atree + end + @test something(roundtrip.result) isa Supposition.Pass + + function evals_match(tree, atree, X) + # Unbuffered ArenaNode eval takes the same generic path as Node, so + # both the values and the `ok` flag must match exactly. + yn, okn = eval_tree_array(copy(tree), X, OPERATORS) + ya, oka = eval_tree_array(atree, X, OPERATORS) + okn == oka || return false + okn && !(yn ≈ ya) && return false + # The buffered plan path's `ok` is best-effort and may differ from the + # generic evaluator (which fuses some shapes without materializing the + # intermediate this path validates; the `is_valid(sum(...))` check can + # also overflow on finite-but-huge values). Values must agree whenever + # both sides report ok. + buffer = ArrayBuffer(zeros(T, 64, size(X, 2)), Ref(0)) + yb, okb = eval_tree_array(atree, X, OPERATORS; eval_options=EvalOptions(; buffer)) + okb && okn && !(yb ≈ yn) && return false + return true + end + + evals = @check function arena_eval_matches_node(tree=tree_gen, X=input_gen) + return evals_match(tree, convert(ArenaNode{T,2}, tree), X) + end + @test something(evals.result) isa Supposition.Pass + + # Valid mutations applied identically to both representations keep them + # equivalent. Mutations are specified positionally (preorder index). + # a single multi-arg map (no Data.OneOf / generator unions, which older + # Supposition versions in the downgrade-compat CI job cannot handle) + mutation_gen = map( + (kind, i, value, op, feature) -> if kind == 1 + (:set_val, i, value) + elseif kind == 2 + (:set_op, i, op) + elseif kind == 3 + (:to_leaf, i, value) + else + (:to_feature, i, feature) + end, + Data.SampledFrom(1:4), + Data.Integers(1, 64), + Data.Floats{T}(; nans=false, infs=false), + Data.Integers(1, 4), + Data.Integers(1, N_FEATURES), + ) + + function apply_mutation!(tree, (kind, i, x)) + nodes = collect(tree) + node = nodes[mod1(i, length(nodes))] + if kind == :set_val + if node.degree == 0 && node.constant + node.val = x + end + elseif kind == :set_op + if node.degree == 2 + node.op = mod1(x, 4) + elseif node.degree == 1 + node.op = mod1(x, 2) + end + elseif kind == :to_leaf + node.degree = 0 + node.constant = true + node.val = x + elseif kind == :to_feature + node.degree = 0 + node.constant = false + node.feature = x + end + return tree + end + + mutations = @check function arena_mutation_matches_node( + tree0=tree_gen, X=input_gen, muts=Data.Vectors(mutation_gen; min_size=1, max_size=5) + ) + tree = copy(tree0) + atree = convert(ArenaNode{T,2}, tree0) + for mut in muts + apply_mutation!(tree, mut) + apply_mutation!(atree, mut) + end + return string_tree(atree, OPERATORS) == string_tree(tree, OPERATORS) && + count_nodes(atree) == count_nodes(tree) && + convert(Node, atree) == tree && + evals_match(tree, atree, X) + end + @test something(mutations.result) isa Supposition.Pass + + # Re-compacting after mutation (via copy) preserves the tree and re-enables + # the flat fast paths. + recompact = @check function arena_copy_recompacts( + tree0=tree_gen, X=input_gen, muts=Data.Vectors(mutation_gen; min_size=1, max_size=3) + ) + tree = copy(tree0) + atree = convert(ArenaNode{T,2}, tree0) + for mut in muts + apply_mutation!(tree, mut) + apply_mutation!(atree, mut) + end + compacted = copy(atree) + return DynamicExpressions.ArenaNodeModule.is_compact_root(compacted) && + compacted == atree && + hash(compacted) == hash(atree) && + evals_match(tree, compacted, X) + end + @test something(recompact.result) isa Supposition.Pass +end diff --git a/test/test_arenanode_allocations.jl b/test/test_arenanode_allocations.jl new file mode 100644 index 00000000..99d8c093 --- /dev/null +++ b/test/test_arenanode_allocations.jl @@ -0,0 +1,85 @@ +using Test +using DynamicExpressions +using DynamicExpressions.NodePreallocationModule: allocate_container, copy_into! + +using DynamicExpressions: ArenaNode, Arena +using DynamicExpressions.ArenaNodeModule: _copy_to_arena!, push_constant! + +operators = OperatorEnum(1 => (sin, cos), 2 => (+, *)) +x1 = DynamicExpressions.Node{Float64}(; feature=1) + +function alloc_push_constant!(arena) + push_constant!(arena, 1.0) + return nothing +end + +function alloc_set_child!(parent, child) + set_child!(parent, child, 1) + return nothing +end + +function alloc_copy_tree!(arena, tree) + _copy_to_arena!(arena, tree) + return nothing +end + +function alloc_copy_into!(dest, tree) + copy_into!(dest, tree) + return nothing +end + +arena_push = Arena{Float64,2}(; capacity=128) + +base_tree = sin(x1) +parent_arena = Arena{Float64,2}(; capacity=128) +parent_idx = _copy_to_arena!(parent_arena, base_tree) +parent = ArenaNode(parent_arena, parent_idx) + +child_tree = x1 * 3.2 +child_arena = Arena{Float64,2}(; capacity=128) +child_idx = _copy_to_arena!(child_arena, child_tree) +child = ArenaNode(child_arena, child_idx) + +tree_large = sin(x1) + x1 * 3.2 + cos(x1) +atree_large = convert(ArenaNode{Float64}, tree_large) +copy_dest = allocate_container(atree_large) +arena_large = Arena{Float64,2}(; capacity=128) + +for _ in 1:5 + alloc_push_constant!(arena_push) + alloc_set_child!(parent, child) + alloc_copy_tree!(arena_large, tree_large) + alloc_copy_into!(copy_dest, atree_large) +end + +arena_push_nodes = arena_push.nodes +arena_push_ptr = pointer(arena_push_nodes) +arena_push_len = length(arena_push_nodes) +alloc_push_constant!(arena_push) + +parent_nodes = parent_arena.nodes +parent_ptr = pointer(parent_nodes) +parent_len = length(parent_nodes) +alloc_set_child!(parent, child) + +arena_large_nodes = arena_large.nodes +arena_large_ptr = pointer(arena_large_nodes) +arena_large_len = length(arena_large_nodes) +alloc_copy_tree!(arena_large, tree_large) + +copy_dest_nodes = copy_dest.nodes +copy_dest_ptr = pointer(copy_dest_nodes) +alloc_copy_into!(copy_dest, atree_large) + +@test arena_push.nodes === arena_push_nodes +@test pointer(arena_push.nodes) == arena_push_ptr +@test length(arena_push.nodes) == arena_push_len + 1 +@test parent_arena.nodes === parent_nodes +@test pointer(parent_arena.nodes) == parent_ptr +@test length(parent_arena.nodes) == parent_len + count_nodes(child) +@test arena_large.nodes === arena_large_nodes +@test pointer(arena_large.nodes) == arena_large_ptr +@test length(arena_large.nodes) == arena_large_len + count_nodes(tree_large) +@test copy_dest.nodes === copy_dest_nodes +@test pointer(copy_dest.nodes) == copy_dest_ptr +@test length(copy_dest.nodes) == count_nodes(atree_large) diff --git a/test/unittest.jl b/test/unittest.jl index 78e0dcd7..7b919083 100644 --- a/test/unittest.jl +++ b/test/unittest.jl @@ -134,3 +134,4 @@ include("test_structured_expression.jl") include("test_zygote_gradient_wrapper.jl") include("test_supposition_consistency.jl") include("test_n_arity_nodes.jl") +include("test_arenanode.jl")