Merge pull request #1219 from pytorch/anuragd/torchtrtc-custom-plugins

feat: Added support for custom torch operators and converters in torchtrtc

Author: narendasan

cpp/bin/torchtrtc/BUILD

@@ -19,6 +19,9 @@ cc_binary(
         "parser_util.h",
         "parser_util.cpp"
     ],
+    linkopts = [
+        "-l:libdl.so"
+    ],
     deps = [
         "//third_party/args",
         "//cpp:torch_tensorrt",

cpp/bin/torchtrtc/CMakeLists.txt

@@ -10,7 +10,7 @@ add_executable(${executable_name}
 if (MSVC)
     target_link_libraries(${executable_name} PRIVATE torch torchtrt)
 else()
-    target_link_libraries(${executable_name} PRIVATE torch "-Wl,--no-as-needed" torchtrt "-Wl,--as-needed")
+    target_link_libraries(${executable_name} PRIVATE torch "-Wl,--no-as-needed -ldl" torchtrt "-Wl,--as-needed")
     set_target_properties(
         ${executable_name}
         PROPERTIES INSTALL_RPATH_USE_LINK_PATH FALSE #

cpp/bin/torchtrtc/README.md

@@ -108,6 +108,8 @@ torchtrtc [input_file_path] [output_file_path]
                                         TorchScript program, save the created
                                         engine to the path specified as the
                                         output path
+      --custom-torch-ops=[lib]          (repeatable) Shared object/DLL containing custom torch operators
+      --custom-converters=[lib]         (repeatable) Shared object/DLL containing custom converters
       input_file_path                   Path to input TorchScript file
       output_file_path                  Path for compiled TorchScript (or
                                         TensorRT engine) file
@@ -131,3 +133,14 @@ e.g.
 ```
 torchtrtc tests/modules/ssd_traced.jit.pt ssd_trt.ts "[(1,3,300,300); (1,3,512,512); (1, 3, 1024, 1024)]@fp16%contiguous" -p f16
 ```
+
+
+To run with custom torch operators
+```
+torchtrtc tests/modules/ssd_traced.jit.pt ssd_trt.ts --custom-torch-ops=<path to custom library> "[(1,3,300,300); (1,3,512,512); (1, 3, 1024, 1024)]@fp16%contiguous" -p f16
+```
+
+To run with custom converters
+```
+torchtrtc tests/modules/ssd_traced.jit.pt ssd_trt.ts --custom-converters=<path to custom library> "[(1,3,300,300); (1,3,512,512); (1, 3, 1024, 1024)]@fp16%contiguous" -p f16
+```

cpp/bin/torchtrtc/main.cpp

@@ -15,6 +15,33 @@
 #include "luts.h"
 #include "parser_util.h"
 
+#if defined(_WIN32)
+#include <windows.h>
+#else
+#include <dlfcn.h>
+#endif
+
+void* load_library(std::string& custom_lib) {
+  void* handle = {nullptr};
+#if defined(_WIN32)
+  handle = LoadLibrary(custom_lib.c_str());
+#else
+  handle = dlopen(custom_lib.c_str(), RTLD_LAZY);
+#endif
+  return handle;
+}
+
+bool unload_library(void* custom_lib) {
+  bool success = false;
+#if defined(_WIN32)
+  // Returns status non-zero for success
+  success = FreeLibrary(custom_lib) ? true : false;
+#else
+  success = dlclose(custom_lib) ? false : true;
+#endif
+  return success;
+}
+
 int main(int argc, char** argv) {
   torchtrt::logging::set_is_colored_output_on(true);
   torchtrt::logging::set_reportable_log_level(torchtrt::logging::Level::kWARNING);
@@ -146,6 +173,18 @@ int main(int argc, char** argv) {
       "save_engine",
       "Instead of compiling a full a TorchScript program, save the created engine to the path specified as the output path",
       {"save-engine"});
+  args::ValueFlagList<std::string> custom_torch_ops(
+      parser,
+      "custom-torch-ops",
+      "(repeatable) Shared object/DLL containing custom torch operators",
+      {"custom-torch-ops"});
+
+  args::ValueFlagList<std::string> custom_converters(
+      parser,
+      "custom-converters",
+      "(repeatable) Shared object/DLL containing custom converters",
+      {"custom-converters"});
+
   args::Positional<std::string> input_path(parser, "input_file_path", "Path to input TorchScript file");
   args::Positional<std::string> output_path(
       parser, "output_file_path", "Path for compiled TorchScript (or TensorRT engine) file");
@@ -173,6 +212,34 @@ int main(int argc, char** argv) {
     torchtrt::logging::set_reportable_log_level(torchtrt::logging::Level::kERROR);
   }
 
+  std::vector<std::pair<std::string, void*>> custom_torch_op, custom_converter_op;
+  if (custom_torch_ops) {
+    for (auto& op : args::get(custom_torch_ops)) {
+      auto* handle = load_library(op);
+      if (handle == nullptr) {
+        torchtrt::logging::log(
+            torchtrt::logging::Level::kERROR, std::string("Could not load custom_torch_ops library " + op));
+      } else {
+        torchtrt::logging::log(torchtrt::logging::Level::kINFO, std::string("Loaded custom_torch_ops library " + op));
+
+        custom_torch_op.push_back({op, handle});
+      }
+    }
+  }
+
+  if (custom_converters) {
+    for (auto& op : args::get(custom_converters)) {
+      auto* handle = load_library(op);
+      if (handle == nullptr) {
+        torchtrt::logging::log(
+            torchtrt::logging::Level::kERROR, std::string("Could not load custom_converter library " + op));
+      } else {
+        torchtrt::logging::log(torchtrt::logging::Level::kINFO, std::string("Loaded custom_converter library " + op));
+        custom_converter_op.push_back({op, handle});
+      }
+    }
+  }
+
   auto real_input_path = torchtrtc::fileio::resolve_path(args::get(input_path));
 
   if (check_method_op_support) {
@@ -188,7 +255,7 @@ int main(int argc, char** argv) {
     auto method = args::get(check_method_op_support);
     auto result = torchtrt::ts::check_method_operator_support(mod, method);
     if (result) {
-      std::cout << "The method is supported end to end by Torch-TensorRT" << std::endl;
+      torchtrt::logging::log(torchtrt::logging::Level::kINFO, "The method is supported end to end by Torch-TensorRT");
       return 0;
     } else {
       torchtrt::logging::log(torchtrt::logging::Level::kERROR, "Method is not currently supported by Torch-TensorRT");
@@ -476,5 +543,29 @@ int main(int argc, char** argv) {
     trt_mod.save(real_output_path);
   }
 
+  if (custom_torch_ops) {
+    for (auto& p : custom_torch_op) {
+      auto status = unload_library(p.second);
+      if (status) {
+        torchtrt::logging::log(torchtrt::logging::Level::kINFO, std::string("Unloaded custom library " + p.first));
+      } else {
+        torchtrt::logging::log(
+            torchtrt::logging::Level::kERROR, std::string("Could not unload custom library " + p.first));
+      }
+    }
+  }
+
+  if (custom_converters) {
+    for (auto& p : custom_converter_op) {
+      auto status = unload_library(p.second);
+      if (status) {
+        torchtrt::logging::log(torchtrt::logging::Level::kINFO, std::string("Unloaded custom library " + p.first));
+      } else {
+        torchtrt::logging::log(
+            torchtrt::logging::Level::kERROR, std::string("Could not unload custom library " + p.first));
+      }
+    }
+  }
+
   return 0;
 }

docsrc/tutorials/torchtrtc.rst

@@ -111,6 +111,8 @@ to standard TorchScript. Load with ``torch.jit.load()`` and run like you would r
                                           TorchScript program, save the created
                                           engine to the path specified as the
                                           output path
+        --custom-torch-ops                (repeatable) Shared object/DLL containing custom torch operators
+        --custom-converters               (repeatable) Shared object/DLL containing custom converters
         input_file_path                   Path to input TorchScript file
         output_file_path                  Path for compiled TorchScript (or
                                           TensorRT engine) file
@@ -132,3 +134,13 @@ e.g.
 .. code-block:: shell
 
     torchtrtc tests/modules/ssd_traced.jit.pt ssd_trt.ts "[(1,3,300,300); (1,3,512,512); (1, 3, 1024, 1024)]@f16%contiguous" -p f16
+
+
+To run with custom torch operators
+.. code-block:: shell
+torchtrtc tests/modules/ssd_traced.jit.pt ssd_trt.ts --custom-torch-ops=<path to custom library> "[(1,3,300,300); (1,3,512,512); (1, 3, 1024, 1024)]@fp16%contiguous" -p f16
+
+
+To run with custom converters
+.. code-block:: shell
+torchtrtc tests/modules/ssd_traced.jit.pt ssd_trt.ts --custom-converters=<path to custom library> "[(1,3,300,300); (1,3,512,512); (1, 3, 1024, 1024)]@fp16%contiguous" -p f16