Skip to content

Commit

Permalink
Improve required allocations dramatically
Browse files Browse the repository at this point in the history
  • Loading branch information
AntonReinhard committed Oct 4, 2024
1 parent c8b9d79 commit 656f17e
Showing 1 changed file with 3 additions and 8 deletions.
11 changes: 3 additions & 8 deletions src/code_gen/tape_machine.jl
Original file line number Diff line number Diff line change
Expand Up @@ -172,7 +172,6 @@ function gen_function_body(tape::Tape; closures_size::Int)
local_inits = gen_local_init.(code_block)

return_symbols = eval.(gen_access_expr.(code_block))
argument_symbols = Set{Symbol}()

ret_symbols_set = Set(return_symbols)
for fc in code_block
Expand All @@ -182,35 +181,31 @@ function gen_function_body(tape::Tape; closures_size::Int)
# symbol won't be defined if it is first calculated in the closure
# so don't add it to the arguments in this case
if !(symbol in ret_symbols_set)
push!(argument_symbols, symbol)
push!(undefined_argument_symbols, symbol)
end
end
end
union!(undefined_argument_symbols, argument_symbols)

intersect!(ret_symbols_set, undefined_argument_symbols)
return_symbols = Symbol[ret_symbols_set...]

argument_symbols = [argument_symbols...] # make sure there is an order (doesn't matter which)

closure = Expr(
:block,
Expr(
:(=),
Expr(:tuple, return_symbols...),
Expr(
:call, # call to the following closure (no arguments)
Expr( # create the closure: (args) -> code block; return (locals)
Expr( # create the closure: () -> code block; return (locals)
:->,
Expr(:tuple, argument_symbols...), # arguments in the closure definition
:(), # closure arguments (none)
Expr( # actual function body of the closure
:block,
local_inits..., # declare local variables with type information inside the closure
expr_from_fc.(code_block)...,
Expr(:return, Expr(:tuple, return_symbols...)),
),
),
argument_symbols..., # arguments to the closure call
),
),
)
Expand Down

0 comments on commit 656f17e

Please sign in to comment.