This document defines a proposal to use a new implementation of the operator graph for MXNet which will supersede the current implementation using NNVM. It will provide a more powerful and extensible implementation that will be easier to maintain and extend, cutting an external dependency to the project, (which we are thankful and have made contributions to)
Current status:
MXNet, as other deep learning frameworks uses a variable and operator graph to schedule computations, track variables and data flow and perform back-propagation and calculate gradients.
...
NNVM has served MXNet well, now it's time to evolve how we process the graph to make it easier to maintain, more powerful so more complex optimizations such as operator fusion can be done, and to enable future applications such as synthesizing graph operations to compiled code or to hardware description languages.
NNVM Graph limitations that will be addressed by these changes:
- External project / subrepository, not part of the core MXNet source.
- Heavy use of type erasure using strings [1], [2] and dmlc any, which is a C++ class that can hold any type. Heavy use of any makes code very difficult to debug as there is no type information that can be used by the debugger, and data is serialized in the heap or the stack in binary.
- Multiple implementations to represent Nodes 1, 2, with the same name and different semantics.
- Unclear data ownership semantics, heavy use of std::shared_ptr. 1 which has memory overhead over unique_ptr and performance impacts on the atomic reference counting mechanism in parallel applications (contention). This also causes difficult to diagnose nested data dependencies between Nodes, NDArray, Engine Variables, NodePtr etc. There's a big potential for simplification and streamlining.
- Untyped fields such as "info" which is used for Autograd doesn't have type nor semantic. There's also all the untyped node attributes.
- We can't add additional fields and structures to the graph, as it's not extensible due to the reasons above.
- Duality of nnvm Indexed graph (indexed by integer key) and the nnvm graph using memory pointers to maintain the graph structure could be improved. On top of that, there's two more additional representations with respect to NDArrays, autograd info and variables belonging to the graph which could be all simplified and coalesced.
- Node index is harcoded to uint32_t. While this seems inofensive, many algorithms are iterating on vector sizes which are of type size_t. These narrowing conversions from unsigned int 64 to unsigned int 32 are most of the time not checked for overflows, and while often we are far below 2^32 nodes, abstracting the type for the indexed graph key allows for an overflow check in a single location and also flexbility for the size type, so we can choose narrower or wider types depending on the application, for our case uint16_t could be enough, and this would be a single change in one file with the proposed changes.
Project goals:
- Have an internalized graph library that can be maintained as part as MXNet. This makes contributions easier and helps evolve the graph without potential design conflicts for a library that is now part of other project with different goals and use cases than MXNet (TVM). In this way we can evolve the library and fit it for MXNet use cases.
- Use a class passed as template argument to the graph for node and edge attributes, where the graph attributes are stored in a maintainable, type safe and compile-time checked way.
- The nodes of the graph also have templated Node data and attributes class which decouple the implementation of generic graph algorithms, traversal etc, from their usage. This will allow for more data fields and structured data and classes to be added to the graph nodes and edges in a propper way.
- The graph is type safe,extensible, reusable, and is debugging friendly.
- Developers can easily extend the graph implementing logic to dump and process the graph in a generic way without having to deal with too many implementation details.
- Generic facilities can be added such as serialization, optimization passes, operation fusion passes etc. which will add tremendous value to MXNet as a flexible framework going forward in the next steps of deep learning developments.
- AGInfo is simplified and such a vital part can be part of the Node attributes class (NodeData in the example below).
Examples of design changes:
A graph is a data structure that holds relationships between nodes which are connected by edges, there are different ways in which we can abstract and code such a data structure so we can then use the generic data structure to specific purposes, the same as a linked list is implemented with pointers to the next item, but can be exposed to users either as a templated class such as std::list, or the pointer can be internal in the user data structure as is commonly done in C.
...