Merge pull request #178 from NVIDIA/public_op_registration

refactor(//core/execution): Move to public APIs for operator registration

Author: narendasan

core/compiler.cpp

@@ -85,7 +85,7 @@ void AddEngineToGraph(torch::jit::script::Module mod, std::shared_ptr<torch::jit
     execute_node_inputs.push_back(engine_node->outputs()[0]);
 
     // Create the actual execution node trt::execute_engine using the assembled inputs
-    auto execute_node = g->create(c10::Symbol::fromQualString("trt::execute_engine"), torch::jit::ArrayRef<torch::jit::Value*>(execute_node_inputs), 1);
+    auto execute_node = g->create(c10::Symbol::fromQualString("tensorrt::execute_engine"), torch::jit::ArrayRef<torch::jit::Value*>(execute_node_inputs), 1);
     g->block()->appendNode(execute_node);
     execute_node->outputs()[0]->setType(c10::ListType::ofTensors());
 

core/execution/register_trt_op.cpp

@@ -46,33 +46,21 @@ std::vector<at::Tensor> RunCudaEngine(nvinfer1::IExecutionContext* ctx, std::pai
     return outputs;
 }
 
-namespace {
-c10::AliasAnalysisKind aliasAnalysisFromSchema() {
-  return c10::AliasAnalysisKind::FROM_SCHEMA;
-}
-
-// Switched to a global operator because op implementations need to be non-capturing lambdas in PYT 1.5.0+
-torch::jit::RegisterOperators jit_registry({
-    torch::jit::Operator(
-        "trt::execute_engine(Tensor[] inputs, __torch__.torch.classes.tensorrt.Engine engine) -> Tensor[]",
-        [](torch::jit::Stack& stack) -> int {
-            // Verify calling convention (right to left or left to right)
-            auto engine = torch::jit::pop(stack).toCustomClass<TRTEngine>();
-            LOG_DEBUG("Attempting to run engine (ID: " << std::hex << engine->name << ")");
+std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intrusive_ptr<TRTEngine> engine) {
+    // Verify calling convention (right to left or left to right)
+    LOG_DEBUG("Attempting to run engine (ID: " << std::hex << engine->name << ")");
 
-            auto inputs = torch::jit::pop(stack).toTensorVector();
+    auto io = engine->num_io;
+    auto ctx = engine->exec_ctx;
+    auto outputs = RunCudaEngine(ctx, io, inputs);
 
-            auto io = engine->num_io;
+    return outputs;
+}
 
-            auto ctx = engine->exec_ctx;
-            auto outputs = RunCudaEngine(ctx, io, inputs);
-            torch::jit::push(stack, std::move(outputs));
-            return 0;
-        },
-        aliasAnalysisFromSchema())
-    });
+TORCH_LIBRARY(tensorrt, m) {
+  m.def("execute_engine", execute_engine);
+}
 
-} // namespace
 } // namespace execution
 } // namespace core
 } // namespace trtorch