Support passing onnxifi backend through the GFLAG/LLVM param.

Summary:
Support manual override of onnxifi backend.
This would allow flexibility when there is a need to use onnxifi with private backend.

Reviewed By: yinghai

Test Plan: CI

Differential Revision: D16215878

Pulled By: rdzhabarov

fbshipit-source-id: 3635ab3dc3ae5be1dac38768cd5b3bf41e9376aa

Author: Roman Dzhabarov

lib/Onnxifi/onnxifiGlow.cpp

@@ -16,6 +16,7 @@
 
 #include "Base.h"
 #include "GlowOnnxifiManager.h"
+#include "llvm/Support/CommandLine.h"
 
 #include "glow/Importer/ONNXIFIModelLoader.h"
 
@@ -28,6 +29,18 @@
 
 #define EXTERNC extern "C"
 
+namespace glow {
+namespace onnxifi {
+
+std::string GlowOnnxifiBackend = "";
+static llvm::cl::opt<std::string, /*external storage*/ true>
+    GlowOnnxifiBackendOpt("glow-onnxifi-backend",
+                          llvm::cl::desc("Glow backend used for ONNXIFI"),
+                          llvm::cl::location(GlowOnnxifiBackend));
+
+} // namespace onnxifi
+} // namespace glow
+
 /**
  * This file contains implementation of the onnxifi interface.
  * Documentation on the functions implementing onnxifi interface in
@@ -57,67 +70,50 @@ GLOW_ONNXIFI_LIBRARY_FUNCTION_WRAPPER(onnxGetBackendIDs)(
   const size_t numBackendsCapacity = *numBackends;
 
   using namespace glow::runtime;
+  using namespace glow::onnxifi;
   const bool withCPU = DeviceManager::numDevices("CPU") > 0;
   const bool withHabana = DeviceManager::numDevices("Habana") > 0;
 
   // Only return quantization backend if GLOW_DUMP_PROFILE.
   if (getenv("GLOW_DUMP_PROFILE")) {
-    *numBackends = 2;
+    *numBackends = 1;
 
     // In case backendIDs is nullptr or does not have enough capacity just
     // return the total number of supported backends.
     if (numBackendsCapacity < *numBackends || !backendIDs) {
       return ONNXIFI_STATUS_FALLBACK;
     }
 
-    auto *quantizationBackendOnnx = manager.createBackend(
-        "Interpreter", /*useOnnx*/ true, /*forQuantization*/ true);
     auto *quantizationBackendC2 = manager.createBackend(
-        "Interpreter", /*useOnnx*/ false, /*forQuantization*/ true);
-
-    backendIDs[0] = quantizationBackendOnnx;
-    backendIDs[1] = quantizationBackendC2;
-  } else if (withCPU || withHabana) {
-    *numBackends = 4;
+        GlowOnnxifiBackend.empty() ? "Interpreter" : GlowOnnxifiBackend,
+        /*useOnnx*/ false, /*forQuantization*/ true);
 
-    auto backendName = withHabana ? "Habana" : "CPU";
-
-    // In case backendIDs is nullptr or does not have enough capacity just
-    // return the total number of supported backends.
-    if (numBackendsCapacity < *numBackends || !backendIDs) {
-      return ONNXIFI_STATUS_FALLBACK;
-    }
-
-    auto *cpuBackendOnnx = manager.createBackend(backendName,
-                                                 /*useOnnx*/ true);
-    auto *interpreterBackendOnnx =
-        manager.createBackend("Interpreter", /*useOnnx*/ true);
-    auto *cpuBackendC2 = manager.createBackend(backendName,
-                                               /*useOnnx*/ false);
-    auto *interpreterBackendC2 =
-        manager.createBackend("Interpreter", /*useOnnx*/ false);
-
-    backendIDs[0] = cpuBackendOnnx;
-    backendIDs[1] = interpreterBackendOnnx;
-    backendIDs[2] = cpuBackendC2;
-    backendIDs[3] = interpreterBackendC2;
+    backendIDs[0] = quantizationBackendC2;
   } else {
-
-    *numBackends = 2;
+    *numBackends = 1;
+
+    auto backendName = GlowOnnxifiBackend;
+
+    if (backendName.empty()) {
+      if (withHabana) {
+        backendName = "Habana";
+      } else if (withCPU) {
+        backendName = "CPU";
+      } else {
+        backendName = "Interpreter";
+      }
+    }
 
     // In case backendIDs is nullptr or does not have enough capacity just
     // return the total number of supported backends.
     if (numBackendsCapacity < *numBackends || !backendIDs) {
       return ONNXIFI_STATUS_FALLBACK;
     }
 
-    auto *interpreterBackendOnnx =
-        manager.createBackend("Interpreter", /*useOnnx*/ true);
-    auto *interpreterBackendC2 =
-        manager.createBackend("Interpreter", /*useOnnx*/ false);
+    auto *executionBackend = manager.createBackend(backendName,
+                                                   /*useOnnx*/ false);
 
-    backendIDs[0] = interpreterBackendOnnx;
-    backendIDs[1] = interpreterBackendC2;
+    backendIDs[0] = executionBackend;
   }
 
   return ONNXIFI_STATUS_SUCCESS;