Operator graph optimization tools

These functions are used to restructure the operator graph so that it can be simulated more efficiently when converted into a TensorFlow graph.

nengo_dl.graph_optimizer.mergeable(op, chosen_ops)

Check if the given op can be merged with the candidate group

Parameters:

op : Operator

The operator to be merged

chosen_ops : list of Operator

The operator group to be merged in to

Returns:

bool

True if op can be merged into chosen_ops, else False

nengo_dl.graph_optimizer.greedy_planner(operators)

Combine mergeable operators into groups that will be executed as a single computation.

Parameters:

operators : list of Operator

All the nengo operators in a model (unordered)

Returns:

list of tuple of :class:`~nengo:nengo.builder.Operator`

Operators combined into mergeable groups and in execution order

Notes

Originally based on nengo_ocl greedy planner

nengo_dl.graph_optimizer.tree_planner(op_list, max_depth=3)

Create merged execution plan through exhaustive tree search.

The max_depth parameter scales the planner between full tree search and greedy search. max_depth==1 is equivalent to greedy_planner(), and max_depth==len(op_list) is full tree search (guaranteed to find the optimal plan, but likely very slow).

Parameters:

op_list : list of Operator

All the nengo operators in a model (unordered)

max_depth : int, optional

The planner will search this many steps ahead before selecting which group to schedule next

Returns:

list of tuple of :class:`~nengo:nengo.builder.Operator`

Operators combined into mergeable groups and in execution order

nengo_dl.graph_optimizer.noop_planner(operators)

Orders operators into a valid execution order, but does not perform any merging.

Parameters:

operators : list of Operator

All the nengo operators in a model (unordered)

Returns:

list of tuple of :class:`~nengo:nengo.builder.Operator`

Operators in execution order

nengo_dl.graph_optimizer.transitive_planner(op_list)

Create merged execution plan through transitive closure construction.

This is something like a middle ground between greedy_planner() and tree_planner(); it can improve simulation time over the greedy planner, but comes with potentially significant build time increases.

Parameters:

op_list : list of Operator

All the nengo operators in a model (unordered)

Returns:

list of tuple of :class:`~nengo:nengo.builder.Operator`

Operators combined into mergeable groups and in execution order

nengo_dl.graph_optimizer.transitive_closure_recurse(dg, ops, trans, builder_type, builder_types, cache)

Computes the transitive closure for the given graph, restricted to the operators with the given builder type.

Parameters:

dg : dict of {int: set of int}

Dependency graph where dg[a] = {b, c} indicates that operators b and c are dependent on a

ops : list of int

The operators for which we want to compute the transitive closure

trans : dict of {int: set of int}

The transitive closure for the graph (will be filled in-place)

builder_type : type

One of the nengo_dl build classes (e.g., CopyBuilder), specifying the type of operators to include in the transitive closure

builder_types : list of type

The build class for each operator

cache : dict of {frozenset of int: set of int}

Stores base sets which trans will reference (to reduce memory usage, since many elements in trans will have the same value)

Notes

This function uses ints to refer to operators, where the int indicates the index of the operator in the overall op list (this is done to save memory). See transitive_planner().

nengo_dl.graph_optimizer.order_signals(plan, n_passes=10)

Orders signals and operators to try to structure reads in contiguous blocks.

Parameters:

plan : list of tuple of Operator

Operator execution plan (e.g., output from greedy_planner)

n_passes : int, optional

Number of repeated passes through the operator reordering stage

Returns:

list of :class:`~nengo:nengo.builder.Signal`

Signals organized into the order in which we want them arranged in memory

list of tuple of :class:`~nengo:nengo.builder.Operator`

Input plan with operators reordered within groups to align with order of signals

nengo_dl.graph_optimizer.hamming_sort(blocks)

Reorder signals using heuristics to try to place signals that are read by the same operators into adjacent positions (giving priority to larger blocks).

Parameters:

blocks : dict of {Signal: frozenset of int}

Dictionary indicating which read blocks each signal is a part of

Returns:

dict of {:class:`~nengo:nengo.builder.Signal`: int}

Indices indicating where each signal should be in the sorted list

nengo_dl.graph_optimizer.sort_ops_by_signals(sorted_reads, sigs, sig_idxs, new_plan, blocks, reads)

Rearrange operators to match the order of signals.

Note: the same operators can be associated with multiple read blocks if they have multiple inputs, so rearranging the operators according to one of those blocks could mess up the order with respect to the other read block. We iterate through the read blocks in increasing size so that the largest blocks win out.

Parameters:

sorted_reads : list of tuple of (Operator, int)

The operators that form each read block, sorted by increasing size of the read block. In the case that a group of operators participate in multiple read blocks, the integer distinguishes which one of those inputs this block is associated with.

sigs : list of Signal

Signals that have been arranged into a given order by other parts of the algorithm

sig_idxs : dict of {Signal: int}

Sorted indices of signals

new_plan : dict of {tuple of Operator: tuple of Operator}

Mapping from original operator group to the sorted operators

blocks : dict of {Signal: frozenset of int}

Indicates which read blocks each signal participates in

reads : dict of {Operator: list of Signal}

The signals read by each operator

Returns:

new_plan : dict of {tuple of Operator: tuple of Operator}

Mapping from original operator group to the sorted operators

sig_idxs : dict of {Signal: int}

Signal indices, possibly updated to match new op order

nengo_dl.graph_optimizer.sort_signals_by_ops(sorted_reads, sigs, sig_idxs, new_plan, blocks, reads)

Attempts to rearrange sigs so that it is in the same order as operator reads, without changing the overall block order.

Parameters:

sorted_reads : list of tuple of (Operator, int)

The operators that form each read block, sorted by increasing size of the read block. In the case that a group of operators participate in multiple read blocks, the integer distinguishes which one of those inputs this block is associated with.

sigs : list of Signal

Signals to be sorted

sig_idxs : dict of {Signal: int}

Sorted indices of signals

new_plan : dict of {tuple of Operator: tuple of Operator}

Mapping from original operator group to the sorted operators

blocks : dict of {Signal: frozenset of int}

Indicates which read blocks each signal participates in

reads : dict of {Operator: list of Signal}

The signals read by each operator

Returns:

sig_idxs : dict of {Signal: int}

Sorted indices of signals

nengo_dl.graph_optimizer.noop_order_signals(plan, **kwargs)

A version of graph_optimizer.order_signals() that doesn’t do any reordering, for debugging.

nengo_dl.graph_optimizer.create_signals(sigs, plan, float_type, minibatch_size)

Groups signal data together into larger arrays, and represent each individual signal as a slice into that array.

Parameters:

sigs : list of Signal

Base signals arranged into the order in which they should reside in memory (e.g., output from order_signals)

plan : list of tuple of Operator

Operator execution plan (only used to get a list of all the operators)

float_type : np.float32 or np.float64

Floating point precision to use for signals

minibatch_size : int

Number of items in each minibatch

Returns:

base_arrays : dict of {object

combined arrays, containing the initial values for all signals

sig_map : dict of {Signal: signals.TensorSignal}

mapping from nengo Signals to nengo_dl TensorSignals (views into the base arrays)

nengo_dl.graph_optimizer.remove_unmodified_resets(operators)

Remove any Reset operators that are targeting a signal that is never modified.

If a signal is reset, but never inced/updated after that, we can just set the default signal value to the reset value and remove the reset. Note: this wouldn’t normally happen, but it can happen if we removed some of the incs (e.g. in remove_zero_incs).

Parameters:

operators : list of Operator

Operators in the model

Returns:

list of :class:`~nengo:nengo.builder.Operator`

Modified list of operators

nengo_dl.graph_optimizer.remove_zero_incs(operators)

Remove any operators where we know the input (and therefore output) is zero.

If the input to a DotInc/ElementwiseInc/Copy is zero then we know that the output of the op will be zero, so we can just get rid of it.

Parameters:

operators : list of Operator

Operators in the model

Returns:

list of :class:`~nengo:nengo.builder.Operator`

Modified list of operators

nengo_dl.graph_optimizer.remove_constant_copies(operators)

Change Copies with constant input to Resets.

If a Copy has no dependencies, or just one Reset() dependency, then we can change it to an op that just directly sets the output signal to the Copy input value.

Parameters:

operators : list of Operator

Operators in the model

Returns:

list of :class:`~nengo:nengo.builder.Operator`

Modified list of operators

nengo_dl.graph_optimizer.remove_identity_muls(operators)

Change y=x*1 ops to y=x Copy ops.

If one of the inputs to a DotInc/ElementwiseInc is 1 then we can skip the multiplication and change it to a Copy op.

Parameters:

operators : list of Operator

Operators in the model

Returns:

list of :class:`~nengo:nengo.builder.Operator`

Modified list of operators

nengo_dl.graph_optimizer.signal_io_dicts(operators)

Organizes operators into dictionaries according to the signals they set/inc/read/update.

Parameters:

operators : list of Operator

Operators in the model

Returns

sets : dict of {Signal: list of Operator}

A dictionary indicating all the Operators that set each signal.

incs : dict of {Signal: list of Operator}

A dictionary indicating all the Operators that inc each signal.

reads : dict of {Signal: list of Operator}

A dictionary indicating all the Operators that read each signal.

updates : dict of {Signal: list of Operator}

A dictionary indicating all the Operators that update each signal.

nengo_dl.graph_optimizer.display_signal_blocks(operators, all_signals)

Creates a visual depiction of the signals blocks read by each operator group.

Parameters:

operators : list of tuple of Operator

Operator execution plan

all_signals : list of Signal

Base signals arranged into some order

Returns:

str

A string where each row corresponds to one operator group, and the non-blank characters in the line indicate that the operator group reads/writes that signal (with a number used to distinguish the different signal blocks within the operator group).