pytorch · narendasan · Sep 9, 2022 · May 31, 2022 · Jul 26, 2022 · Jul 27, 2022
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -435,6 +435,7 @@ commands:
             mkdir -p /tmp/artifacts/test_results
             cd tests/py
             pytest --junitxml=/tmp/artifacts/test_results/api/api_test_results.xml api/
+            pytest --junitxml=/tmp/artifacts/test_results/models/models_test_results.xml models/
             pytest --junitxml=/tmp/artifacts/test_results/integrations/integrations_test_results.xml integrations/
             cd ~/project
 

diff --git a/.github/workflows/docgen.yml b/.github/workflows/docgen.yml
@@ -31,7 +31,7 @@ jobs:
       - name: Set up Python 3.9.4
         uses: actions/setup-python@v2
         with:
-          python-version: 3.9.4 
+          python-version: 3.9.4
       - uses: actions/checkout@v2
         with:
           ref: ${{github.head_ref}}

diff --git a/.github/workflows/linter.yml b/.github/workflows/linter.yml
@@ -39,7 +39,7 @@ jobs:
           pip3 install -r $GITHUB_WORKSPACE/.github/scripts/requirements.txt
           pip3 install -r $GITHUB_WORKSPACE/requirements-dev.txt
       - name: Lint C++
-        run: | 
+        run: |
           cd $GITHUB_WORKSPACE
           python3 $GITHUB_WORKSPACE/.github/scripts/run_cpp_linter.py
         env:

diff --git a/core/conversion/converters/impl/max.cpp b/core/conversion/converters/impl/max.cpp
@@ -13,47 +13,95 @@ namespace conversion {
 namespace converters {
 namespace impl {
 namespace {
-auto max_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(
-    {"aten::max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)",
-     [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
-       auto self = args[0].ITensorOrFreeze(ctx);
-       auto dim = args[1].unwrapToInt();
-       auto keep_dims = args[2].unwrapToBool();
-       auto selfDim = util::toVec(self->getDimensions());
-       if (dim < 0) {
-         dim = selfDim.size() + dim;
-       }
-       uint32_t shiftDim = 1 << dim;
-       auto TopKOperation = nvinfer1::TopKOperation::kMAX;
-       auto topk_layer = ctx->net->addTopK(*self, TopKOperation, 1, shiftDim);
-       TORCHTRT_CHECK(topk_layer, "Unable to create max layer from node: " << *n);
-       auto topk_dims = util::toVec(topk_layer->getOutput(0)->getDimensions());
-
-       nvinfer1::ITensor* out0 = nullptr;
-       nvinfer1::ITensor* out1 = nullptr;
-       if (!keep_dims) {
-         if (topk_dims[dim] == 1) {
-           auto squeeze_layer = ctx->net->addShuffle(*topk_layer->getOutput(0));
-           squeeze_layer->setReshapeDimensions(util::squeezeDims(topk_layer->getOutput(0)->getDimensions(), dim));
-           TORCHTRT_CHECK(squeeze_layer, "Unable to create squeeze_layer layer from node: " << *n);
-           out0 = ctx->AssociateValueAndTensor(n->outputs()[0], squeeze_layer->getOutput(0));
-
-           auto squeeze_layer_indices = ctx->net->addShuffle(*topk_layer->getOutput(1));
-           squeeze_layer_indices->setReshapeDimensions(
-               util::squeezeDims(topk_layer->getOutput(1)->getDimensions(), dim));
-           TORCHTRT_CHECK(squeeze_layer_indices, "Unable to create squeeze_layer_indices layer from node: " << *n);
-           out1 = ctx->AssociateValueAndTensor(n->outputs()[1], squeeze_layer_indices->getOutput(0));
-         }
-       } else {
-         out0 = ctx->AssociateValueAndTensor(n->outputs()[0], topk_layer->getOutput(0));
-         out1 = ctx->AssociateValueAndTensor(n->outputs()[1], topk_layer->getOutput(1));
-       }
-
-       LOG_DEBUG("Output tensor(0) shape: " << out0->getDimensions());
-       LOG_DEBUG("Output tensor(1) shape: " << out1->getDimensions());
-
-       return true;
-     }});
+
+bool min_max_dim(ConversionCtx* ctx, const torch::jit::Node* n, args& args, nvinfer1::TopKOperation topKOperation) {
+  auto self = args[0].ITensorOrFreeze(ctx);
+  auto dim = args[1].unwrapToInt();
+  auto keep_dims = args[2].unwrapToBool();
+  auto selfDim = util::toVec(self->getDimensions());
+  if (dim < 0) {
+    dim = selfDim.size() + dim;
+  }
+  uint32_t reduce_axes_mask = 1 << dim;
+  auto topk_layer = ctx->net->addTopK(*self, topKOperation, 1, reduce_axes_mask);
+  TORCHTRT_CHECK(topk_layer, "Unable to create topk layer from node: " << *n);
+  auto topk_dims = util::toVec(topk_layer->getOutput(0)->getDimensions());
+
+  nvinfer1::ITensor* out0 = nullptr;
+  nvinfer1::ITensor* out1 = nullptr;
+  if (!keep_dims) {
+    TORCHTRT_CHECK(topk_dims[dim] == 1, "Unexpected size in squeeze dimension. Expected: 1 Actual: " << topk_dims[dim]);
+    auto squeeze_layer = ctx->net->addShuffle(*topk_layer->getOutput(0));
+    squeeze_layer->setReshapeDimensions(util::squeezeDims(topk_layer->getOutput(0)->getDimensions(), dim));
+    TORCHTRT_CHECK(squeeze_layer, "Unable to create squeeze_layer layer from node: " << *n);
+    out0 = ctx->AssociateValueAndTensor(n->outputs()[0], squeeze_layer->getOutput(0));
+
+    auto squeeze_layer_indices = ctx->net->addShuffle(*topk_layer->getOutput(1));
+    squeeze_layer_indices->setReshapeDimensions(util::squeezeDims(topk_layer->getOutput(1)->getDimensions(), dim));
+    TORCHTRT_CHECK(squeeze_layer_indices, "Unable to create squeeze_layer_indices layer from node: " << *n);
+    out1 = ctx->AssociateValueAndTensor(n->outputs()[1], squeeze_layer_indices->getOutput(0));
+  } else {
+    out0 = ctx->AssociateValueAndTensor(n->outputs()[0], topk_layer->getOutput(0));
+    out1 = ctx->AssociateValueAndTensor(n->outputs()[1], topk_layer->getOutput(1));
+  }
+
+  LOG_DEBUG("Output tensor(0) shape: " << out0->getDimensions());
+  LOG_DEBUG("Output tensor(1) shape: " << out1->getDimensions());
+
+  return true;
+}
+
+bool arg_min_max(ConversionCtx* ctx, const torch::jit::Node* n, args& args, nvinfer1::TopKOperation topKOperation) {
+  auto self = args[0].ITensorOrFreeze(ctx);
+  auto dim = args[1].unwrapToInt();
+  auto keep_dims = args[2].unwrapToBool();
+  auto selfDim = util::toVec(self->getDimensions());
+  if (dim < 0) {
+    dim = selfDim.size() + dim;
+  }
+  uint32_t reduce_axes_mask = 1 << dim;
+  auto topk_layer = ctx->net->addTopK(*self, topKOperation, 1, reduce_axes_mask);
+  TORCHTRT_CHECK(topk_layer, "Unable to create topk layer from node: " << *n);
+  auto topk_dims = util::toVec(topk_layer->getOutput(0)->getDimensions());
+
+  nvinfer1::ITensor* out = nullptr;
+  if (!keep_dims) {
+    TORCHTRT_CHECK(topk_dims[dim] == 1, "Unexpected size in squeeze dimension. Expected: 1 Actual: " << topk_dims[dim]);
+    auto squeeze_layer_indices = ctx->net->addShuffle(*topk_layer->getOutput(1));
+    squeeze_layer_indices->setReshapeDimensions(util::squeezeDims(topk_layer->getOutput(1)->getDimensions(), dim));
+    TORCHTRT_CHECK(squeeze_layer_indices, "Unable to create squeeze_layer_indices layer from node: " << *n);
+    out = ctx->AssociateValueAndTensor(n->outputs()[0], squeeze_layer_indices->getOutput(0));
+  } else {
+    out = ctx->AssociateValueAndTensor(n->outputs()[0], topk_layer->getOutput(1));
+  }
+
+  LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+
+  return true;
+}
+
+auto max_registrations TORCHTRT_UNUSED =
+    RegisterNodeConversionPatterns()
+        .pattern(
+            {"aten::max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               return min_max_dim(ctx, n, args, nvinfer1::TopKOperation::kMAX);
+             }})
+        .pattern(
+            {"aten::min.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               return min_max_dim(ctx, n, args, nvinfer1::TopKOperation::kMIN);
+             }})
+        .pattern(
+            {"aten::argmax(Tensor self, int dim, bool keepdim=False) -> (Tensor)",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               return arg_min_max(ctx, n, args, nvinfer1::TopKOperation::kMAX);
+             }})
+        .pattern(
+            {"aten::argmin(Tensor self, int dim, bool keepdim=False) -> (Tensor)",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               return arg_min_max(ctx, n, args, nvinfer1::TopKOperation::kMIN);
+             }});
 } // namespace
 } // namespace impl
 } // namespace converters

diff --git a/core/conversion/converters/impl/unary.cpp b/core/conversion/converters/impl/unary.cpp
@@ -49,6 +49,21 @@ auto abs_registration TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern
        }
      }});
 
+auto reciprocal_registration TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(
+    {"aten::reciprocal(Tensor self) -> Tensor", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+       auto in = args[0].ITensorOrFreeze(ctx);
+       if (in->getType() == nvinfer1::DataType::kINT32) {
+         // pytorch implicitly casts to float for aten::reciprocal(int)
+         in = castITensor(ctx, in, nvinfer1::DataType::kFLOAT);
+       }
+       auto unary_layer = ctx->net->addUnary(*in, nvinfer1::UnaryOperation::kRECIP);
+       TORCHTRT_CHECK(unary_layer, "Unable to create recip layer from node: " << *n);
+       unary_layer->setName(util::node_info(n).c_str());
+       auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], unary_layer->getOutput(0));
+       LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+       return true;
+     }});
+
 #define convert(unary, trt_type)                                                               \
   auto unary##_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(       \
       {"aten::" #unary "(Tensor self) -> Tensor",                                              \
@@ -74,7 +89,6 @@ convert(sinh, kSINH);
 convert(tan, kTAN);
 convert(atan, kATAN);
 convert(floor, kFLOOR);
-convert(reciprocal, kRECIP);
 convert(log, kLOG);
 convert(ceil, kCEIL);
 convert(sqrt, kSQRT);

diff --git a/core/partitioning/shape_analysis.cpp b/core/partitioning/shape_analysis.cpp
@@ -167,7 +167,7 @@ void getSegmentsOutputByRunning(
       }
       if (cur_ivalue.toTensor().sizes().size() == 0) {
         // handle Scalar types, which has sizes of []
-        input_shapes.push_back(util::toVec(util::toDims(c10::List<long int>({1}))));
+        input_shapes.push_back(util::toVec(util::toDims(c10::List<int64_t>({1}))));
       } else {
         input_shapes.push_back(util::toVec(util::toDims(cur_ivalue.toTensor().sizes())));
       }

diff --git a/cpp/bin/torchtrtc/main.cpp b/cpp/bin/torchtrtc/main.cpp
@@ -35,7 +35,7 @@ bool unload_library(void* custom_lib) {
   bool success = false;
 #if defined(_WIN32)
   // Returns status non-zero for success
-  success = FreeLibrary(custom_lib) ? true : false;
+  success = FreeLibrary((HMODULE)custom_lib) ? true : false;
 #else
   success = dlclose(custom_lib) ? false : true;
 #endif

diff --git a/cpp/include/torch_tensorrt/torch_tensorrt.h b/cpp/include/torch_tensorrt/torch_tensorrt.h
@@ -365,7 +365,7 @@ class TensorFormat {
  * signifying a static input shape or a set of three input shapes representing
  * the min, optiminal and max input shapes allowed for the engine.
  */
-struct TORCHTRT_API Input : torch::CustomClassHolder {
+struct Input : torch::CustomClassHolder {
   /// Minimum acceptable input size into the engine
   std::vector<int64_t> min_shape;
   /// Optimal input size into the engine (size optimized for given kernels accept any size in min max range)
@@ -520,7 +520,7 @@ struct TORCHTRT_API Input : torch::CustomClassHolder {
  *
  * This struct can either hold a complex inputs of shape or a flattened one,
  */
-struct TORCHTRT_API GraphInputs {
+struct GraphInputs {
   torch::jit::IValue input_signature; // nested Input, full input spec
   std::vector<Input> inputs; // flatten input spec
 };
@@ -592,14 +592,14 @@ struct CompileSpec {
    *
    * @param inputs
    */
-  CompileSpec(std::vector<Input> inputs);
+  TORCHTRT_API CompileSpec(std::vector<Input> inputs);
 
   /**
    * @brief Construct a new Compile Spec  object from IValue which represents the nesting of input tensors for a module.
    *
    * @param input_signature
    */
-  CompileSpec(torch::jit::IValue input_signature);
+  TORCHTRT_API CompileSpec(torch::jit::IValue input_signature);
   // Defaults should reflect TensorRT defaults for BuilderConfig
 
   /**

diff --git a/docs/_cpp_api/classtorch__tensorrt_1_1DataType.html b/docs/_cpp_api/classtorch__tensorrt_1_1DataType.html
@@ -199,7 +199,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/classtorch__tensorrt_1_1Device_1_1DeviceType.html b/docs/_cpp_api/classtorch__tensorrt_1_1Device_1_1DeviceType.html
@@ -199,7 +199,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/classtorch__tensorrt_1_1TensorFormat.html b/docs/_cpp_api/classtorch__tensorrt_1_1TensorFormat.html
@@ -199,7 +199,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/classtorch__tensorrt_1_1ptq_1_1Int8CacheCalibrator.html b/docs/_cpp_api/classtorch__tensorrt_1_1ptq_1_1Int8CacheCalibrator.html
@@ -199,7 +199,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/classtorch__tensorrt_1_1ptq_1_1Int8Calibrator.html b/docs/_cpp_api/classtorch__tensorrt_1_1ptq_1_1Int8Calibrator.html
@@ -199,7 +199,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/define_macros_8h_1a18d295a837ac71add5578860b55e5502.html b/docs/_cpp_api/define_macros_8h_1a18d295a837ac71add5578860b55e5502.html
@@ -199,7 +199,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/define_macros_8h_1a282fd3c0b1c3a215148ae372070e1268.html b/docs/_cpp_api/define_macros_8h_1a282fd3c0b1c3a215148ae372070e1268.html
@@ -199,7 +199,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/define_macros_8h_1a31398a6d4d27e28817afb0f0139e909e.html b/docs/_cpp_api/define_macros_8h_1a31398a6d4d27e28817afb0f0139e909e.html
@@ -199,7 +199,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/define_macros_8h_1a35703561b26b1a9d2738ad7d58b27827.html b/docs/_cpp_api/define_macros_8h_1a35703561b26b1a9d2738ad7d58b27827.html
@@ -199,7 +199,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/define_macros_8h_1abd1465eb38256d3f22cc1426b23d516b.html b/docs/_cpp_api/define_macros_8h_1abd1465eb38256d3f22cc1426b23d516b.html
@@ -199,7 +199,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/define_macros_8h_1abe87b341f562fd1cf40b7672e4d759da.html b/docs/_cpp_api/define_macros_8h_1abe87b341f562fd1cf40b7672e4d759da.html
@@ -199,7 +199,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/define_macros_8h_1ad19939408f7be171a74a89928b36eb59.html b/docs/_cpp_api/define_macros_8h_1ad19939408f7be171a74a89928b36eb59.html
@@ -199,7 +199,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/define_macros_8h_1adad592a7b1b7eed529cdf6acd584c883.html b/docs/_cpp_api/define_macros_8h_1adad592a7b1b7eed529cdf6acd584c883.html
@@ -199,7 +199,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/dir_cpp.html b/docs/_cpp_api/dir_cpp.html
@@ -197,7 +197,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/dir_cpp_include.html b/docs/_cpp_api/dir_cpp_include.html
@@ -197,7 +197,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/dir_cpp_include_torch_tensorrt.html b/docs/_cpp_api/dir_cpp_include_torch_tensorrt.html
@@ -197,7 +197,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
 
 

diff --git a/docs/_cpp_api/enum_logging_8h_1a130f65408ad8cbaee060f05e8db69558.html b/docs/_cpp_api/enum_logging_8h_1a130f65408ad8cbaee060f05e8db69558.html
@@ -199,7 +199,7 @@
 
 
                 <div class="version">
-                  master (1.2.0a0+51a991e)
+                  master (1.2.0a0+096fd41)
                 </div>
-Original file line number
+Diff line change
@@ Expand Up / @@ -199,7 +199,7 @@ @@
                     <div class="version">
-                      master (1.2.0a0+51a991e)
+                      master (1.2.0a0+096fd41)
                     </div>
@@ Expand Down @@