Merge branch 'release/1.0' into rebased_1.10_update

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);
@@ -508,24 +526,37 @@ bool VerifyConverterSupportForBlock(const torch::jit::Block* b, bool suppress_er
     unsupported_msg << "https://www.github.com/nvidia/Torch-TensorRT/issues" << std::endl;
     unsupported_msg << std::endl << "In Module:" << std::endl;
 
-    if (suppress_errors) {
+    if (!suppress_errors) {
       LOG_ERROR(unsupported_msg.str());
     }
 
+    std::unordered_map<std::string, std::unordered_set<std::string>> unsupported_node_locations;
     for (const auto n : b->nodes()) {
       auto schema = n->maybeSchema();
       if (schema) {
         for (const auto& x : unsupported_ops) {
           if (x.first == schema->operator_name()) {
-            if (suppress_errors) {
-              LOG_ERROR(
-                  "Unsupported operator: " << *schema << std::endl
-                                           << torch_tensorrt::core::util::GetPyTorchSourceCode(n) << std::endl);
+            auto loc = unsupported_node_locations.find(x.second);
+            if (loc == unsupported_node_locations.end()) {
+              unsupported_node_locations.insert({x.second, {torch_tensorrt::core::util::GetPyTorchSourceCode(n)}});
+            } else  {
+              loc->second.insert(torch_tensorrt::core::util::GetPyTorchSourceCode(n));
             }
           }
         }
       }
     }
+
+    for (const auto& type : unsupported_node_locations) {
+      std::stringstream traceback;
+      traceback << "Unsupported operator: " << type.first << std::endl;
+      for (const auto& str : type.second) {
+        traceback << str;
+      }
+      auto tb_str = traceback.str();
+      LOG_ERROR(tb_str);
+    }
+
     return false;
   }
 
@@ -537,7 +568,7 @@ bool VerifyConverterSupportForBlock(const torch::jit::Block* b, bool suppress_er
     unsupported_msg
         << "This may be because there are no operators that can be added to the TensorRT graph or all operators have a resolved compile time value."
         << std::endl;
-    if (suppress_errors) {
+    if (!suppress_errors) {
       LOG_ERROR(unsupported_msg.str());
     }
     return false;

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) {

docker/README.md

@@ -1,9 +1,16 @@
+# Building a Torch-TensorRT container
 
-## Use the container for pytorch1.10+cuda11.1+trt8.0.3.4
+### Install Docker and NVIDIA Container Toolkit
+
+https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html
+
+### Build Container
+
+> From root of Torch-TensorRT repo
 
 ```
 # Build:
-docker build -f docker/Dockerfile -t torch_tensorrt1.0:latest .
+docker build --build-arg BASE={PyTorch Base Container Version} -f docker/Dockerfile -t torch_tensorrt1.0:latest .
 
 # Run:
 docker run --gpus all -it \

notebooks/README.md

@@ -14,18 +14,20 @@ git clone https://github.com/NVIDIA/Torch-TensorRT
 Next, build the NVIDIA Torch-TensorRT container (from repo root):
 
 ```
+cd Torch-TensorRT
 docker build -t torch_tensorrt -f ./docker/Dockerfile .
 ```
 
 Then launch the container with:
 
 ```
-docker run --runtime=nvidia -it --rm --ipc=host --net=host torch_tensorrt
+docker run --gpus=all --rm -it -v $PWD:/Torch-TensorRT --net=host torch_tensorrt bash
 ```
 
 Within the docker interactive bash session, start Jupyter with
 
 ```
+cd /Torch-TensorRT/notebooks
 jupyter notebook --allow-root --ip 0.0.0.0 --port 8888
 ```
 
@@ -38,8 +40,7 @@ in, for example:
 ```http://[host machine]:8888/?token=aae96ae9387cd28151868fee318c3b3581a2d794f3b25c6b```
 
 
-Within the container, the notebooks themselves are located at `/workspace/torch_tensorrt/notebooks`. To reach them in Jupyter, click on the folder marked 
-`torch_tensorrt`, then the folder marked `notebooks`.
+Within the container, the notebooks themselves are located at `/Torch-TensorRT/notebooks`.
 
 ## 2. Notebook list