feat(to_backend): Updating backend integration preproc function

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

py/setup.py

@@ -44,10 +44,12 @@ def is_exe(fpath):
     return None
 
 
-BAZEL_EXE = which("bazel")
+BAZEL_EXE = which("bazelisk")
 
 if BAZEL_EXE is None:
-    sys.exit("Could not find bazel in PATH")
+    BAZEL_EXE = which("bazel")
+    if BAZEL_EXE is None:
+        sys.exit("Could not find bazel in PATH")
 
 
 def build_libtrtorch_pre_cxx11_abi(develop=True, use_dist_dir=True, cxx11_abi=False):
@@ -207,7 +209,7 @@ def run(self):
       long_description=long_description,
       ext_modules=ext_modules,
       install_requires=[
-          'torch>=1.8.0+cu111,<1.9.0',
+          'torch>=1.9.0+cu111,<1.10.0',
       ],
       setup_requires=[],
       cmdclass={

py/trtorch/csrc/tensorrt_backend.cpp

@@ -10,17 +10,17 @@
 namespace trtorch {
 namespace backend {
 
-c10::IValue TensorRTBackend::preprocess(c10::IValue mod, c10::impl::GenericDict method_compile_spec) {
-  auto spec = c10::impl::toTypedDict<std::string, at::IValue>(method_compile_spec);
-
-  for (auto it = spec.begin(), end = spec.end(); it != end; ++it) {
+namespace {
+c10::IValue preprocess(const torch::jit::Module& mod, const c10::Dict<c10::IValue, c10::IValue>& method_compile_spec) {
+  for (auto it = method_compile_spec.begin(), end = method_compile_spec.end(); it != end; ++it) {
     TRTORCH_CHECK(
-        core::CheckMethodOperatorSupport(mod.toModule(), it->key()),
-        "Method " << it->key() << "cannot be compiled by TRTorch");
+        core::CheckMethodOperatorSupport(mod, it->key().toStringRef()),
+        "Method " << it->key().toStringRef() << "cannot be compiled by TRTorch");
   }
 
-  return mod;
+  return mod._ivalue();
 }
+} // namespace
 
 c10::impl::GenericDict TensorRTBackend::compile(c10::IValue mod_val, c10::impl::GenericDict method_compile_spec) {
   auto mod = mod_val.toModule();
@@ -78,6 +78,7 @@ c10::impl::GenericList TensorRTBackend::execute(c10::IValue handle, c10::impl::G
 
 namespace {
 static auto reg = torch::jit::backend<TensorRTBackend>("tensorrt");
+static auto preproc_reg = torch::jit::backend_preprocess_register("tensorrt", &preprocess);
 }
 
 } // namespace backend

py/trtorch/csrc/tensorrt_backend.h

@@ -1,6 +1,8 @@
 #pragma once
 #include "torch/csrc/jit/api/module.h"
 #include "torch/csrc/jit/backends/backend.h"
+#include "torch/csrc/jit/backends/backend_debug_handler.h"
+#include "torch/csrc/jit/backends/backend_preprocess.h"
 
 namespace trtorch {
 namespace backend {
@@ -10,7 +12,10 @@ class TensorRTBackend : public torch::jit::PyTorchBackendInterface {
   explicit TensorRTBackend() {}
   virtual ~TensorRTBackend() = default;
 
-  c10::IValue preprocess(c10::IValue mod, c10::impl::GenericDict method_compile_spec) override;
+  bool is_available() override {
+    return true;
+  }
+
   c10::impl::GenericDict compile(c10::IValue processed_mod, c10::impl::GenericDict method_compile_spec) override;
   c10::impl::GenericList execute(c10::IValue handle, c10::impl::GenericList inputs) override;
 };