Understanding Memory Aliasing for Speed and Correctness¶
The aggressive reuse of memory is one of the ways through which Aesara makes code fast, and it is important for the correctness and speed of your program that you understand how Aesara might alias buffers.
This section describes the principles based on which Aesara handles memory, and explains when you might want to alter the default behaviour of some functions and methods for faster performance.
The Memory Model: Two Spaces¶
There are some simple principles that guide Aesara’s handling of memory. The main idea is that there is a pool of memory managed by Aesara, and Aesara tracks changes to values in that pool.
- Aesara manages its own memory space, which typically does not overlap with the memory of normal Python variables that non-Aesara code creates.
- Aesara functions only modify buffers that are in Aesara’s memory space.
- Aesara’s memory space includes the buffers allocated to store
sharedvariables and the temporaries used to evaluate functions.
- Physically, Aesara’s memory space may be spread across the host, a GPU device(s), and in the future may even include objects on a remote machine.
- The memory allocated for a
sharedvariable buffer is unique: it is never aliased to another
- Aesara’s managed memory is constant while Aesara functions are not running and Aesara’s library code is not running.
- The default behaviour of a function is to return user-space values for outputs, and to expect user-space values for inputs.
The distinction between Aesara-managed memory and user-managed memory can be
broken down by some Aesara functions (e.g.
get_value and the
Out) by using a
This can make those methods faster (by avoiding copy operations) at the expense
of risking subtle bugs in the overall program (by aliasing memory).
The rest of this section is aimed at helping you to understand when it is safe
to use the
borrow=True argument and reap the benefits of faster code.
Borrowing when Constructing Function Objects¶
borrow argument can also be provided to the
that control how
aesara.function handles its argument[s] and return value[s].
import aesara import aesara.tensor as at from aesara.compile.io import In, Out x = at.matrix() y = 2 * x f = aesara.function([In(x, borrow=True)], Out(y, borrow=True))
Borrowing an input means that Aesara will treat the argument you provide as if
it were part of Aesara’s pool of temporaries. Consequently, your input
may be reused as a buffer (and overwritten!) during the computation of other variables in the
course of evaluating that function (e.g.
Borrowing an output means that Aesara will not insist on allocating a fresh
output buffer every time you call the function. It will possibly reuse the same one as
on a previous call, and overwrite the old content. Consequently, it may overwrite
old return values through side-effect.
Those return values may also be overwritten in
the course of evaluating another compiled function (for example, the output
may be aliased to a
shared variable). So be careful to use a borrowed return
value right away before calling any more Aesara functions.
The default is of course to not borrow internal results.
It is also possible to pass a
return_internal_type=True flag to the
variable which has the same interpretation as the
get_value function. Unlike
borrow=True arguments to
Out() are not guaranteed to avoid copying an output value. They are just
hints that give more flexibility to the compilation and rewriting of the
Take home message:
When an input
x to a function is not needed after the function
returns and you would like to make it available to Aesara as
additional workspace, then consider marking it with
In(x, borrow=True). It
may make the function faster and reduce its memory requirement. When a return
y is large (in terms of memory footprint), and you only need to read
from it once, right away when it’s returned, then consider marking it with an