View source on GitHub |
Optimizes a Program
's variable allocation strategy.
tfp.experimental.auto_batching.allocation_strategy.optimize(
program
)
The variable allocation strategies determine how much memory the Program
consumes, and how costly its memory access operations are (see
instructions.VariableAllocation
). In general, a variable holding data with
a longer or more complex lifetime will need a more expensive storage strategy.
This analysis examines variables' liveness and opportunistically selects
inexpensive sound allocation strategies.
Specifically, the algorithm is to:
- Run liveness analysis to determine the lifespan of each variable.
- Assume optimistically that no variable needs to be stored at all
(
instructions.VariableAllocation.NULL
). - Traverse the instructions and pattern-match conditions that require
some storage:
- If a variable is read by an instruction, it must be at least
instructions.VariableAllocation.TEMPORARY
. - If a variable is live out of some block (i.e., crosses a block boundary),
it must be at least
instructions.VariableAllocation.REGISTER
. This is because temporaries do not appear in the loop state inexecute
. - If a variable is alive across a call to an autobatched
Function
, it must beinstructions.VariableAllocation.FULL
, because thatFunction
may push values to it that must not overwrite the value present at the call point. (This can be improved by examining the call graph to see whether the callee really does push values to this variable, but that's future work.)
- If a variable is read by an instruction, it must be at least
Args | |
---|---|
program
|
Program to optimize.
|