Fix batch insert (#17158)

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

Because of Reshape op, batch size can be changed. This diff addresses first order issue raised from multiple batch size system. We need to export different real_batch_size for different max_batch_size input and attach it to the right output.

It also fixes a false exception.

Reviewed By: ipiszy

Differential Revision: D14099541

fbshipit-source-id: 0fa9e86826f417a11d2b5dd2ee60dff64a7ce8c4

4 files changed, 43 insertions, 30 deletions

diff --git a/caffe2/opt/bound_shape_inferencer.cc b/caffe2/opt/bound_shape_inferencer.cc
index f0b72053e8..c1b3c4a014 100644
--- a/caffe2/opt/bound_shape_inferencer.cc
+++ b/caffe2/opt/bound_shape_inferencer.cc
@@ -46,7 +46,7 @@ void BoundShapeInferencer::InferBoundShapeAndType(
   shape_info_ = info;
 
   for (const auto& op : net.op()) {
-    LOG(INFO) << op.type();
+    VLOG(1) << op.type();
     if (op.type() == "SparseLengthsSum" ||
         op.type() == "SparseLengthsSumFused8BitRowwise" ||
         op.type() == "SparseLengthsWeightedSum" ||
diff --git a/caffe2/opt/onnxifi_transformer.cc b/caffe2/opt/onnxifi_transformer.cc
index 085546a7ff..aa7c425985 100644
--- a/caffe2/opt/onnxifi_transformer.cc
+++ b/caffe2/opt/onnxifi_transformer.cc
@@ -185,13 +185,11 @@ std::unordered_set<string> ToHashSet(
 int64_t GetBlob1stDimSize(
     const ShapeInfo& shape_info,
     const string& blob_name) {
-  CAFFE_ENFORCE(
-      shape_info.shape.dims_size() > 0 && shape_info.shape.dims(0) > 0,
-      "Tensor " + blob_name +
-          " is type BATCH/SEQ, however the batch_size is unknown. " +
-          "Dims size: " + to_string(shape_info.shape.dims_size()) +
-          ", dim[0] = " + to_string(shape_info.shape.dims(0)));
-  return shape_info.shape.dims(0);
+  if (shape_info.shape.dims_size() == 0) {
+    return 0;
+  } else {
+    return shape_info.shape.dims(0);
+  }
 }
 
 // Generates AdjustBatchOps for external inputs/outputs with type BATCH or
@@ -223,21 +221,26 @@ std::unordered_map<std::string, std::string> AddAdjustBatchOps(
                        << input_blob;
           continue;
         }
-        string real_batch_size_blob = "";
+        std::string real_batch_size_blob = "";
+        auto max_batch_size = 0;
         if (shape_info_it->second.dim_type == ShapeInfo::DimType::BATCH) {
-          real_batch_size_blob = kRealBatchSizeBlob;
+          max_batch_size = GetBlob1stDimSize(shape_info_it->second, input_blob);
+          real_batch_size_blob =
+              kRealBatchSizeBlob + "_" + c10::to_string(max_batch_size);
         } else if (shape_info_it->second.dim_type == ShapeInfo::DimType::SEQ) {
+          max_batch_size = GetBlob1stDimSize(shape_info_it->second, input_blob);
           real_batch_size_blob = MakeSeqSizeBlob(input_blob);
         } else {
           continue;
         }
+
         auto output_blob = MakeOutputForAdjustBatchOp(input_blob);
         auto ret = real_batch_size_blobs.emplace(real_batch_size_blob);
         if (post_adjust_inputs.emplace(output_blob).second) {
           input_ops->push_back(MakeAdjustBatchOp(
               input_blob,
               output_blob,
-              GetBlob1stDimSize(shape_info_it->second, input_blob),
+              max_batch_size,
               ret.second ? real_batch_size_blob : "",
               true /* adjust_to_max_batch_size */));
         }
@@ -257,27 +260,36 @@ std::unordered_map<std::string, std::string> AddAdjustBatchOps(
     for (auto& output_blob : *(op.mutable_output())) {
       if (external_outputs.count(output_blob)) {
         auto shape_info_it = shape_hints.find(output_blob);
-        if (shape_info_it == shape_hints.end()) {
-          continue;
-        }
+        CAFFE_ENFORCE(
+            shape_info_it != shape_hints.end(),
+            "Cannot find shape info for ",
+            output_blob,
+            " for AdjustBatchOp insertion");
         if (shape_info_it->second.dim_type == ShapeInfo::DimType::BATCH) {
-          if (!real_batch_size_blobs.count(kRealBatchSizeBlob)) {
-            continue;
-          }
+          auto max_batch_size =
+              GetBlob1stDimSize(shape_info_it->second, output_blob);
+          std::string real_size_blob =
+              kRealBatchSizeBlob + "_" + c10::to_string(max_batch_size);
+          CAFFE_ENFORCE(
+              real_batch_size_blobs.count(real_size_blob),
+              output_blob,
+              ": Cannot find ",
+              real_size_blob,
+              " to make AdjustBatchOp");
           auto input_blob = MakeInputForAdjustBatchOp(output_blob);
           output_ops->push_back(MakeAdjustBatchOp(
               input_blob,
               output_blob,
-              GetBlob1stDimSize(shape_info_it->second, output_blob),
-              kRealBatchSizeBlob,
+              max_batch_size,
+              real_size_blob,
               false /* adjust_to_max_batch_size */));
           renaming_map[output_blob] = input_blob;
           output_blob = input_blob;
-        } else {
-          CAFFE_ENFORCE(
-              shape_info_it->second.dim_type != ShapeInfo::DimType::SEQ,
-              "Output tensor " + output_blob +
-                  " should never have dim_type SEQ.");
+        } else if (shape_info_it->second.dim_type == ShapeInfo::DimType::SEQ) {
+          LOG(WARNING) << "It's unusual that output tesnor " << output_blob
+                       << " is of dim_type SEQ. "
+                       << "AdjustBatchOp won't attached "
+                       << "and it might degrade the performance";
         }
       }
     }
diff --git a/caffe2/python/onnx/onnxifi.py b/caffe2/python/onnx/onnxifi.py
index 9a859cbf60..ccda2bb3e0 100644
--- a/caffe2/python/onnx/onnxifi.py
+++ b/caffe2/python/onnx/onnxifi.py
@@ -22,7 +22,8 @@ def onnxifi_caffe2_net(
         max_batch_size=1,
         max_seq_size=1,
         debug=False,
-        use_onnx=True):
+        use_onnx=True,
+        black_list=None):
     """
     Transform the caffe2_net by collapsing ONNXIFI-runnable nodes into Onnxifi c2 ops
     """
@@ -31,6 +32,7 @@ def onnxifi_caffe2_net(
         shape_hints[k] = v
     pred_net_str = C.onnxifi(pred_net.SerializeToString(),
                              shape_hints,
+                             black_list if black_list else [],
                              max_batch_size,
                              max_seq_size,
                              debug,
diff --git a/caffe2/python/pybind_state.cc b/caffe2/python/pybind_state.cc
index aec3551992..0bff1e03ea 100644
--- a/caffe2/python/pybind_state.cc
+++ b/caffe2/python/pybind_state.cc
@@ -1605,6 +1605,7 @@ void addGlobalMethods(py::module& m) {
       "onnxifi",
       [](const py::bytes& pred_net_str,
          const std::unordered_map<std::string, std::vector<int>>& shapes,
+         const std::vector<int>& black_list,
          int max_batch_size,
          int max_seq_size,
          bool debug_builder,
@@ -1626,13 +1627,11 @@ void addGlobalMethods(py::module& m) {
         opts.use_onnx = use_onnx;
         OnnxifiTransformer ts(opts);
         Workspace* curr_ws = GetCurrentWorkspace();
+        std::unordered_set<int> blacklist_set(
+            black_list.begin(), black_list.end());
         auto weight_names = curr_ws->Blobs();
         ts.transform(
-            curr_ws,
-            &pred_net,
-            weight_names,
-            tensor_shapes,
-            std::unordered_set<int>());
+            curr_ws, &pred_net, weight_names, tensor_shapes, blacklist_set);
         std::string pred_net_str2;
         pred_net.SerializeToString(&pred_net_str2);
         return py::bytes(pred_net_str2);