Merge branch 'inocsin-fix_fallback_inputs' into 'release/1.0'

feat: support setting input types of subgraph in fallback, handle Tensor type...

See merge request adlsa/TRTorch!10

Author: narendasan

core/conversion/conversion.cpp

@@ -8,7 +8,9 @@
 #include "core/util/prelude.h"
 
 #include "c10/util/intrusive_ptr.h"
+#include "core/conversion/converters/converter_util.h"
 #include "core/conversion/tensorcontainer/TensorContainer.h"
+#include "core/util/trt_util.h"
 
 namespace torch_tensorrt {
 namespace core {
@@ -212,6 +214,21 @@ void MarkOutputs(ConversionCtx* ctx, at::ArrayRef<const torch::jit::Value*> outp
           LOG_INFO(
               ctx->logger, "Marking Output " << out->debugName() << " named " << name << " in engine (ctx.MarkOutput)");
           ctx->num_outputs += 1;
+        } else if (out_ivalue.isTuple()) {
+          TORCHTRT_THROW_ERROR("Tuple type. Only a single tensor or a TensorList type is supported.");
+        } else if (out_ivalue.isList()) {
+          TORCHTRT_THROW_ERROR("List type. Only a single tensor or a TensorList type is supported.");
+        } else if (out_ivalue.isScalar()) {
+          TORCHTRT_THROW_ERROR("Scalar type. Only a single tensor or a TensorList type is supported.");
+        } else if (out_ivalue.isTensor()) {
+          // prim::NumToTensor will go to here
+          std::string name = std::string("output_") + std::to_string(ctx->num_outputs);
+          auto out_tensor = converters::tensor_to_const(ctx, out_ivalue.toTensor(), "");
+          out_tensor->setName(name.c_str());
+          ctx->net->markOutput(*out_tensor);
+          LOG_INFO(
+              ctx->logger, "Marking Output " << out->debugName() << " named " << name << " in engine (ctx.MarkOutput)");
+          ctx->num_outputs += 1;
         } else {
           TORCHTRT_THROW_ERROR("Unknown output type. Only a single tensor or a TensorList type is supported.");
         }
@@ -364,6 +381,7 @@ void ConvertBlockToNetDef(
     ConversionInfo& build_info,
     ir::StaticParams& static_params) {
   LOG_INFO(ctx->logger, "Converting Block");
+  LOG_DEBUG(ctx->logger, *b->owningGraph());
 
   auto inputs = b->inputs();
   AddParamsToCtxValueMap(ctx, static_params);

core/partitioning/PartitionInfo.h

@@ -12,6 +12,7 @@ struct PartitionInfo {
   bool enabled = false;
   uint64_t min_block_size = 1;
   std::vector<std::string> forced_fallback_operators;
+  bool truncate_long_and_double;
 };
 
 std::ostream& operator<<(std::ostream& os, const PartitionInfo& s);

core/partitioning/partitioning.cpp

@@ -404,7 +404,7 @@ PartitionedGraph Partition(
   registerSegmentsOutputs(segmented_blocks, block);
 
   // run shape analysis on each segmented block
-  runShapeAnalysis(segmented_blocks, example_tensor_map);
+  runShapeAnalysis(segmented_blocks, example_tensor_map, partition_info);
 
   LOG_INFO(segmented_blocks);
 

core/partitioning/shape_analysis.cpp

@@ -34,7 +34,8 @@ std::unordered_map<const torch::jit::Value*, torch::jit::IValue> generateRandomI
 
 void getSegmentsOutputByRunning(
     SegmentedBlock& seg_block,
-    std::unordered_map<const torch::jit::Value*, torch::jit::IValue>& ivalues_maps) {
+    std::unordered_map<const torch::jit::Value*, torch::jit::IValue>& ivalues_maps,
+    const PartitionInfo& partition_info) {
   // create a module to run the graph
   auto g = seg_block.g();
   auto copy_g = g->copy();
@@ -108,7 +109,28 @@ void getSegmentsOutputByRunning(
   std::vector<at::ScalarType> input_types;
   for (auto& i : seg_block.raw_inputs()) {
     if (ivalues_maps[i].isTensor()) {
-      input_shapes.push_back(util::toVec(util::toDims(ivalues_maps[i].toTensor().sizes())));
+      // set the input_shape and data_type
+      at::ScalarType t = ivalues_maps[i].toTensor().scalar_type();
+      if (!partition_info.truncate_long_and_double && (t == at::kLong || t == at::kDouble)) {
+        TORCHTRT_THROW_ERROR(
+            "Unable to process subgraph input type of at::kLong/at::kDouble, try to compile model with truncate_long_and_double enabled");
+      } else if (partition_info.truncate_long_and_double && t == at::kLong) {
+        ivalues_maps[i] = ivalues_maps[i].toTensor().to(at::kInt);
+        LOG_WARNING("Truncating graph input type from at::kLong to at::kInt");
+      } else if (partition_info.truncate_long_and_double && t == at::kDouble) {
+        ivalues_maps[i] = ivalues_maps[i].toTensor().to(at::kFloat);
+        LOG_WARNING("Truncating graph input type from at::kDouble to at::kFloat");
+      }
+      c10::optional<nvinfer1::DataType> dtype = util::optTypeMetaToTRTDataType(ivalues_maps[i].toTensor().dtype());
+      if (dtype == c10::nullopt) {
+        TORCHTRT_THROW_ERROR("Unsupported input data type " << ivalues_maps[i].toTensor().dtype());
+      }
+      if (ivalues_maps[i].toTensor().sizes().size() == 0) {
+        // handle Scalar types, which has sizes of []
+        input_shapes.push_back(util::toVec(util::toDims(c10::List<long int>({1}))));
+      } else {
+        input_shapes.push_back(util::toVec(util::toDims(ivalues_maps[i].toTensor().sizes())));
+      }
       input_types.push_back(ivalues_maps[i].toTensor().scalar_type());
     }
   }
@@ -119,11 +141,12 @@ void getSegmentsOutputByRunning(
 
 void runShapeAnalysis(
     std::vector<SegmentedBlock>& segmented_blocks,
-    std::unordered_map<const torch::jit::Value*, torch::jit::IValue>& example_tensor_map) {
+    std::unordered_map<const torch::jit::Value*, torch::jit::IValue>& example_tensor_map,
+    const PartitionInfo& partition_info) {
   // register every segment's input shape, and it's running output IValues
   for (auto& seg_block : segmented_blocks) {
     torch::jit::ConstantPooling(seg_block.g());
-    getSegmentsOutputByRunning(seg_block, example_tensor_map);
+    getSegmentsOutputByRunning(seg_block, example_tensor_map, partition_info);
   }
   return;
 }

core/partitioning/shape_analysis.h

@@ -12,7 +12,8 @@ std::unordered_map<const torch::jit::Value*, torch::jit::IValue> generateRandomI
 
 void runShapeAnalysis(
     std::vector<SegmentedBlock>& segmented_blocks,
-    std::unordered_map<const torch::jit::Value*, torch::jit::IValue>& ivalues_maps);
+    std::unordered_map<const torch::jit::Value*, torch::jit::IValue>& ivalues_maps,
+    const PartitionInfo& partition_info);
 
 } // namespace partitioning
 } // namespace core

core/util/trt_util.cpp

@@ -239,8 +239,7 @@ const std::unordered_map<at::ScalarType, nvinfer1::DataType>& get_at_trt_type_ma
       {at::kHalf, nvinfer1::DataType::kHALF},
       {at::kInt, nvinfer1::DataType::kINT32},
       {at::kChar, nvinfer1::DataType::kINT8},
-      {at::kBool, nvinfer1::DataType::kBOOL},
-  };
+      {at::kBool, nvinfer1::DataType::kBOOL}};
   return at_trt_type_map;
 }
 

cpp/src/compile_spec.cpp

@@ -57,6 +57,7 @@ torchtrt::core::CompileSpec to_internal_compile_spec(CompileSpec external) {
   internal.partition_info.enabled = !external.require_full_compilation;
   internal.partition_info.min_block_size = external.min_block_size;
   internal.partition_info.forced_fallback_operators = std::move(external.torch_executed_ops);
+  internal.partition_info.truncate_long_and_double = external.truncate_long_and_double;
   internal.lower_info.forced_fallback_modules = std::move(external.torch_executed_modules);
 
   switch (external.device.device_type) {