Optimized Cpp runtime

Author: cehongwang

core/runtime/register_jit_hooks.cpp

@@ -92,6 +92,10 @@ static auto TORCHTRT_UNUSED TRTEngineTSRegistrtion =
         .def("reset_captured_graph", &TRTEngine::reset_captured_graph)
         .def("set_unowned_output_tensor", &TRTEngine::set_unowned_output_tensor)
         .def("is_unowned_output_tensor", &TRTEngine::is_unowned_output_tensor)
+        .def(
+            "run_engine",
+            [](const c10::intrusive_ptr<TRTEngine>& self, const std::vector<at::Tensor>& input_tensors)
+                -> std::vector<at::Tensor> { return execute_engine(input_tensors, self); })
         .def_readwrite("use_pre_allocated_outputs", &TRTEngine::use_pre_allocated_outputs)
         .def_readwrite("use_output_allocator_outputs", &TRTEngine::use_output_allocator_outputs)
         .def_property(

py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py

@@ -332,9 +332,10 @@ def forward(self, *inputs: Any) -> torch.Tensor | Tuple[torch.Tensor, ...]:
             for i in inputs
         ]
 
-        outputs: List[torch.Tensor] = torch.ops.tensorrt.execute_engine(
-            list(input_tensors), self.engine
-        )
+        # outputs: List[torch.Tensor] = torch.ops.tensorrt.execute_engine(
+        #     list(input_tensors), self.engine
+        # )
+        outputs: List[torch.Tensor] = self.engine.run_engine(list(input_tensors))
 
         if len(outputs) == 1:
             return outputs[0]