Renamed to be less confusing

Author: cehongwang

core/runtime/TRTEngine.cpp

@@ -289,12 +289,12 @@ void TRTEngine::enable_profiling() {
   exec_ctx->setProfiler(trt_engine_profiler.get());
 }
 
-void TRTEngine::set_unowned_output_tensor(bool enable) {
-  this->unowned_output_tensor = enable;
+void TRTEngine::set_output_tensors_as_unowned(bool enable) {
+  this->output_tensors_are_unowned = enable;
 }
 
-bool TRTEngine::is_unowned_output_tensor() {
-  return this->unowned_output_tensor;
+bool TRTEngine::are_output_tensors_unowned() {
+  return this->output_tensors_are_unowned;
 }
 
 void TRTEngine::set_profile_format(std::string format) {

core/runtime/TRTEngine.h

@@ -105,7 +105,7 @@ struct TRTEngine : torch::CustomClassHolder {
   std::pair<uint64_t, uint64_t> num_io;
   uint64_t io_size;
   std::map<std::string, bool> isShapeInferenceIO;
-  bool unowned_output_tensor = false;
+  bool output_tensors_are_unowned = false;
   std::string name;
   RTDevice device_info;
 
@@ -162,8 +162,8 @@ struct TRTEngine : torch::CustomClassHolder {
   int64_t get_automatic_device_memory_budget();
   std::vector<at::Tensor> infer_outputs(std::vector<std::vector<int64_t>> input_shapes);
   void set_pre_allocated_outputs(bool enable);
-  void set_unowned_output_tensor(bool enable);
-  bool is_unowned_output_tensor();
+  void set_output_tensors_as_unowned(bool enable);
+  bool are_output_tensors_unowned();
   TorchTRTRuntimeStates runtime_states;
   friend std::ostream& operator<<(std::ostream& os, const TRTEngine& engine);
   static const char BINDING_DELIM = '%';

core/runtime/execute_engine.cpp

@@ -249,7 +249,8 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
       if (can_use_pre_allocated_outputs) {
         outputs = compiled_engine->pre_allocated_outputs;
       } else {
-        if (compiled_engine->allocated_outputs.size() == 0 or compiled_engine->unowned_output_tensor or shape_changed) {
+        if (compiled_engine->allocated_outputs.size() == 0 or compiled_engine->output_tensors_are_unowned or
+            shape_changed) {
           compiled_engine->allocated_outputs = create_output_tensors(compiled_engine);
           new_outputs = true;
         }

core/runtime/register_jit_hooks.cpp

@@ -90,8 +90,8 @@ static auto TORCHTRT_UNUSED TRTEngineTSRegistrtion =
         .def("get_engine_layer_info", &TRTEngine::get_engine_layer_info)
         .def("infer_outputs", &TRTEngine::infer_outputs)
         .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("set_output_tensors_as_unowned", &TRTEngine::set_output_tensors_as_unowned)
+        .def("are_output_tensors_unowned", &TRTEngine::are_output_tensors_unowned)
         .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/_compiler.py

@@ -1073,7 +1073,7 @@ def preserve_module_specs(
 
     # Only set the requires_unique_output flag for the last TRT Module when user has access to the output tensor
     if trt_module:
-        trt_module.set_unowned_output_tensor(True)
+        trt_module.set_output_tensors_as_unowned(True)
 
     # Parse the graph I/O and store it in dryrun tracker
     parse_graph_io(gm, dryrun_tracker)

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

@@ -221,28 +221,28 @@ def __init__(
         self.use_output_allocator_outputs = False
         self.device = torch.cuda.current_device()
         self.cudagraphs_enabled = torch_tensorrt.runtime.get_cudagraphs_mode()
-        # If the output tensor is not owned by the engine (unowned_output_tensor=True), we need to create a new output tensor in each forward pass
-        self.unowned_output_tensor = False
+        # If the output tensor is not owned by the engine (output_tensors_are_unowned=True), we need to create a new output tensor in each forward pass
+        self.output_tensors_are_unowned = False
         if self.serialized_engine is not None and not self.settings.lazy_engine_init:
             self.setup_engine()
         self.is_shape_inference_io = {
             input_name: self.engine.is_shape_inference_io(input_name)
             for input_name in self.input_names
         }
 
-    def set_unowned_output_tensor(self, enabled: bool) -> None:
+    def set_output_tensors_as_unowned(self, enabled: bool) -> None:
         """
-        Set the flag to indicate if the output tensor is unowned by the engine.
-        If self.unowned_output_tensor=True, the engine will create a new output tensor in each forward pass.
-        This would be slower but is required when users need to manipulate the output tensor after each forward pass.
-        Therefore, this should be set to True only for the last module in a graph and leave to False for intermediate modules,
-        which users don't have access to.
+        Flag to set if the output tensors of this engine are solely owned by the Torch-TensorRT Runtime or if they might be shared with a user.
+        If the tensors are not owned by the runtime, then they must be recreated on every forward call which may have implications for performance.
+        Typically only the final engine in a graph requires output tensors to be unowned and there are performance gains to be had for intermediate engines to manage their own standing memory.
+        Therefore this should only be set to True for the final module in a graph and leave false for intermediate modules.
+
         Args:
             enabled: bool
                 Whether to set the flag to True.
 
         """
-        self.unowned_output_tensor = enabled
+        self.output_tensors_are_unowned = enabled
 
     def get_streamable_device_memory_budget(self) -> Any:
         return self.engine.streamable_weights_size
@@ -528,7 +528,7 @@ def run_standard_execution() -> torch.Tensor | Tuple[torch.Tensor, ...]:
                         )
                     if (
                         self.output_tensors is None
-                        or self.unowned_output_tensor
+                        or self.output_tensors_are_unowned
                         or shape_changed
                     ):
                         self.output_tensors = self.create_output_tensors()

py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py

@@ -289,8 +289,8 @@ def set_extra_state(self, state: SerializedTorchTensorRTModuleFmt) -> None:
             metadata = TorchTensorRTModule.decode_metadata(serialized_metadata)
             self.settings = metadata["settings"]
             self.weight_name_map = metadata["weight_name_map"]
-            self.unowned_output_tensor = metadata["unowned_output_tensor"]
-            self.engine.set_unowned_output_tensor(self.unowned_output_tensor)
+            self.output_tensors_are_unowned = metadata["output_tensors_are_unowned"]
+            self.engine.set_output_tensors_as_unowned(self.output_tensors_are_unowned)
 
         else:
             self.engine = None
@@ -362,11 +362,11 @@ def enable_profiling(
         self.engine.enable_profiling()
         self.engine.set_profile_format(profile_format)
 
-    def set_unowned_output_tensor(self, enabled: bool) -> None:
-        self.engine.set_unowned_output_tensor(enabled)
+    def set_output_tensors_as_unowned(self, enabled: bool) -> None:
+        self.engine.set_output_tensors_as_unowned(enabled)
 
-    def is_unowned_output_tensor(self) -> bool:
-        return self.engine.is_unowned_output_tensor()  # type: ignore[no-any-return]
+    def are_output_tensors_unowned(self) -> bool:
+        return self.engine.are_output_tensors_unowned()  # type: ignore[no-any-return]
 
     def disable_profiling(self) -> None:
         """Disable the profiler"""