Subgraph API for integrating accelerators with MXNet #12157

reminisce · 2018-08-14T05:54:24Z

Description

This is a joint work of @reminisce and @zheng-da on implementing subgraph API for integrating backend accelerators (TVM, MKLDNN, TensorRT, nGraph, etc.) with MXNet. See design proposal.

Highlighted deliverables

A graph partitioning algorithm accepting custom partitioning rules (through subclassing SubgraphProperty).
A naive graph partitioning rule called DefaultSubgraphProperty partitioning graphs based upon provided operator names. For example, in the left graph, operators B and C are designated to be grouped into a subgraph operator node. After applying the partitioning algorithm, nodes B and C are replaced by a subgraph operator node in the graph on the right-hand-side.
DefaultSubgraphOperator. A stateful operator storing a subgraph for execution by CachedOp.
NDArray version number for tracking mutations of NDArrays. This is useful, for example, in a stateful MKLDNN subgraph operator for checking whether input data has changed since the last forward call to decide if to trigger data format conversion.
Graph partitioning enabled in GraphExecutor for the forward pass. Triggered by an environment variable MXNET_SUBGRAPH_BACKEND.

Coming soon after this PR

Backward pass support in DefaultSubgraphOperator.
Support graph partitioning for Gluon/CacheOp.
Tutorial for developers on integrating backend accelerators using the Subgraph API.
Real backend accelerator integration with MXNet.

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

The PR title starts with [MXNET-$JIRA_ID], where $JIRA_ID refers to the relevant JIRA issue created (except PRs with tiny changes)
Changes are complete (i.e. I finished coding on this PR)
All changes have test coverage:
Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
Nightly tests are added for complicated/long-running ones (e.g. changing distributed kvstore)
Build tests will be added for build configuration changes (e.g. adding a new build option with NCCL)
Code is well-documented:
For user-facing API changes, API doc string has been updated.
For new C++ functions in header files, their functionalities and arguments are documented.
For new examples, README.md is added to explain the what the example does, the source of the dataset, expected performance on test set and reference to the original paper if applicable
Check the API doc at http://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

pengzhao-intel · 2018-08-15T03:08:07Z

Thanks @reminisce it's really helpful for us :)

@ZhennanQin @lvtao, please take a look the compatibility of current code with our interface :)

ZhennanQin · 2018-08-14T06:16:36Z

src/operator/subgraph/default_subgraph_property.h

+    nnvm::NodePtr n = nnvm::Node::Create();
+    n->attrs.op = Op::Get("_default_subgraph_op");
+    n->attrs.name = "_default_subgraph_op" + std::to_string(subgraph_id);
+    n->attrs.subgraphs.push_back(std::make_shared<nnvm::Symbol>(sym));


subgraph symbol is pushed back into n->attrs.subgraphs, how can we guarantee it is at index 0? Because for DefaultSubgraphOpNumInputs(), it always try to get symbol from index 0.

n is a newly created node and its attrs.subgraphs is empty before push_back is called. It's actually up to you to decide which index is for the sym in your specific subgraph property. Just make sure you get the correct subgraph in your subgraph operator implementation.

I heard that control dependence will use attrs.subgraphs as well. I'm not sure how this attribute will be used for dependence purpose. Do we have any instruction to define the sequence for attrs.subgraphs?Eg. always put subgraph symbol in index 0, and following dependence purpose node.

As far as I know, there is no fixed convention of doing this. @zheng-da
I don't think you need to worry about this if you can ensure the consistency between your specific subgraph property and subgraph operator.

We can ensure the consistency between subgraph property and subgraph operator, and we also need to ensure nobody changes it after subgraph partition. It's better to add some check for this consistency in case it breaks.

One thing you can do to enforce the consistency is defining some enums as indices of forward graph, backward graph, fused graph, etc in attrs.subgraphs, and always use those enums to get the desired subgraphs. It also increases the code readability.

anirudh2290 · 2018-08-21T07:30:45Z

src/operator/subgraph/subgraph_property.h

+
+#define MXNET_REGISTER_SUBGRAPH_PROPERTY(Name, SubgraphPropertyType) \
+  static DMLC_ATTRIBUTE_UNUSED auto __make_ ## SubgraphPropertyType ## _ ## Name ## __ = \
+    SubgraphPropertyRegistry::Get()->__REGISTER__(#Name, &SubgraphPropertyType::Create)


Can we implement and call __REGISTER_OR_GET__ here (which checks if it is part of map else calls __REGISTER__). This may be needed because it is possible that __REGISTER__ may be called twice on the same name leading to exception.

This is intended. The registration can only happen once.

With __REGISTER_OR_GET__ also register will happen only once. To explain an use case where register wont work: lets say I create a custom subgraph property and call MXNET_REGISTER_SUBGRAPH_PROPERTY from inside the namespace. Now i include this custom subgraph property header file in multiple source files. If they are part of different translation units __REGISTER__ called multiple times and fails since the static variable scope is limited to a single translation unit.

That's a good catch. I will fix it. Thanks.

zheng-da · 2018-08-22T22:53:35Z

src/imperative/imperative_utils.h

@@ -436,7 +436,8 @@ inline void PushFComputeEx(const FComputeEx& fn,
      }
    };

-  if (exec_type == ExecType::kCrossDeviceCopy) {
+  if (exec_type == ExecType::kCrossDeviceCopy
+      || exec_type == ExecType::kSubgraphExec) {


we don't need kSubgraphExec here. kSubgraphExec only works for stateful operators.

I think this line is added by you. Nevertheless, isn't _default_subgraph_op a stateful operator? This should be needed to make sure this operator executes in the main thread in CachedOp.

Yes, I added this line. It's actually not necessary because this is PushFComputeEx, which only pushes FComputeEx. could you help me remove the line?

zheng-da · 2018-08-22T22:53:58Z

src/ndarray/ndarray.cc

@@ -39,6 +39,7 @@
 #include "../operator/tensor/matrix_op-inl.h"
 #include "../operator/tensor/init_op.h"
 #include "../operator/nn/mkldnn/mkldnn_base-inl.h"
+#include "../engine/engine_impl.h"


do we need to include this?

zheng-da · 2018-08-22T23:03:26Z

src/operator/subgraph/default_subgraph_op.cc

+  for (size_t i = 0; i < outputs.size(); ++i) {
+    LOG(INFO) << "outputs[" << i << "].version = " << outputs[i].version();
+  }
+#endif


should we clean this?

As discussed offline, we are going to keep the debugging code for a while.

Ah, I see. Removed now.

zheng-da · 2018-08-22T23:18:05Z

src/operator/subgraph/partition_graph.cc

+    PrintNodeEntry(*entries[i]);
+  }
+}
+#endif


delete the debug code?

anirudh2290 · 2018-08-23T22:46:32Z

src/engine/naive_engine.cc

@@ -165,14 +178,22 @@ class NaiveEngine final : public Engine {
    }
    CHECK(this->req_completed_)
        << "NaiveEngine only support synchronize Push so far";
+    // increment var version
+    for (auto var : mutable_vars) {


In case of deletevariable, will this increment the var version after it is deleted ?

Good catch. Should move this before exec_fun runs.

anirudh2290 · 2018-08-23T22:50:32Z

src/executor/graph_executor.cc

+      ++i1;
+    }
+  }
+  return PartitionGraph(src, prop_name, arg_shapes, arg_dtypes, arg_stypes,


should we check for prop_name here. Currently only default is supported correct ?

nice catch. we should check in here https://github.com/apache/incubator-mxnet/pull/12157/files/8435acff465725e406bbd467902c3e6bb5c83139#diff-d8b5a5b027d00584737fb6486cba38b9R1467

There is a check inside CreateSubgraphProperty.

anirudh2290 · 2018-08-23T22:57:50Z

src/operator/subgraph/default_subgraph_property.cc

+  }
+  virtual SubgraphSelectorPtr CreateSubgraphSelector() const {
+    return std::make_shared<ContainOpSelector>(
+        this->GetAttr<std::unordered_set<std::string>>("op_names"));


Can we override GetAttr for DefaultSubgraphProperty, so that it informs the user that the set has to be provided by the user using C API to test default subgraph property? This will also be a good example for developers writing custom subgraph property.

this is a templated method. I don't think we can make it a virtual method. but I don't really know what you mean here.

To explain my use case: Lets say a user forgot to call setattr for op_names when using the default_subgraph_property. It would fial with the error that op_names cannot be found in subgraph property. A better error message would have been useful. There may be such cases where classes subclassing SubgraphProperty may want to also implement logic for SetAttr and GetAttr in the subclass.

It's not meant for users to use and is there only for the flexible testing purpose. Furthermore, template method cannot be virtual. But I still cannot understand your use case. How different could SetAttr and GetAttr be in a subclass than the base?

for example you only want certain attributes to be set for the subgraph property and you want to do some sanity check on whether the attrs can be supported or not. For example supported op_names.

for your use case, i think it's fine without customization. Users of subgraph API can write their own code to do sanity check. The main problem is a templated method can't be virtual.

zheng-da · 2018-08-24T16:19:25Z

src/executor/graph_executor.cc

+                                   const std::map<std::string, Context>& ctx_map,
+                                   const std::vector<Context>& in_arg_ctxes,
+                                   const std::vector<Context>& aux_state_ctxes) {
+  auto subgraph_prop = op::SubgraphPropertyRegistry::Get()->CreateSubgraphProperty(prop_name);


If prop_name doesn't exist, we need to handle the failure of creating a subgraph property.

There is a check inside CreateSubgraphProperty.

zheng-da · 2018-08-24T16:24:10Z

src/executor/graph_executor.cc

+      ++i1;
+    }
+  }
+  return PartitionGraph(src, prop_name, arg_shapes, arg_dtypes, arg_stypes,


nice catch. we should check in here https://github.com/apache/incubator-mxnet/pull/12157/files/8435acff465725e406bbd467902c3e6bb5c83139#diff-d8b5a5b027d00584737fb6486cba38b9R1467

zheng-da · 2018-08-24T16:26:16Z

src/operator/subgraph/default_subgraph_property.cc

+  }
+  virtual SubgraphSelectorPtr CreateSubgraphSelector() const {
+    return std::make_shared<ContainOpSelector>(
+        this->GetAttr<std::unordered_set<std::string>>("op_names"));


this is a templated method. I don't think we can make it a virtual method. but I don't really know what you mean here.

piiswrong · 2018-08-24T19:06:46Z

Let's not put this in example since the CAPIs are not meant to be called by users

eric-haibin-lin

Nice work!

eric-haibin-lin · 2018-08-26T17:18:18Z

src/executor/graph_executor.cc

@@ -42,6 +43,7 @@ using namespace mxnet::common;
 GraphExecutor::GraphExecutor() {
  log_verbose_ = dmlc::GetEnv("MXNET_EXEC_VERBOSE_LOGGING", false);
  need_grad_ = false;
+  subgraph_property_ = dmlc::GetEnv("MXNET_SUBGRAPH_BACKEND", std::string());


Shall we also add MXNET_SUBGRAPH_BACKEND to env_var.md?

Right now, there is no real accelerator integrated with MXNet. Intel team is going to submit their work on MKLDNN using this API set. We can add this evn var along with their PR.

eric-haibin-lin · 2018-08-26T17:51:32Z

src/operator/subgraph/common.h

+  return ret;
+}
+
+inline std::vector<ResourceRequest> DefaultSubgraphOpResourceRequest(const nnvm::NodeAttrs& attrs) {


This function aggregates all resources requested by the subgraph. What are these acquired resources distributed to the individual operators in the subgraph?

This is a convenience function for developers to register with their own subgraph operators. They may or may not use this function for registration depending on their implementation. For the _defafult_subgraph_operator, this is actually not necessary to be there. I have removed it from the registering _defafult_subgraph_operator.

eric-haibin-lin · 2018-08-26T17:54:25Z

src/operator/subgraph/subgraph_property.h

+  // Determine if the node should be selected for a subgraph.
+  virtual bool Select(const nnvm::Node &n) = 0;
+  // Determine if the input node should be selected for a subgraph.
+  virtual bool SelectInput(const nnvm::Node &n, const nnvm::Node &new_node) = 0;


Do you document what are n and new_node?

Review comments addressed

anirudh2290

Good work guys ! @reminisce @zheng-da

anirudh2290 · 2018-08-28T18:52:12Z

src/operator/subgraph/partition_graph.cc

+  const size_t max_num_retry = simple_nodes.size() * simple_nodes.size();
+  size_t count = 0;
+  bool success = false;
+  while (!success && count < max_num_retry) {


Can you explain the logic behind ssetting max_num_retry to square of simple_nodes.size(). Shouldnt it just be simple_nodes.size()

Theoretically, yes. But just want to be safe here.

anirudh2290 · 2018-08-28T18:54:01Z

src/operator/subgraph/partition_graph.cc

+      }
+      for (const auto& entry : top->inputs) {
+        // when searching for the ancestor, the path cannot cross any subgraph node
+        auto it = std::find(snodes.begin(), snodes.end(), entry.node.get());


nit: maybe add check for it == snodes.end() since it wont happen that it will find the input node in the subgraph nodes ?

Either case is possible here, so cannot add the CHECK.

eric-haibin-lin · 2018-08-29T03:15:58Z

LGTM

Graph partitioner and subgraph op Fix duplicate entry bugs (apache#11767) Make subgraph var node name unique (apache#11876) [DO NOT REVIEW] Fix bug of eliminating cycles (apache#11907) * Fix cycle bug * Fix decycle bug * Fix comment [DO NOT REVIEW] Subgraph API (apache#12104) * Initial commit * Add unit tests * Fix lint * Fix lint * Clean up * Add graph partitiong to Bind * Add property name to graph partitioning c api * Fix unit test gpu context * Address cr * Move subgraph to attrs.subgraphs and fix the example * Fix lint * Add var version unit test * Address cr * Enable unit test that was flaky

eric-haibin-lin · 2018-08-31T02:13:23Z

Looks like all comments are addressed. Merging it now

…12157)" This reverts commit a64cf7d.

* Graph partitioner and subgraph op (apache#11251) Graph partitioner and subgraph op Fix duplicate entry bugs (apache#11767) Make subgraph var node name unique (apache#11876) [DO NOT REVIEW] Fix bug of eliminating cycles (apache#11907) * Fix cycle bug * Fix decycle bug * Fix comment [DO NOT REVIEW] Subgraph API (apache#12104) * Initial commit * Add unit tests * Fix lint * Fix lint * Clean up * Add graph partitiong to Bind * Add property name to graph partitioning c api * Fix unit test gpu context * Address cr * Move subgraph to attrs.subgraphs and fix the example * Fix lint * Add var version unit test * Address cr * Enable unit test that was flaky * Clean up * Clean up * Clean up * Change version return type in NDArray * Clean up * Add register or get for subgraph prop registry * Address cr * Remove unnecessary code * Handle var version issue in naive engine * Delete example * Remove registration of resource request for default subgraph op * Add doc string * Improve doc string

reminisce requested review from piiswrong, tqchen and eric-haibin-lin August 14, 2018 05:54

reminisce requested review from anirudh2290 and szha as code owners August 14, 2018 05:54

ZhennanQin reviewed Aug 15, 2018

View reviewed changes

anirudh2290 previously requested changes Aug 21, 2018

View reviewed changes

reminisce force-pushed the subgraph_api_to_master branch from 2f1434d to 3c93b12 Compare August 21, 2018 19:08

zheng-da reviewed Aug 22, 2018

View reviewed changes

reminisce force-pushed the subgraph_api_to_master branch from 3c93b12 to 54971ba Compare August 23, 2018 03:53

anirudh2290 reviewed Aug 23, 2018

View reviewed changes

zheng-da reviewed Aug 24, 2018

View reviewed changes

zheng-da approved these changes Aug 24, 2018

View reviewed changes

reminisce force-pushed the subgraph_api_to_master branch from 8435acf to 964d6e2 Compare August 24, 2018 18:26

eric-haibin-lin reviewed Aug 26, 2018

View reviewed changes

zheng-da mentioned this pull request Aug 27, 2018

[MXNET-876] make CachedOp a normal operator #11641

Merged

7 tasks

anirudh2290 reviewed Aug 28, 2018

View reviewed changes

reminisce force-pushed the subgraph_api_to_master branch from a1aef90 to 3a722c4 Compare August 28, 2018 19:41

reminisce added 8 commits August 29, 2018 16:24

Clean up

8b10cbd

Clean up

9fd8ae9

Clean up

0b764e7

Change version return type in NDArray

9228e25

Clean up

db03355

Add register or get for subgraph prop registry

085de6a

Address cr

ddf4334

reminisce added 6 commits August 29, 2018 16:24

Remove unnecessary code

59ca92c

Handle var version issue in naive engine

c375803

Delete example

1e96e5b

Remove registration of resource request for default subgraph op

958f7dc

Add doc string

bdb582d

Improve doc string

ae2b480

reminisce force-pushed the subgraph_api_to_master branch from 3a722c4 to ae2b480 Compare August 29, 2018 23:25

piiswrong approved these changes Aug 30, 2018

View reviewed changes

eric-haibin-lin merged commit a64cf7d into apache:master Aug 31, 2018

lebeg added a commit to lebeg/incubator-mxnet that referenced this pull request Sep 3, 2018

Revert "Subgraph API for integrating accelerators with MXNet (apache#…

ceaac8c

…12157)" This reverts commit a64cf7d.

This was referenced Sep 3, 2018

Master build fails due to Subgraph API commit #12442

Closed

Revert "Subgraph API for integrating accelerators with MXNet (#12157)" #12443

Closed

reminisce mentioned this pull request Oct 5, 2018

Features to Announce for the MXNet v1.4 #12741

Closed

reminisce mentioned this pull request Mar 4, 2019

Enhance PartitionGraph #14277

Merged

7 tasks

Subgraph API for integrating accelerators with MXNet #12157

Subgraph API for integrating accelerators with MXNet #12157

Conversation

reminisce commented Aug 14, 2018 • edited Loading

Description

Highlighted deliverables

Coming soon after this PR

Checklist

Essentials

pengzhao-intel commented Aug 15, 2018

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

anirudh2290 Aug 21, 2018 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

anirudh2290 Aug 24, 2018 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

piiswrong commented Aug 24, 2018

eric-haibin-lin left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

anirudh2290 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eric-haibin-lin commented Aug 29, 2018

eric-haibin-lin commented Aug 31, 2018

reminisce commented Aug 14, 2018 •

edited

Loading

anirudh2290 Aug 21, 2018 •

edited

Loading

anirudh2290 Aug 24, 2018 •

edited

Loading