chore: [collection] rename ConversionInfo.collection_inputs to ConversionInfo.collection_input_spec_map

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

Author: inocsin

core/compiler.cpp

@@ -294,13 +294,13 @@ void MapInputsAndDetermineDTypes(
     std::shared_ptr<torch::jit::Graph>& g,
     ir::StaticParams& static_params,
     ir::CollectionTypeMap& first_use_type_map) {
-    cfg.convert_info.collection_inputs = std::move(ir::associate_specs_with_collection_inputs(g, cfg.graph_inputs, static_params));
+    cfg.convert_info.collection_input_spec_map = std::move(ir::associate_specs_with_collection_inputs(g, cfg.graph_inputs, static_params));
 
     auto collection_inputs = ir::get_collection_inputs(g, static_params);
     LOG_DEBUG("In MapInputsAndDetermineDTypes, the g->inputs() size is " << g->inputs().size() << ", CollectionInputSpecMap size is" << collection_inputs.size());
 
     for (auto in : collection_inputs) {
-      std::vector<ir::Input>& spec = cfg.convert_info.collection_inputs.find(in)->second;
+      std::vector<ir::Input>& spec = cfg.convert_info.collection_input_spec_map.find(in)->second;
       std::vector<c10::optional<at::ScalarType>> est_type_opt;
 
       auto est_it = first_use_type_map.find(in);
@@ -327,21 +327,21 @@ void MapInputsAndDetermineDTypes(
             LOG_INFO("Cannot infer input tensor dtype in graph, compiler is going to use the user setting");
             std::stringstream ss;
             ss << "For input " << in->debugName() << ", found user specified input dtype as ";
-            ss << cfg.convert_info.collection_inputs.find(in)->second[i].dtype;
-            ss << ". The compiler is going to use the user setting " << cfg.convert_info.collection_inputs.find(in)->second[i].dtype;
+            ss << cfg.convert_info.collection_input_spec_map.find(in)->second[i].dtype;
+            ss << ". The compiler is going to use the user setting " << cfg.convert_info.collection_input_spec_map.find(in)->second[i].dtype;
             auto warn_str = ss.str();
             LOG_WARNING(warn_str);
             // Overwrite type map with user settings
-            first_use_type_map[in][i] = {util::TRTDataTypeToScalarType(cfg.convert_info.collection_inputs.find(in)->second[i].dtype)};
+            first_use_type_map[in][i] = {util::TRTDataTypeToScalarType(cfg.convert_info.collection_input_spec_map.find(in)->second[i].dtype)};
 
           } else {
-            if (util::TRTDataTypeToScalarType(cfg.convert_info.collection_inputs.find(in)->second[i].dtype) != est_type_opt[i].value()) {
+            if (util::TRTDataTypeToScalarType(cfg.convert_info.collection_input_spec_map.find(in)->second[i].dtype) != est_type_opt[i].value()) {
               std::stringstream ss;
               ss << "For input " << in->debugName() << ", found user specified input dtype as ";
-              ss << cfg.convert_info.collection_inputs.find(in)->second[i].dtype;
+              ss << cfg.convert_info.collection_input_spec_map.find(in)->second[i].dtype;
               ss << ", however when inspecting the graph, the input type expected was inferred to be ";
               ss << est_type_opt[i].value() << std::endl;
-              ss << "The compiler is going to use the user setting " << cfg.convert_info.collection_inputs.find(in)->second[i].dtype;
+              ss << "The compiler is going to use the user setting " << cfg.convert_info.collection_input_spec_map.find(in)->second[i].dtype;
               ss << "\nThis conflict may cause an error at runtime due to partial compilation being enabled and therefore\n";
               ss << "compatibility with PyTorch's data type convention is required.\n";
               ss << "If you do indeed see errors at runtime either:\n";
@@ -350,7 +350,7 @@ void MapInputsAndDetermineDTypes(
               auto warn_str = ss.str();
               LOG_WARNING(warn_str);
               // Overwrite type map with user settings
-              first_use_type_map[in][i] = {util::TRTDataTypeToScalarType(cfg.convert_info.collection_inputs.find(in)->second[i].dtype)};
+              first_use_type_map[in][i] = {util::TRTDataTypeToScalarType(cfg.convert_info.collection_input_spec_map.find(in)->second[i].dtype)};
             }
           }
         } else {
@@ -436,7 +436,7 @@ torch::jit::Module CompileGraph(const torch::jit::Module& mod, CompileSpec cfg)
           !(cfg.lower_info.forced_fallback_modules.size() == 0 &&
             cfg.partition_info.forced_fallback_operators.size() == 0 &&
             conversion::VerifyConverterSupportForBlock(g->block(), false))) {
-        auto collection_input_ivalues_map = partitioning::generateRandomInputs(cfg.convert_info.collection_inputs, first_use_types);
+        auto collection_input_ivalues_map = partitioning::generateRandomInputs(cfg.convert_info.collection_input_spec_map, first_use_types);
         auto graph_and_mapping = ConstructFallbackGraph(new_mod, g->block(), collection_input_ivalues_map, cfg, static_params);
         new_g = graph_and_mapping.first;
         LOG_INFO("Segmented Graph: " << *new_g);

core/compiler.h

@@ -14,12 +14,8 @@ namespace torch_tensorrt {
 namespace core {
 
 struct CompileSpec {
-  CompileSpec(std::vector<ir::Input> inputs) : inputs(inputs) {
-    // graph_inputs = ir::GraphInputs(inputs);
-  }
+  CompileSpec(std::vector<ir::Input> inputs) : inputs(inputs) {}
   CompileSpec(torch::jit::IValue& input_signature) {
-    // graph_inputs = ir::GraphInputs(input_signature);
-    // inputs = graph_inputs.flattened_inputs;
     graph_inputs.input_signature = input_signature;
   }
   ir::GraphInputs graph_inputs;

core/conversion/conversion.cpp

@@ -134,9 +134,8 @@ void AddInputs(
     ConversionCtx* ctx,
     c10::ArrayRef<const torch::jit::Value*> inputs,
     ConversionInfo& conversion_info) {
-    // std::unordered_map<const torch::jit::Value*, ir::Input>& input_specs) {
   std::unordered_map<const torch::jit::Value*, ir::Input>& input_specs = conversion_info.inputs;
-  std::unordered_map<const torch::jit::Value*, std::vector<ir::Input>> collection_input_spec = conversion_info.collection_inputs;
+  std::unordered_map<const torch::jit::Value*, std::vector<ir::Input>> collection_input_spec = conversion_info.collection_input_spec_map;
 
   std::vector<const torch::jit::Value*> input_tensors;
   for (auto in : inputs) {
@@ -396,7 +395,6 @@ void ConvertBlockToNetDef(
 
   auto inputs = b->inputs();
   AddParamsToCtxValueMap(ctx, static_params);
-  // AddInputs(ctx, inputs, build_info.inputs);
   AddInputs(ctx, inputs, build_info);
 
   auto nodes = b->nodes();

core/conversion/conversion.h

@@ -13,7 +13,7 @@ namespace conversion {
 
 struct ConversionInfo {
   ir::InputSpecMap inputs;
-  ir::CollectionInputSpecMap collection_inputs;
+  ir::CollectionInputSpecMap collection_input_spec_map;
   BuilderSettings engine_settings;
 };