Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Re-use allocations in mutation loop #387

Merged
merged 8 commits into from
Dec 13, 2024
+18 −6
Merged

Re-use allocations in mutation loop #387

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
c189509
deps: bump DE with new buffer tools
MilesCranmer Dec 12, 2024
530689f
feat: reduce allocations within `next_generation`
MilesCranmer Dec 12, 2024
bf1c521
feat: re-use allocations within mutation loop
MilesCranmer Dec 12, 2024
399b86a
fix: ensure right size for preallocated storage
MilesCranmer Dec 12, 2024
adcdf15
fix: allocate node storage for actual size
MilesCranmer Dec 12, 2024
693ef01
fix: prealloc for template expressions
MilesCranmer Dec 12, 2024
a397dad
fix: prealloc was unused
MilesCranmer Dec 13, 2024
19eeb11
refactor: move allocation utility to DynamicExpressions.jl
MilesCranmer Dec 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ Dates = "1"
DifferentiationInterface = "0.5, 0.6"
DispatchDoctor = "^0.4.17"
Distributed = "<0.0.1, 1"
DynamicExpressions = "1.6.0"
DynamicExpressions = "~1.8"
DynamicQuantities = "1"
Enzyme = "0.12, 0.13"
JSON3 = "1"
Expand Down
9 changes: 9 additions & 0 deletions src/ComposableExpression.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,15 @@ end
function CO.count_constants_for_optimization(ex::AbstractComposableExpression)
return CO.count_constants_for_optimization(convert(Expression, ex))
end
function DE.allocate_container(
prototype::ComposableExpression, n::Union{Nothing,Integer}=nothing
)
return (; tree=DE.allocate_container(get_contents(prototype), n))
end
function DE.copy_into!(dest::NamedTuple, src::ComposableExpression)
new_tree = DE.copy_into!(dest.tree, get_contents(src))
return DE.with_contents(src, new_tree)
end

@implements(
ExpressionInterface{all_ei_methods_except(())}, ComposableExpression, [Arguments()]
Expand Down
13 changes: 8 additions & 5 deletions src/Mutate.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,13 @@ module MutateModule

using DynamicExpressions:
AbstractExpression,
copy_into!,
get_tree,
preserve_sharing,
count_scalar_constants,
simplify_tree!,
combine_operators
combine_operators,
allocate_container
using ..CoreModule:
AbstractOptions,
AbstractMutationWeights,
Expand Down Expand Up @@ -187,13 +189,14 @@ function next_generation(
successful_mutation = false
attempts = 0
max_attempts = 10
node_storage = allocate_container(member.tree)

#############################################
# Mutations
#############################################
local tree
while (!successful_mutation) && attempts < max_attempts
tree = copy(member.tree)
tree = copy_into!(node_storage, member.tree)

mutation_result = _dispatch_mutations!(
tree,
Expand Down Expand Up @@ -238,7 +241,7 @@ function next_generation(
mutation_accepted = false
return (
PopMember(
copy(member.tree),
copy_into!(node_storage, member.tree),
beforeScore,
beforeLoss,
options,
Expand Down Expand Up @@ -267,7 +270,7 @@ function next_generation(
mutation_accepted = false
return (
PopMember(
copy(member.tree),
copy_into!(node_storage, member.tree),
beforeScore,
beforeLoss,
options,
Expand Down Expand Up @@ -310,7 +313,7 @@ function next_generation(
mutation_accepted = false
return (
PopMember(
copy(member.tree),
copy_into!(node_storage, member.tree),
beforeScore,
beforeLoss,
options,
Expand Down
Loading