Minor improvements to ONNXIFI transform (#17964)

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

1. Make the output of TensorFill outputs CONSTANT during shape inference
2. Add option to avoid adding BatchAdjust ops
3. Add option to avoid lowering subgraph that's smaller than a limit

Reviewed By: hyuen

Differential Revision: D14360903

fbshipit-source-id: b3c5966b44e7cd0d56428acd6cc97f529b36b171

diff --git a/caffe2/opt/backend_transformer_base.cc b/caffe2/opt/backend_transformer_base.cc
index a4db8b6..333c6ae 100644 (file)
--- a/caffe2/opt/backend_transformer_base.cc
+++ b/caffe2/opt/backend_transformer_base.cc
@@ -15,8 +15,18 @@ void AnnotateOpIndex(NetDef* net) {
 
 std::string BackendTransformerBase::getModelId(const NetDef& net) {
   static std::atomic<size_t> seq_id{0};
-  auto model_id =
-      ArgumentHelper(net).GetSingleArgument<std::string>(kModelId, "");
+  std::string model_id;
+  for (const auto& arg : net.arg()) {
+    if (arg.name() == kModelId) {
+      if (arg.has_s()) {
+        model_id = arg.s();
+      } else if (arg.has_i()) {
+        model_id = c10::to_string(arg.i());
+      }
+      break;
+    }
+  }
+
   if (model_id.empty()) {
     model_id = "unnamed_" + c10::to_string(seq_id++);
   }

diff --git a/caffe2/opt/bound_shape_inferencer.cc b/caffe2/opt/bound_shape_inferencer.cc
index e56c8d2..dea1eeb 100644 (file)
--- a/caffe2/opt/bound_shape_inferencer.cc
+++ b/caffe2/opt/bound_shape_inferencer.cc
@@ -62,9 +62,13 @@ void BoundShapeInferencer::InferBoundShapeAndType(
     } else if (op.type() == "LengthsRangeFill") {
       InferLengthsRangeFill(op);
     } else if (
-        caffe2::StartsWith(op.type(), "GivenTensor") &&
-        caffe2::EndsWith(op.type(), "Fill")) {
+        (caffe2::StartsWith(op.type(), "GivenTensor") &&
+         caffe2::EndsWith(op.type(), "Fill")) ||
+        op.type() == "ConstantFill" || op.type() == "Int8GivenTensorFill" ||
+        op.type() == "Int8GivenIntTensorFill") {
       InferGivenTensorFill(op);
+    } else if (op.type() == "Shape") {
+      InferShape(op);
     } else {
       InferCommonOp(op);
     }
@@ -215,6 +219,14 @@ void BoundShapeInferencer::InferSparseLengthsSum(const OperatorDef& op) {
       TensorProto_DataType_FLOAT);
 }
 
+void BoundShapeInferencer::InferShape(const OperatorDef& op) {
+  InferCommonOp(op);
+  // old_shape should be a constant
+  if (op.output_size() > 0 && shape_info_.count(op.output(0))) {
+    shape_info_[op.output(0)].dim_type = ShapeInfo::DimType::CONSTANT;
+  }
+}
+
 void BoundShapeInferencer::InferReshape(const OperatorDef& op) {
   InferCommonOp(op);
   // old_shape should be a constant

diff --git a/caffe2/opt/bound_shape_inferencer.h b/caffe2/opt/bound_shape_inferencer.h
index ef6fa04..333eca3 100644 (file)
--- a/caffe2/opt/bound_shape_inferencer.h
+++ b/caffe2/opt/bound_shape_inferencer.h
@@ -67,6 +67,7 @@ class CAFFE2_API BoundShapeInferencer {
   void InferSparseLengthsSum(const OperatorDef& op);
   void InferFC(const OperatorDef& op);
   void InferConcat(const OperatorDef& op);
+  void InferShape(const OperatorDef& op);
   void InferReshape(const OperatorDef& op);
   void InferLengthsRangeFill(const OperatorDef& op);
 

diff --git a/caffe2/opt/onnxifi_transformer.cc b/caffe2/opt/onnxifi_transformer.cc
index 3289f62..442e7d9 100644 (file)
--- a/caffe2/opt/onnxifi_transformer.cc
+++ b/caffe2/opt/onnxifi_transformer.cc
@@ -409,6 +409,15 @@ NetDef OnnxifiTransformer::SubnetToOnnxifiOpViaC2(
     const caffe2::NetDef& net,
     const std::unordered_set<std::string>& weights_in_ws,
     const ShapeInfoMap& shape_hints) {
+  if (opts_.min_ops > net.op_size()) {
+    return net;
+  }
+  int onnxifi_op_id = onnxifi_op_id_;
+
+  if (opts_.debug) {
+    WriteProtoToTextFile(
+        net, "debug_original_net_" + c10::to_string(onnxifi_op_id) + ".pb_txt");
+  }
   // We already have all the ops and external inputs and outputs!
   NetDef onnxifi_net(net);
 
@@ -437,8 +446,11 @@ NetDef OnnxifiTransformer::SubnetToOnnxifiOpViaC2(
   // names for external_inputs/outputs and input_shape_info for the onnxifi_net.
   vector<OperatorDef> input_ops;
   vector<OperatorDef> output_ops;
-  auto renaming_map =
-      AddAdjustBatchOps(shape_hints, &onnxifi_net, &input_ops, &output_ops);
+  std::unordered_map<std::string, std::string> renaming_map;
+  if (opts_.add_adjust_batch_ops) {
+    renaming_map =
+        AddAdjustBatchOps(shape_hints, &onnxifi_net, &input_ops, &output_ops);
+  }
 
   // Figure out weights and add it to external_inputs too
   std::unordered_set<std::string> initialization_list;
@@ -498,14 +510,11 @@ NetDef OnnxifiTransformer::SubnetToOnnxifiOpViaC2(
   // Debugging stuff
   if (opts_.debug) {
     WriteProtoToTextFile(
-        net,
-        "debug_original_net_" + c10::to_string(onnxifi_op_id_) + ".pb_txt");
-    WriteProtoToTextFile(
         onnxifi_net,
-        "debug_onnxifi_net_" + c10::to_string(onnxifi_op_id_) + ".pb_txt");
+        "debug_onnxifi_net_" + c10::to_string(onnxifi_op_id) + ".pb_txt");
     WriteProtoToTextFile(
         net_opt,
-        "debug_optimized_net_" + c10::to_string(onnxifi_op_id_) + ".pb_txt");
+        "debug_optimized_net_" + c10::to_string(onnxifi_op_id) + ".pb_txt");
   }
   return net_opt;
 }
@@ -516,6 +525,9 @@ NetDef OnnxifiTransformer::SubnetToOnnxifiOpViaOnnx(
     Workspace* ws,
     onnx::OnnxExporter* exporter,
     ShapeInfoMap* shape_hints) {
+  if (opts_.min_ops > net.op_size()) {
+    return net;
+  }
   ::ONNX_NAMESPACE::ModelProto onnx_model;
   FillModelInfo(&onnx_model);
 

diff --git a/caffe2/opt/onnxifi_transformer.h b/caffe2/opt/onnxifi_transformer.h
index 915accd..c93fd7f 100644 (file)
--- a/caffe2/opt/onnxifi_transformer.h
+++ b/caffe2/opt/onnxifi_transformer.h
@@ -21,9 +21,20 @@ struct OnnxifiTransformerOptions {
 
   // Dump onnx model for debugging
   bool debug{false};
+
   // Pass serialized onnx model if true, otherwise pass serialized c2 model
   bool use_onnx{true};
 
+  // Whether to attach AdjustBatch ops or not. In order to maintain static
+  // shapes to the backend, most of the time, we need to add AdjustBatch ops to
+  // the inputs/outputs of the Onnxifi op. But if backend itself supports max
+  // batch size, we don't need to do it.
+  bool add_adjust_batch_ops{true};
+
+  // Minimum number of ops to create an onnxifi op. If the subgraph is too
+  // small, it doesn't make sense to lower it to backend.
+  size_t min_ops{1};
+
   // Bound shape spec
   BoundShapeSpec bound_shape_spec;
 };