Page History

...

From this starting node, SelectInput and SelectOutput are called to determine if a neighbor node (that hasn't been selected by another subgraph before) should be selected as a candidate for the subgraph. The search continues when a new node is selected as a candidate, and terminates when no more nodes are found. When the process of searching for candidate nodes ends, all of the candidate nodes will passed to Filter to finalize the subgraph. The filter step gives the user the last opportunity to drop out some of the candidate nodes (by default, Filter returns all nodes as the subgraph nodes). The selector is stateful. It can change its own state when seeing a new node.

class SubgraphSelector {
 public:
  // Select a starting node for a subgraph.
  virtual bool Select(const nnvm::Node &n) = 0;
  // The two methods below select neighbors in the incoming or outgoing links as candidate nodes of a subgraph.
  virtual bool SelectInput(const nnvm::Node &curr_node, const nnvm::Node &new_node) = 0;
  virtual bool SelectOutput(const nnvm::Node &curr_node, const nnvm::Node &new_node) = 0;
  // After collecting all candidate nodes for a subgraph, users can use this method to prune some of the nodes
  // to finalize the subgraph.
  virtual std::vector<nnvm::Node*> Filter(nnvm::Graph* g, const std::vector<nnvm::Node*>& candidates) = 0;
};

We provide a default selector called ContainOpSelector, which extracts a subgraph with operators supported by a backend. This selector or its variant is commonly needed by most of the backends.

When a subgraph is found, the partitioning algorithm invokes SubgraphProperty::CreateSubgraphNode to create a new node for the subgraph and connects the new node back to the original graph to replace the subgraph. The subgraph passed to CreateSubgraphNode contains shape/dtype/storage information of the input nodes. This is important because backends, such as TVM and TensorRT, need these information to compile and select the best kernel. CreateSubgraphNode allows any customization on the subgraph and the node, which gives users many opportunities to optimize the subgraph. For example, some backends, such as TVM and MKLDNN, can perform another level of graph partitioning in CreateSubgraphNode for operator fusion. To perform graph partitioning, we attach a graph property (a class that implements SubgraphProperty) and invoke PartitionGraph.

g.attrs["subgraph_property"] = std::make_shared<nnvm::any>(std::move(property));
g = ApplyPass(std::move(g), "PartitionGraph");

Each backend that uses the subgraph mechanism for integration needs to register a subgraph property with a macro MXNET_REGISTER_SUBGRAPH_PROPERTY. For example, MKLDNN can register its subgraph property as below. The first argument of this macro is the backend name and MKLDNNSubgraphProperty is a class name of the MKLDNN subgraph property.

MXNET_REGISTER_SUBGRAPH_PROPERTY("mkldnn", MKLDNNSubgraphProperty);

The subgraph properties registered by MXNET_REGISTER_SUBGRAPH_PROPERTY is used for graph partitioning in bind/simple_bind This should be done in bind for Symbol executor and in CachedOp for Gluon hybridize. We only need to partition the forward graph. The backward graph will be generated accordingly. Currently, our design only supports graph partitioning for one backend. The backend is determined by an environment variable MXNET_SUBGRAPH_BACKEND. In the next step, we'll support graph partitioning for multiple backends.

Step 2: subgraph operator (function call)
Although there are two levels of graph partitioning, we only need to handle one level of subgraphs in the executor because the subgraphs in the second level are fused into operators. We can execute these subgraphs inside special operators, which is specific to the backend.

...