ONNXIFI extension & e2e tests. (#17478)

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

Enable onnxifi_ext in glow and build an e2e test in caffe2.

Reviewed By: yinghai

Differential Revision: D14190136

fbshipit-source-id: 26245278b487b551623109b14432f675279b17b5

diff --git a/caffe2/operators/onnxifi_op.cc b/caffe2/operators/onnxifi_op.cc
index 6ed9209..03a4289 100644 (file)
--- a/caffe2/operators/onnxifi_op.cc
+++ b/caffe2/operators/onnxifi_op.cc
@@ -151,42 +151,69 @@ bool OnnxifiOp<float, CPUContext>::RunOnDevice() {
     SetOutputTensorDescriptorTypeAndBuffer(
         type, output_tensor, &tensor_descriptor);
   }
-
-  CAFFE_ENFORCE_EQ(
-      lib_->onnxSetGraphIO(
-          graph_,
-          input_desc_.size(),
-          input_desc_.data(),
-          output_desc_.size(),
-          output_desc_.data()),
-      ONNXIFI_STATUS_SUCCESS);
-
+  bool ext_supported = false;
   onnxMemoryFenceV1 input_fence;
-  input_fence.tag = ONNXIFI_TAG_MEMORY_FENCE_V1;
-  input_fence.type = ONNXIFI_SYNCHRONIZATION_EVENT;
-  CAFFE_ENFORCE_EQ(
-      lib_->onnxInitEvent(backend_, &input_fence.event),
-      ONNXIFI_STATUS_SUCCESS);
   onnxMemoryFenceV1 output_fence;
-  output_fence.tag = ONNXIFI_TAG_MEMORY_FENCE_V1;
-  output_fence.type = ONNXIFI_SYNCHRONIZATION_EVENT;
-
-  // Call the asycn run on backend, singal event on input fence and wait for the
-  // event on output fence
-  CAFFE_ENFORCE_EQ(
-      lib_->onnxSignalEvent(input_fence.event), ONNXIFI_STATUS_SUCCESS);
-  CAFFE_ENFORCE_EQ(
-      lib_->onnxRunGraph(graph_, &input_fence, &output_fence),
-      ONNXIFI_STATUS_SUCCESS);
-  CAFFE_ENFORCE_EQ(
-      lib_->onnxWaitEvent(output_fence.event), ONNXIFI_STATUS_SUCCESS);
-
-  // Destroy the event objects
-  CAFFE_ENFORCE_EQ(
-      lib_->onnxReleaseEvent(input_fence.event), ONNXIFI_STATUS_SUCCESS);
-  CAFFE_ENFORCE_EQ(
-      lib_->onnxReleaseEvent(output_fence.event), ONNXIFI_STATUS_SUCCESS);
-
+#ifdef ONNXIFI_ENABLE_EXT
+  /**
+   * If onnxifi extension mode is enabled,
+   * and onnxSetIOAndRunGraph is supported in backend,
+   * then we run throw this workflow;
+   * Else we fallback to non-onnxifi-extension workflow.
+   **/
+  if (onnxSetIOAndRunGraphPointer_ != nullptr) {
+    ext_supported = true;
+    output_fence.tag = ONNXIFI_TAG_MEMORY_FENCE_V1;
+    output_fence.type = ONNXIFI_SYNCHRONIZATION_EVENT;
+    CAFFE_ENFORCE_EQ(
+        (*onnxSetIOAndRunGraphPointer_)(
+            graph_,
+            input_desc_.size(),
+            input_desc_.data(),
+            output_desc_.size(),
+            output_desc_.data(),
+            &output_fence),
+        ONNXIFI_STATUS_SUCCESS);
+    CAFFE_ENFORCE_EQ(
+        lib_->onnxWaitEvent(output_fence.event), ONNXIFI_STATUS_SUCCESS);
+    CAFFE_ENFORCE_EQ(
+        lib_->onnxReleaseEvent(output_fence.event), ONNXIFI_STATUS_SUCCESS);
+  }
+#endif
+  if (!ext_supported) {
+    CAFFE_ENFORCE_EQ(
+        lib_->onnxSetGraphIO(
+            graph_,
+            input_desc_.size(),
+            input_desc_.data(),
+            output_desc_.size(),
+            output_desc_.data()),
+        ONNXIFI_STATUS_SUCCESS);
+
+    input_fence.tag = ONNXIFI_TAG_MEMORY_FENCE_V1;
+    input_fence.type = ONNXIFI_SYNCHRONIZATION_EVENT;
+    CAFFE_ENFORCE_EQ(
+        lib_->onnxInitEvent(backend_, &input_fence.event),
+        ONNXIFI_STATUS_SUCCESS);
+    output_fence.tag = ONNXIFI_TAG_MEMORY_FENCE_V1;
+    output_fence.type = ONNXIFI_SYNCHRONIZATION_EVENT;
+
+    // Call the async run on backend, signal event on input fence and wait for
+    // the event on output fence
+    CAFFE_ENFORCE_EQ(
+        lib_->onnxRunGraph(graph_, &input_fence, &output_fence),
+        ONNXIFI_STATUS_SUCCESS);
+    CAFFE_ENFORCE_EQ(
+        lib_->onnxSignalEvent(input_fence.event), ONNXIFI_STATUS_SUCCESS);
+    CAFFE_ENFORCE_EQ(
+        lib_->onnxWaitEvent(output_fence.event), ONNXIFI_STATUS_SUCCESS);
+
+    // Destroy the event objects
+    CAFFE_ENFORCE_EQ(
+        lib_->onnxReleaseEvent(input_fence.event), ONNXIFI_STATUS_SUCCESS);
+    CAFFE_ENFORCE_EQ(
+        lib_->onnxReleaseEvent(output_fence.event), ONNXIFI_STATUS_SUCCESS);
+  }
   return true;
 }
 

diff --git a/caffe2/operators/onnxifi_op.h b/caffe2/operators/onnxifi_op.h
index 0d4a207..1fb5806 100644 (file)
--- a/caffe2/operators/onnxifi_op.h
+++ b/caffe2/operators/onnxifi_op.h
@@ -192,6 +192,19 @@ class OnnxifiOp final : public Operator<Context> {
     backend_id_ = backend_graph_shared_ptr_->backend_id;
     backend_ = backend_graph_shared_ptr_->backend;
     graph_ = backend_graph_shared_ptr_->graph;
+
+// Set up function pointer if onnxifi_ext is enabled
+#ifdef ONNXIFI_ENABLE_EXT
+    onnxExtensionFunctionPointer p;
+    if (lib_->onnxGetExtensionFunctionAddress(
+            backend_id_, "onnxSetIOAndRunGraphFunction", &p) !=
+        ONNXIFI_STATUS_SUCCESS) {
+      onnxSetIOAndRunGraphPointer_ = nullptr;
+      return;
+    }
+    onnxSetIOAndRunGraphPointer_ =
+        reinterpret_cast<decltype(onnxSetIOAndRunGraphPointer_)>(p);
+#endif
   }
 
   std::vector<onnxTensorDescriptorV1> buildInitializationList(
@@ -214,6 +227,17 @@ class OnnxifiOp final : public Operator<Context> {
   std::vector<onnxTensorDescriptorV1> input_desc_;
   std::vector<onnxTensorDescriptorV1> output_desc_;
 
+#ifdef ONNXIFI_ENABLE_EXT
+  // onnxifi extension mode function pointer
+  onnxStatus (*onnxSetIOAndRunGraphPointer_)(
+      onnxGraph,
+      uint32_t,
+      const onnxTensorDescriptorV1*,
+      uint32_t,
+      const onnxTensorDescriptorV1*,
+      onnxMemoryFenceV1*);
+#endif
+
   // We bind the op input/output by position while ONNXIFI binds input/output by
   // names. In addition, op input/output names can be writtten by, for example,
   // memonger. We cache the original input/output name of ONNX object here and