refactor: [collection] fuse Input with GraphInputs

Signed-off-by: inocsin <[email protected]>

Author: inocsin

core/compiler.h

@@ -14,12 +14,13 @@ namespace torch_tensorrt {
 namespace core {
 
 struct CompileSpec {
-  CompileSpec(std::vector<ir::Input> inputs) : inputs(inputs) {}
+  CompileSpec(std::vector<ir::Input> inputs) {
+    graph_inputs.inputs = inputs;
+  }
   CompileSpec(torch::jit::IValue& input_signature) {
     graph_inputs.input_signature = input_signature;
   }
   ir::GraphInputs graph_inputs;
-  std::vector<ir::Input> inputs; // can be replaced by graph_inputs
   conversion::ConversionInfo convert_info;
   lowering::LowerInfo lower_info;
   partitioning::PartitionInfo partition_info;

core/ir/ir.h

@@ -41,7 +41,7 @@ struct Input : torch::CustomClassHolder {
 // Add to spec
 struct GraphInputs {
   torch::jit::IValue input_signature;  // nested Input, full input spec
-  std::vector<Input> flattened_inputs;  // flattend Input
+  std::vector<Input> inputs;  // flattend Input
   std::vector<std::vector<Input>> collection_inputs; // only support two layer nesting, e.g. ((a, b), [c, d], e)
 };
 

cpp/include/torch_tensorrt/torch_tensorrt.h

@@ -517,10 +517,11 @@ struct TORCHTRT_API Input : torch::CustomClassHolder{
 /**
  * @brief A struct to hold complex inputs
  *
- * This struct can either hold a conplex inputs of shape or a flattened one,
+ * This struct can either hold a complex inputs of shape or a flattened one,
  */
 struct TORCHTRT_API GraphInputs {
-  torch::jit::IValue input_signature;   // nested Input, full input spec
+  torch::jit::IValue input_signature;  // nested Input, full input spec
+  std::vector<Input> inputs; // flatten input spec
 };
 
 /**
@@ -590,25 +591,17 @@ struct TORCHTRT_API CompileSpec {
    *
    * @param inputs
    */
-  CompileSpec(std::vector<Input> inputs) : inputs(std::move(inputs)) {}
+  CompileSpec(std::vector<Input> inputs);
 
   /**
    * @brief Construct a new Extra Info object from IValue.
    * The IValue store a complex Input
    *
-   * @param inputs
+   * @param input_signature
    */
   CompileSpec(torch::jit::IValue input_signature);
   // Defaults should reflect TensorRT defaults for BuilderConfig
 
-  /**
-   * @brief Specifications for inputs to the engine, can either be a single size or a range defined by min, opt and max
-   * sizes Users can also specify expected input type as well as tensor memory format
-   *
-   * Order in vector should match call order for the function
-   */
-  std::vector<Input> inputs;
-
   /**
    * @brief Specifications for inputs to the engine, can store a IValue which has stored complex Input
    *  or a flatened Input

cpp/src/compile_spec.cpp

@@ -18,22 +18,26 @@ torchtrt::core::runtime::CudaDevice to_internal_cuda_device(Device device);
 namespace torchscript {
 CompileSpec::CompileSpec(std::vector<c10::ArrayRef<int64_t>> fixed_sizes) {
   for (auto in : fixed_sizes) {
-    inputs.push_back(Input(in));
+    graph_inputs.inputs.push_back(Input(in));
   }
-  // graph_inputs.flattened_inputs = inputs;
 }
 
 CompileSpec::CompileSpec(std::vector<std::vector<int64_t>> fixed_sizes) {
   for (auto in : fixed_sizes) {
-    inputs.push_back(Input(in));
+    graph_inputs.inputs.push_back(Input(in));
   }
-  // graph_inputs.flattened_inputs = inputs;
+}
+
+CompileSpec::CompileSpec(std::vector<Input> inputs) {
+    graph_inputs.inputs = std::move(inputs);
 }
 
 CompileSpec::CompileSpec(torch::jit::IValue input_signature) {
     graph_inputs.input_signature = input_signature;
 }
 
+
+
 void flatten_dfs(std::vector<torchtrt::core::ir::Input>& flattened_inputs, std::vector<std::vector<torchtrt::core::ir::Input>>& collection_inputs, 
                  torch::jit::IValue input_ivalue, torch::jit::IValue& converted_ivalue, int level, int index) {
     if (input_ivalue.isTuple()) {
@@ -59,7 +63,6 @@ void flatten_dfs(std::vector<torchtrt::core::ir::Input>& flattened_inputs, std::
       }
       c10::TypePtr type = input_list[0].type();
       auto converted_elements = c10::impl::GenericList(type);
-      // std::vector<torch::jit::IValue> converted_elements;
       int idx = 0;
       for (auto item: input_list) {
         int cur_idx = level < 1 ? idx: index;
@@ -95,7 +98,7 @@ torch_tensorrt::core::ir::GraphInputs to_internal_graph_inputs(GraphInputs exter
 
     torch::jit::IValue converted_input_signature;
     flatten_dfs(flattened_inputs, collection_inputs, external_graph_input.input_signature, converted_input_signature, 0, 0);
-    internal_graph_input.flattened_inputs = flattened_inputs;
+    internal_graph_input.inputs = flattened_inputs;
     internal_graph_input.input_signature = converted_input_signature;
     internal_graph_input.collection_inputs = collection_inputs;
 
@@ -105,17 +108,15 @@ torch_tensorrt::core::ir::GraphInputs to_internal_graph_inputs(GraphInputs exter
 }
 
 torchtrt::core::CompileSpec to_internal_compile_spec(CompileSpec external) {
-  torchtrt::core::CompileSpec internal(to_vec_internal_inputs(external.inputs));
-  if (internal.inputs.size() == 0) {
+  torchtrt::core::CompileSpec internal(to_vec_internal_inputs(external.graph_inputs.inputs));
+  if (internal.graph_inputs.inputs.size() == 0) {
     LOG_DEBUG("GraphInput.inputs size == 0, using GraphInput.input_signature to get Input spec");
     internal.graph_inputs = to_internal_graph_inputs(external.graph_inputs);
-    internal.inputs = internal.graph_inputs.flattened_inputs;
   } else {
     LOG_DEBUG("GraphInput.inputs size != 0, using GraphInput.inputs to get Input spec");
-    internal.graph_inputs.collection_inputs.resize(internal.inputs.size());
-    for (int i = 0; i < internal.inputs.size(); i++) {
-      internal.graph_inputs.collection_inputs[i].push_back(internal.inputs[i]);
-      internal.graph_inputs.flattened_inputs = internal.inputs;
+    internal.graph_inputs.collection_inputs.resize(internal.graph_inputs.inputs.size());
+    for (int i = 0; i < internal.graph_inputs.inputs.size(); i++) {
+      internal.graph_inputs.collection_inputs[i].push_back(internal.graph_inputs.inputs[i]);
     }
   }