[caffe2][redo] Reimplement RemoveOpsByType with SSA (pytorch#41606)

Summary:
Pull Request resolved: pytorch#41606

The previous diff (D22220798 (pytorch@59294fb) and D22220797) was recently reverted (D22492356 (pytorch@28291d3), D22492355) because of a bug associated with the op AsyncIf. The AsyncIf op has net_defs as args and the SSA rewriting didn't take that into account. It has a special path for the op If, but not for AsyncIf. Several changes I made to fix the bug:
1) Add op AsyncIf to the special path for If op in SSA rewriting
2) clear inputs/outputs of the netdefs that are args in If/AsyncIf ops because they're no longer valid
3) revert renamed inputs/outputs in the arg netdefs that are in the external_outputs in the parent netdef

2) and 3) are existing bugs in the `SsaRewrite` function that were just never exposed before.

The algorithm for `RemoveOpsByType` is the same as in my previous diff D22220798 (pytorch@59294fb). The only new changes in this diff are in `onnx::SsaRewrite` and a few newly added unit tests.

(Note: this ignores all push blocking failures!)

Reviewed By: yinghai

Differential Revision: D22588652

fbshipit-source-id: ebb68ecd1662ea2bae14d4be8f61a75cd8b7e3e6

Author: Hao Lu

caffe2/onnx/onnx_exporter.cc

@@ -103,31 +103,6 @@ NodeProto AddShapeNode(const std::string& input, const std::string& output) {
   return shape_node;
 }
 
-} // namespace
-
-::ONNX_NAMESPACE::TensorProto::DataType Caffe2TypeToOnnxType(
-    caffe2::TensorProto::DataType t) {
-#define CAFFE2_TO_ONNX_TYPE(x)   \
-  case (caffe2::TensorProto::x): \
-    return ::ONNX_NAMESPACE::TensorProto::x
-  switch (t) {
-    CAFFE2_TO_ONNX_TYPE(FLOAT);
-    CAFFE2_TO_ONNX_TYPE(BOOL);
-    CAFFE2_TO_ONNX_TYPE(INT8);
-    CAFFE2_TO_ONNX_TYPE(UINT8);
-    CAFFE2_TO_ONNX_TYPE(UINT16);
-    CAFFE2_TO_ONNX_TYPE(INT16);
-    CAFFE2_TO_ONNX_TYPE(INT32);
-    CAFFE2_TO_ONNX_TYPE(INT64);
-    CAFFE2_TO_ONNX_TYPE(FLOAT16);
-    default:
-      LOG(WARNING) << "Unsupported Caffe2 tensor type: " << t
-                   << ", fallback to FLOAT";
-      return ::ONNX_NAMESPACE::TensorProto::FLOAT;
-  }
-#undef CAFFE2_TO_ONNX_TYPE
-}
-
 void collectExternalsFromIfOpSubnet(
     const NetDef* net,
     std::vector<std::string>* input,
@@ -158,6 +133,9 @@ void rewriteSubnet(
     Argument* arg,
     std::map<std::string, std::string> oldname_to_newname) {
   NetDef* net = arg->mutable_n();
+  // clear external inputs and outputs since they're no longer valid
+  net->mutable_external_input()->Clear();
+  net->mutable_external_output()->Clear();
   for (auto& op : *(net->mutable_op())) {
     for (auto& input : *(op.mutable_input())) {
       if (oldname_to_newname.find(input) != oldname_to_newname.end()) {
@@ -245,6 +223,72 @@ void ssaRewriteForIfOp(
   }
 }
 
+void revertRenamedExternalOutput(
+    OperatorDef* op,
+    const std::unordered_map<std::string, std::string>&
+        renamed_external_outputs) {
+  for (auto& input : *(op->mutable_input())) {
+    const auto it = renamed_external_outputs.find(input);
+    if (it != renamed_external_outputs.end()) {
+      input = it->second;
+    }
+  }
+  for (auto& output : *(op->mutable_output())) {
+    const auto it = renamed_external_outputs.find(output);
+    if (it != renamed_external_outputs.end()) {
+      output = it->second;
+    }
+  }
+}
+
+void revertRenamedExternalOutputForIfOp(
+    OperatorDef* if_op,
+    const std::unordered_map<std::string, std::string>&
+        renamed_external_outputs) {
+  ArgumentHelper helper(*if_op);
+  Argument *then_arg = nullptr, *else_arg = nullptr;
+
+  if (helper.HasSingleArgumentOfType<NetDef>("then_net")) {
+    then_arg = getArgumentFromName(if_op, "then_net");
+    NetDef* net = then_arg->mutable_n();
+    for (auto& op : *(net->mutable_op())) {
+      revertRenamedExternalOutput(&op, renamed_external_outputs);
+    }
+  }
+  if (helper.HasSingleArgumentOfType<NetDef>("else_net")) {
+    else_arg = getArgumentFromName(if_op, "else_net");
+    NetDef* net = else_arg->mutable_n();
+    for (auto& op : *(net->mutable_op())) {
+      revertRenamedExternalOutput(&op, renamed_external_outputs);
+    }
+  }
+}
+
+} // namespace
+
+::ONNX_NAMESPACE::TensorProto::DataType Caffe2TypeToOnnxType(
+    caffe2::TensorProto::DataType t) {
+#define CAFFE2_TO_ONNX_TYPE(x)   \
+  case (caffe2::TensorProto::x): \
+    return ::ONNX_NAMESPACE::TensorProto::x
+  switch (t) {
+    CAFFE2_TO_ONNX_TYPE(FLOAT);
+    CAFFE2_TO_ONNX_TYPE(BOOL);
+    CAFFE2_TO_ONNX_TYPE(INT8);
+    CAFFE2_TO_ONNX_TYPE(UINT8);
+    CAFFE2_TO_ONNX_TYPE(UINT16);
+    CAFFE2_TO_ONNX_TYPE(INT16);
+    CAFFE2_TO_ONNX_TYPE(INT32);
+    CAFFE2_TO_ONNX_TYPE(INT64);
+    CAFFE2_TO_ONNX_TYPE(FLOAT16);
+    default:
+      LOG(WARNING) << "Unsupported Caffe2 tensor type: " << t
+                   << ", fallback to FLOAT";
+      return ::ONNX_NAMESPACE::TensorProto::FLOAT;
+  }
+#undef CAFFE2_TO_ONNX_TYPE
+}
+
 std::unordered_map<std::string, std::string> SsaRewrite(
     caffe2::NetDef* init_net,
     caffe2::NetDef* pred_net) {
@@ -288,13 +332,13 @@ std::unordered_map<std::string, std::string> SsaRewrite(
         }
       }
       // Special SSA Rewrite for subnet of If Operator
-      if (op.type() == "If") {
+      if (op.type() == "If" || op.type() == "AsyncIf") {
         ssaRewriteForIfOp(&op, &blob_versions, &is_initialized_tensor);
       }
       for (auto& output : *op.mutable_output()) {
         auto it = blob_versions.find(output);
         if (it != blob_versions.end()) {
-          if (op.type() != "If") {
+          if (op.type() != "If" && op.type() != "AsyncIf") {
             if (is_initialized_tensor.count(output) == 0) {
               it->second += 1;
             } else {
@@ -338,17 +382,11 @@ std::unordered_map<std::string, std::string> SsaRewrite(
     // Use the mapping to find if the input or output of an op was a renamed
     // external output. If so replace it with its original name.
     for (auto& op : *pred_net->mutable_op()) {
-      for (auto& input : *op.mutable_input()) {
-        const auto it = renamed_external_outputs.find(input);
-        if (it != renamed_external_outputs.end()) {
-          input = it->second;
-        }
-      }
-      for (auto& output : *op.mutable_output()) {
-        const auto it = renamed_external_outputs.find(output);
-        if (it != renamed_external_outputs.end()) {
-          output = it->second;
-        }
+      // If/AsyncIf needs special handling
+      if (op.type() == "If" || op.type() == "AsyncIf") {
+        revertRenamedExternalOutputForIfOp(&op, renamed_external_outputs);
+      } else {
+        revertRenamedExternalOutput(&op, renamed_external_outputs);
       }
     }
   }

caffe2/opt/tvm_transformer.cc

@@ -156,20 +156,15 @@ void TvmTransformer::transform(
   std::unordered_set<std::string> weights(
       weight_names.begin(), weight_names.end());
 
-  // SSA Rewrite the net
-  auto shape_hints_mapped =
-      ssaRewriteAndMapNames(ws, pred_net, input_shape_hints);
-
-  // Populate shape info
-  Workspace mapped_ws(ws, input_mapping_);
+  // input_shape_hints should only contain shapes of inputs and not activations
   ShapeInfoMap shape_hints;
   if (!opts_.profiling_based_jit) {
-    shape_hints = inferShapes(
-        &mapped_ws, pred_net, shape_hints_mapped, opts_.bound_shape_spec);
+    shape_hints =
+        inferShapes(ws, pred_net, input_shape_hints, opts_.bound_shape_spec);
   }
 
   if (opts_.debug) {
-    dumpNet(*pred_net, shape_hints, "debug_ssa_net.pbtxt");
+    dumpNet(*pred_net, shape_hints, "debug_net.pbtxt");
   }
 
   // We are ready to transform the net

caffe2/predictor/InferenceGraph.h

@@ -20,5 +20,7 @@ struct InferenceGraph {
   std::vector<std::string> input_names;
   std::vector<std::string> output_names;
   std::vector<std::string> parameter_names;
+
+  bool predictor_net_ssa_rewritten{false};
 };
 } // namespace caffe2