Add rule based filtering for ONNXIFI transformation (pytorch#17198)

Summary:
Pull Request resolved: pytorch#17198

We come to the point that we need to apply some rules to bind certain ops together to avoid un-inferrable intermediate shapes. We either lower them together to backend or neither. This diff adds a pass for us to add rules like this. The first one is to bind `Gather` with `SparseLengthsWeighted*`.

Reviewed By: ipiszy

Differential Revision: D14118326

fbshipit-source-id: 14bc62e1feddae02a3dd8eae93b8f553d52ac951

Author: Yinghai Lu

caffe2/onnx/onnxifi_graph_info.h

@@ -16,21 +16,26 @@ struct BackendGraphInfo {
   onnxBackend backend;
   onnxGraph graph;
   onnxifi_library* lib{nullptr};
+
   BackendGraphInfo(
       onnxBackendID backend_id,
       onnxBackend backend,
       onnxGraph graph,
       onnxifi_library* lib)
       : backend_id(backend_id), backend(backend), graph(graph), lib(lib) {}
+
   BackendGraphInfo(const BackendGraphInfo& other) = delete;
+
   BackendGraphInfo& operator=(const BackendGraphInfo& other) = delete;
+
   BackendGraphInfo(BackendGraphInfo&& other) noexcept {
     backend_id = other.backend_id;
     backend = other.backend;
     graph = other.graph;
     lib = other.lib;
     other.backend_id = other.backend = other.graph = other.lib = nullptr;
   }
+
   BackendGraphInfo& operator=(BackendGraphInfo&& other) {
     backend_id = other.backend_id;
     backend = other.backend;
@@ -39,6 +44,7 @@ struct BackendGraphInfo {
     other.backend_id = other.backend = other.graph = other.lib = nullptr;
     return *this;
   }
+
   ~BackendGraphInfo() {
     if (lib) {
       onnxStatus err;
@@ -74,14 +80,17 @@ class OnnxBackendGraphMap {
   OnnxBackendGraphMap(OnnxBackendGraphMap&&) = delete;
   OnnxBackendGraphMap operator=(const OnnxBackendGraphMap&) = delete;
   OnnxBackendGraphMap operator=(OnnxBackendGraphMap&&) = delete;
+
   SharedPtrBackendGraphInfo lookup(const std::string& key);
+
   // If acquisition of graph_ptr fails then graph already exists. And the
   // corresponding SharedPtrBackendGraphInfo is returned. Otherwise graph_ptr is
   // inserted in the map and the wrapper SharedPtrBackendGraphInfo is
   // returned.teebfhlbnheuk
   SharedPtrBackendGraphInfo insert(
       const std::string& key,
       BackendGraphInfo graph);
+
   void remove(const std::string& key);
 
  private:

caffe2/operators/onnxifi_op.h

@@ -171,19 +171,28 @@ class OnnxifiOp final : public Operator<Context> {
         }
         lib_->onnxReleaseBackendID(backend_ids[i]);
       }
-      CAFFE_ENFORCE_EQ(
-          lib_->onnxInitGraph(
-              backend_,
-              nullptr,
-              onnx_model_str.size(),
-              (const void*)(onnx_model_str.c_str()),
-              weight_descs.size(),
-              weight_descs.data(),
-              &graph_),
-          ONNXIFI_STATUS_SUCCESS);
-      backend_graph_shared_ptr_ = backend_graph_map_ptr_->insert(
-          op_id_string_,
-          onnx::BackendGraphInfo(backend_id_, backend_, graph_, lib_));
+
+      // Lookup the backend first, if it's not there, create our own and try
+      // submitting it to the backend_graph_map
+      backend_graph_shared_ptr_ = backend_graph_map_ptr_->lookup(op_id_string_);
+      if (!backend_graph_shared_ptr_) {
+        LOG(INFO) << "Creating backend for " << op_id_string_;
+        CAFFE_ENFORCE_EQ(
+            lib_->onnxInitGraph(
+                backend_,
+                nullptr,
+                onnx_model_str.size(),
+                (const void*)(onnx_model_str.c_str()),
+                weight_descs.size(),
+                weight_descs.data(),
+                &graph_),
+            ONNXIFI_STATUS_SUCCESS);
+        backend_graph_shared_ptr_ = backend_graph_map_ptr_->insert(
+            op_id_string_,
+            onnx::BackendGraphInfo(backend_id_, backend_, graph_, lib_));
+      } else {
+        LOG(INFO) << "Got cached backend for " << op_id_string_;
+      }
       // This checks if our insertion was successful or some other thread did
       // the insert in the meantime.
       if (backend_graph_shared_ptr_->backend_id != backend_id_ ||

caffe2/opt/backend_transformer_base.cc

@@ -25,7 +25,7 @@ std::string BackendTransformerBase::getModelId(const NetDef& net) {
 
 TensorProto BackendTransformerBase::wrapShapeInfoIntoTensorProto(
     const std::string& name,
-    const ShapeInfo& shape_info) {
+    const ShapeInfo& shape_info) const {
   TensorProto t;
   t.set_name(name);
   t.set_data_type(shape_info.shape.data_type());

caffe2/opt/backend_transformer_base.h

@@ -52,7 +52,7 @@ class BackendTransformerBase {
   // Wrap TensorShape into TensorProto
   TensorProto wrapShapeInfoIntoTensorProto(
       const std::string& name,
-      const ShapeInfo& shape_info);
+      const ShapeInfo& shape_info) const;
 
   // Do bound shape inference and collect shape infos
   ShapeInfoMap inferShapes(