Fix out-of-topological-order issue in Nomnigraph (#19458)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/19458

The algorithm in https://fburl.com/ggh9iyvc fails to really ensure topological ordering of nodes. The fix is ugly but effective. I think we need a real topological sort to fix this issue more nicely. Mikhail Zolotukhin, Bram Wasti.

Differential Revision: D15011893

fbshipit-source-id: 130c3aa442f5d578adfb14fbe5f16aa722434942

3 files changed, 32 insertions, 14 deletions

diff --git a/caffe2/core/nomnigraph/Representations/NeuralNet.cc b/caffe2/core/nomnigraph/Representations/NeuralNet.cc
index fbac8f1a1f..174e02b0bb 100644
--- a/caffe2/core/nomnigraph/Representations/NeuralNet.cc
+++ b/caffe2/core/nomnigraph/Representations/NeuralNet.cc
@@ -309,21 +309,26 @@ void coalesceInsertedDataDependencies(repr::NNModule* m) {
   // Finally we reconcile any data dependency issues (if we can).
   for (auto& bbNode : m->controlFlow.getMutableNodes()) {
     auto bb = bbNode->mutableData();
-    std::unordered_set<repr::NNGraph::NodeRef> seen;
-    for (auto instr_iter = bb->getMutableInstructions()->begin();
-         instr_iter != bb->getMutableInstructions()->end();
-         ++instr_iter) {
-      // This cannot be auto& because *iter is pure R-ref
-      auto instr = *instr_iter;
-      for (auto& output : getOutputs(instr)) {
-        for (auto& consumer : getConsumers(output)) {
-          if (seen.count(consumer)) {
-            bb->moveInstructionBefore(instr, consumer);
+    int permutation;
+    do {
+      permutation = 0;
+      std::unordered_set<repr::NNGraph::NodeRef> seen;
+      for (auto instr_iter = bb->getMutableInstructions()->begin();
+           instr_iter != bb->getMutableInstructions()->end();
+           ++instr_iter) {
+        // This cannot be auto& because *iter is pure R-ref
+        auto instr = *instr_iter;
+        for (auto& output : getOutputs(instr)) {
+          for (auto& consumer : getConsumers(output)) {
+            if (seen.count(consumer)) {
+              bb->moveInstructionBefore(instr, consumer);
+              ++permutation;
+            }
           }
         }
+        seen.insert(instr);
       }
-      seen.insert(instr);
-    }
+    } while (permutation);
   }
 }
 
diff --git a/caffe2/opt/backend_cutting.cc b/caffe2/opt/backend_cutting.cc
index c4dd792470..b48d4961ce 100644
--- a/caffe2/opt/backend_cutting.cc
+++ b/caffe2/opt/backend_cutting.cc
@@ -44,12 +44,24 @@ void DumpGraph(NNGraph* g) {
     assert(node->data() && "Node doesn't have data, can't render it");
     if (isa<NeuralNetOperator>(node->data())) {
       auto* op = dyn_cast<NeuralNetOperator>(node->data().get());
+      const auto& op_def =
+          dyn_cast<Caffe2Annotation>(op->getAnnotation())->getOperatorDef();
+      int pos = -1;
+      for (const auto& arg : op_def.arg()) {
+        if (arg.name() == "net_pos") {
+          if (arg.has_i()) {
+            pos = arg.i();
+          }
+          break;
+        }
+      }
       labelMap["label"] =
           op->getName() + " (" + c10::to_string((unsigned long long)node) + ")";
       auto* annotation = op->getAnnotation();
       if (annotation && isa<Caffe2Annotation>(annotation)) {
         auto device_annotation = dyn_cast<Caffe2Annotation>(annotation);
-        labelMap["label"] += "\\n[" + device_annotation->getDevice() + "]";
+        labelMap["label"] += "\\n[" + device_annotation->getDevice() +
+            ", pos=" + c10::to_string(pos) + "]";
         auto hash = std::hash<std::string>{}(device_annotation->getDevice());
         std::stringstream hex_stream;
         hex_stream << std::hex << hash;
diff --git a/caffe2/opt/onnxifi_transformer.cc b/caffe2/opt/onnxifi_transformer.cc
index 0c0e8fe357..13b860e1eb 100644
--- a/caffe2/opt/onnxifi_transformer.cc
+++ b/caffe2/opt/onnxifi_transformer.cc
@@ -841,7 +841,8 @@ NetDef OnnxifiTransformer::TransformViaC2(
         return SubnetToOnnxifiOpViaC2(net, weights, shape_hints);
       };
 
-  return opt::OptimizeForBackend(*pred_net, c2_supports, c2_converter);
+  return opt::OptimizeForBackend(
+      *pred_net, c2_supports, c2_converter, opts_.debug);
 }
 
 NetDef OnnxifiTransformer::TransformViaOnnx(