Changed the how unown_tensor attribute is set on TRT mod

Author: cehongwang

core/runtime/execute_engine.cpp

@@ -96,7 +96,7 @@ void setup_input_tensors(
     std::vector<at::Tensor> inputs,
     c10::intrusive_ptr<TRTEngine> compiled_engine,
     bool cudagraphs_enabled,
-    bool need_cudagraphs_record) {
+    bool shape_changed) {
   // this is a buffer to store shape tensor input addresses throughout the runtime scope
   std::list<std::vector<int64_t>> inputShapeTensorValues;
   std::list<at::Tensor> formatted_inputs(compiled_engine->num_io.first);
@@ -140,12 +140,14 @@ void setup_input_tensors(
     } else {
       at::Tensor contig_input = inputs[i].view(shape).contiguous();
       formatted_inputs.emplace_back(std::move(contig_input));
-
+      bool need_cudagraphs_record = cudagraphs_enabled &&
+          (!compiled_engine->runtime_states.old_cudagraphs || shape_changed ||
+           compiled_engine->runtime_states.context_changed);
       if (need_cudagraphs_record) {
         // Create a new persistent input buffer
         compiled_engine->input_buffers[i] = std::move(formatted_inputs.back().clone());
       }
-      if (need_cudagraphs_record or compiled_engine->allocated_outputs.size() == 0) {
+      if (shape_changed || compiled_engine->allocated_outputs.size() == 0) {
         TORCHTRT_CHECK(
             compiled_engine->exec_ctx->setInputShape(name.c_str(), dims), "Error while setting the input shape");
       }
@@ -226,7 +228,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
         input_profiler_guard =
             std::make_unique<torch::autograd::profiler::RecordProfile>(compiled_engine->input_profile_path);
       }
-      setup_input_tensors(inputs, compiled_engine, cudagraphs_enabled, need_cudagraphs_record);
+      setup_input_tensors(inputs, compiled_engine, cudagraphs_enabled, shape_changed);
       // Check if input shapes can be inferred.
       int32_t const io_size{compiled_engine->io_size};
       std::vector<char const*> names(io_size);

py/torch_tensorrt/dynamo/_compiler.py

@@ -949,7 +949,7 @@ def preserve_module_specs(
     for attr in dir(gm):
         if attr.startswith("_frozen_param"):
             delattr(gm, attr)
-    trt_module = None
+
     for name, _ in partitioned_module.named_children():
         submodule = getattr(partitioned_module, name)
         # filter on the GraphModule
@@ -1082,8 +1082,12 @@ def preserve_module_specs(
             trt_module = getattr(partitioned_module, name)
             trt_module.setup_engine()
 
-    if trt_module:
-        trt_module.set_output_tensors_as_unowned(True)
+    output_node = list(partitioned_module.graph.nodes)[-1]
+    for arg in output_node.args:
+        target = arg[0].target
+        if "_run_on_acc" not in str(target):
+            continue
+        getattr(partitioned_module, target).set_output_tensors_as_unowned(True)
 
     # Reset settings object to user specification after fallback to global partitioning mode
     if fast_partitioner_failed:

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

@@ -383,8 +383,13 @@ def setup_input_tensors(
         self,
         contiguous_inputs: List[torch.Tensor],
         cudagraphs_enabled: bool,
-        need_cudagraphs_record: bool,
+        shape_changed: bool = True,
     ) -> None:
+        need_cudagraphs_record = cudagraphs_enabled and (
+            not self.runtime_states.old_cudagraphs
+            or shape_changed
+            or self.runtime_states.context_changed
+        )
         for i, input_name in enumerate(self.input_names):
             if not contiguous_inputs[i].is_cuda:
                 logger.warning(
@@ -417,9 +422,7 @@ def setup_input_tensors(
                 inputs_cpu = contiguous_inputs[i].cpu().to(torch.int64).numpy().copy()
                 self.context.set_tensor_address(input_name, inputs_cpu.ctypes.data)
             else:
-                if (
-                    need_cudagraphs_record or self.output_tensors is None
-                ):  # First time execution:
+                if shape_changed or self.output_tensors is None:
                     self.context.set_input_shape(
                         input_name, tuple(contiguous_inputs[i].shape)
                     )
@@ -490,9 +493,7 @@ def run_standard_execution() -> torch.Tensor | Tuple[torch.Tensor, ...]:
                 ), f"Wrong number of inputs, expect {len(self.input_names)} get {len(contiguous_inputs)}."
 
                 self.setup_input_tensors(
-                    contiguous_inputs,
-                    self.cudagraphs_enabled,
-                    need_cudagraphs_record,
+                    contiguous_inputs, self.cudagraphs_enabled, shape_changed
                 )
 
                 if shape_changed:
@@ -807,3 +808,6 @@ def validate_input_shapes(self, inputs: Sequence[torch.Tensor]) -> bool:
             return True
 
         return False
+
+    def are_output_tensors_unowned(self) -> bool:
+        return self.output_tensors_are_unowned

setup.py

@@ -195,10 +195,10 @@ def build_libtorchtrt_cxx11_abi(
     else:
         cmd.append("//:libtorchtrt")
 
-    # if develop:
-    #     cmd.append("--compilation_mode=dbg")
-    # else:
-    cmd.append("--compilation_mode=opt")
+    if develop:
+        cmd.append("--compilation_mode=dbg")
+    else:
+        cmd.append("--compilation_mode=opt")
     if use_dist_dir:
         if IS_AARCH64:
             cmd.append("--distdir=third_party/dist_dir/aarch64-linux-gnu")