Page History

...

In order to support Gluon hybridization, all control flow operators need to have two versions: one for NDArrays and the other for symbols. The NDArray version is easy to implement. We can simply use Python control flow to implement the operator. The discussion of this section focuses on the implementation of symbol control flow operators.

There are two possible approaches of implementing the symbol control flow operators.

Approach 1: The approach is to keep all computation in a single graph and insert some special operators in the graph to change the executor behavior (e.g., avoid executing a sequence of operators, repeat execution of a sequence of operators). These special operators resembles the jump instruction in CPU. TensorFlow takes this approach. The detailed design can be found in http://download.tensorflow.org/paper/white_paper_tf_control_flow_implementation_2017_11_1.pdf

Approach 2: The approach is to maintain multiple computation graphs. There is a main graph that contains control flow operators. Each control flow operator contains its own computation graph(s) and is responsible for executing the computation graphs inside the operator. This approach represents computation more like how a high-level programming language handles control flows (the code is broken into basic block for compilation and execution).

Both approaches have their pros and cons. I think the second approach will be preferred in MXNet because the execution of MXNet is more static than TensorFlow (e.g., MXNet requires static shape and data type inference and static memory planning). The second approach also allows easier graph-level optimization and easier implementation. From here on, I'll discuss the second approach in more details.

Given the API (shown in the end of the proposal), the first step for approach 2 is to build a subgraph in the user-defined functions of a control flow operator and pass them to the operator. I'll use ``foreach’’ as an example to describe the implementation of a symbolic control flow operator. To help Gluon hybridization, we also implement an NDArray version of these operators.

Building a subgraph: a control flow operator invokes the user-defined functions (UDFs) once to create a computation subgraph. The subgraph of ``foreach’’ has four sets of inputs: slices from the input arrays we iterate over, the states from the previous iteration or the initial states, the variables outside the UDF and the variables defined inside the UDF. The initial states and the variables defined outside the UDF reference to the symbols in the main graph and, thus, connect the subgraph in the UDF with the main graph. However, the main graph and the subgraph have to be completely disjoint (otherwise, stateful operators may have different behaviors). To build a completely separate graph, we create new variable symbols to replace the ones in ``state'' and in the closure. It's easy to find the symbols passed in as ``state''. However, it's more difficult to identify the symbols in the closure and the ones defined inside the UDF, which aren't the inputs of the control flow operator.

...