Enabled alias IO for python runtime and fixed comments

Author: cehongwang

core/runtime/TRTEngine.cpp

@@ -284,11 +284,6 @@ TRTEngine::TRTEngine(
   }
 
   has_dynamic_inputs = engine_has_dynamic_inputs(cuda_engine.get(), in_binding_names);
-  // Cache input binding metadata used by the runtime hot path, then reconstruct
-  // optimization-profile info from the TRT API so multi-profile selection works
-  // for any loaded engine.
-  this->setup_input_binding_infos();
-  this->setup_optimization_profiles();
 
   // Store the build-time aliased_io map as the starting point. Then reconcile
   // against ICudaEngine::getAliasedInputTensor — that API is the source of
@@ -308,6 +303,17 @@ TRTEngine::TRTEngine(
     if (it == this->aliased_io.end()) {
       this->aliased_io[out_name] = AliasedIOSpec{std::string(aliased_in), AliasKind::kKVCacheUpdate};
       LOG_DEBUG("aliased_io reconciliation: discovered " << out_name << " -> " << aliased_in << " (kv_cache_update)");
+    } else if (it->second.kind == AliasKind::kUser) {
+      // A kUser alias is declared by Torch-TensorRT, not TRT. If the engine
+      // also reports an alias for this output the two views should agree on the
+      // source input; keep the kUser kind either way so the provenance recorded
+      // at build time is not silently rewritten to kv_cache_update.
+      if (it->second.input_binding_name != std::string(aliased_in)) {
+        LOG_WARNING(
+            "aliased_io: user-declared alias for output " << out_name << " (input: " << it->second.input_binding_name
+                                                          << ") disagrees with engine-reported input " << aliased_in
+                                                          << "; keeping the user-declared value.");
+      }
     } else if (it->second.input_binding_name != std::string(aliased_in)) {
       LOG_WARNING(
           "aliased_io: build-time map disagrees with engine for output "
@@ -324,6 +330,19 @@ TRTEngine::TRTEngine(
                           << " is incompatible with dynamic output allocator. Aliasing requires fixed output shape.");
   }
 
+  // Precompute the set of input binding names that are the alias source of some
+  // output so the per-call input-setup loop can test membership in O(1).
+  this->aliased_input_binding_names.clear();
+  for (const auto& kv : this->aliased_io) {
+    this->aliased_input_binding_names.insert(kv.second.input_binding_name);
+  }
+
+  // Cache input binding metadata used by the runtime hot path (depends on
+  // aliased_input_binding_names above), then reconstruct optimization-profile
+  // info from the TRT API so multi-profile selection works for any loaded engine.
+  this->setup_input_binding_infos();
+  this->setup_optimization_profiles();
+
 #ifndef NDEBUG
   this->enable_profiling();
 #endif
@@ -611,7 +630,8 @@ void TRTEngine::setup_input_binding_infos() {
     input_binding_infos.push_back(
         {name,
          util::TRTDataTypeToScalarType(cuda_engine->getTensorDataType(name.c_str())),
-         cuda_engine->isShapeInferenceIO(name.c_str())});
+         cuda_engine->isShapeInferenceIO(name.c_str()),
+         aliased_input_binding_names.find(name) != aliased_input_binding_names.end()});
   }
 }
 

core/runtime/TRTEngine.h

@@ -6,6 +6,7 @@
 #include <memory>
 #include <mutex>
 #include <unordered_map>
+#include <unordered_set>
 #include <utility>
 
 #include "ATen/core/function_schema.h"
@@ -110,7 +111,9 @@ struct TorchTRTRuntimeStates {
     if (new_cudagraphs && (!old_cudagraphs || shape_changed || context_changed)) {
       need_cudagraphs_record = true;
     }
-    // Pre-allocated output can be used when previous and current state are true without shape change
+    // Pre-allocated output can be used when previous and current state are true without shape change.
+    // Note: engines with aliased I/O disable pre-allocation entirely; that gating lives in
+    // execute_engine (which has access to the engine's aliased_io map) rather than here.
     if (old_pre_allocated_outputs && new_pre_allocated_output && !shape_changed) {
       can_use_pre_allocated_outputs = true;
     }
@@ -181,6 +184,12 @@ struct TRTEngine : torch::CustomClassHolder {
   // the same device pointer as the source input.
   std::unordered_map<std::string, AliasedIOSpec> aliased_io = {};
 
+  // The set of input binding names that are the alias source of some output.
+  // Derived once from aliased_io at construction so the per-call input-setup
+  // loop can test membership in O(1) instead of scanning aliased_io on every
+  // input, every execution.
+  std::unordered_set<std::string> aliased_input_binding_names = {};
+
   bool hardware_compatible = false; // Whether the engine was compiled in hardware compatible mode
   std::string serialized_metadata; // This is a base64 encoded pkl object used to store metadata such as settings used
                                    // in compilation
@@ -265,6 +274,10 @@ struct TRTEngine : torch::CustomClassHolder {
     std::string name;
     at::ScalarType expected_type;
     bool is_shape_tensor;
+    // True when this input is the alias source of some output binding. Precomputed
+    // here so the per-call input-setup loop avoids an aliased_input_binding_names
+    // lookup on every input, every execution.
+    bool is_aliased_input;
   };
   std::vector<InputBindingInfo> input_binding_infos = {};
   std::vector<at::Tensor> active_input_tensors = {};

core/runtime/execute_engine.cpp

@@ -145,14 +145,10 @@ void setup_input_tensors(
       // aliased input we deliberately bind to the user's tensor instead, so
       // the engine writes through to the user's storage. The user is already
       // required to pass stable input addresses under cudagraphs, so the
-      // aliasing contract is compatible.
-      bool is_aliased_input = false;
-      for (const auto& kv : compiled_engine->aliased_io) {
-        if (kv.second.input_binding_name == name) {
-          is_aliased_input = true;
-          break;
-        }
-      }
+      // aliasing contract is compatible. Membership is precomputed at engine
+      // construction (InputBindingInfo::is_aliased_input) to avoid rescanning
+      // aliased_io on every input, every execution.
+      const bool is_aliased_input = binding.is_aliased_input;
 
       if (need_cudagraphs_record && !is_aliased_input) {
         // Create a persistent CUDA graph input staging buffer with a stable replay address.
@@ -346,6 +342,23 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
     bool can_use_pre_allocated_outputs = std::get<1>(result);
     bool need_cudagraphs_reset = std::get<2>(result);
 
+    // Pre-allocated outputs are disabled entirely for engines with aliased I/O.
+    // An aliased output reuses the caller's input storage, so there is no
+    // allocation to amortize by caching, and caching would pin a caller-owned
+    // tensor across calls. Only the non-aliased outputs could benefit, and that
+    // win is too small to justify the extra staleness/lifetime surface. See the
+    // cache-refresh guard below, which is the counterpart to this one.
+    if (!compiled_engine->aliased_io.empty()) {
+      if (compiled_engine->use_pre_allocated_outputs) {
+        LOG_WARNING(
+            "pre_allocated_outputs is enabled but this engine has aliased I/O; "
+            "pre-allocation is disabled for aliased engines (aliased outputs reuse "
+            "the caller's input storage, so there is nothing to pre-allocate). "
+            "Outputs are allocated fresh each call.");
+      }
+      can_use_pre_allocated_outputs = false;
+    }
+
     if (need_cudagraphs_reset) {
       compiled_engine->cudagraph.reset();
     }
@@ -367,8 +380,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
             std::make_unique<torch::autograd::profiler::RecordProfile>(compiled_engine->input_profile_path);
       }
 
-      setup_input_tensors(
-          inputs, compiled_engine, effective_cudagraphs, need_cudagraphs_record, bound_inputs_by_name);
+      setup_input_tensors(inputs, compiled_engine, effective_cudagraphs, need_cudagraphs_record, bound_inputs_by_name);
       // Check if input shapes can be inferred.
       int32_t const io_size{compiled_engine->cuda_engine->getNbIOTensors()};
       std::vector<char const*> names(io_size);
@@ -387,6 +399,9 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
             std::make_unique<torch::autograd::profiler::RecordProfile>(compiled_engine->output_profile_path);
       }
       if (can_use_pre_allocated_outputs) {
+        // Never reached for engines with aliased I/O: can_use_pre_allocated_outputs
+        // is forced false above when aliased_io is non-empty, so cached outputs
+        // never hold an aliased (caller-owned) slot.
         outputs = compiled_engine->pre_allocated_outputs;
       } else {
         outputs = create_output_tensors(compiled_engine, bound_inputs_by_name);
@@ -493,16 +508,14 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
     compiled_engine->clear_active_input_tensors();
 
     // When the pre-allocated output mode is turned on, for intermediate modules, we only create the output in the first
-    // execution or when shape is changed. If the engine has aliased outputs we
-    // disable pre-allocation entirely: aliased outputs share storage with
-    // user-supplied inputs that may change between calls, so caching the
-    // tensor reference would lead to writes against stale storage.
-    if (compiled_engine->use_pre_allocated_outputs && !compiled_engine->aliased_io.empty()) {
-      LOG_DEBUG(
-          "Skipping pre_allocated_outputs cache because engine has aliased I/O; "
-          "aliased outputs reuse the user's input storage on every call.");
-    } else if (
-        compiled_engine->use_pre_allocated_outputs &&
+    // execution or when shape is changed.
+    //
+    // Disabled entirely for engines with aliased I/O: an aliased output reuses
+    // the caller's input storage (no allocation to amortize), and caching would
+    // pin a caller-owned tensor across calls. This mirrors the
+    // can_use_pre_allocated_outputs gate above so the cache is never populated
+    // for aliased engines in the first place.
+    if (compiled_engine->use_pre_allocated_outputs && compiled_engine->aliased_io.empty() &&
         (compiled_engine->pre_allocated_outputs.size() == 0 || compiled_engine->output_tensors_are_unowned ||
          shape_changed)) {
       compiled_engine->pre_allocated_outputs = create_output_tensors(compiled_engine, bound_inputs_by_name);