Merge branch 'master' into param_input

Author: bowang007

.github/code-owners.yml

@@ -9,7 +9,6 @@
 
 "component: build system":
   - "narendasan"
-  - "andi4191"
 
 "component: conversion":
   - "narendasan"
@@ -29,7 +28,6 @@
   - "peri044"
 
 "component: execution":
-  - "andi4191"
   - "narendasan"
 
 "component: lowering":
@@ -48,15 +46,12 @@
   - "peri044"
 
 "component: runtime":
-  - "andi4191"
   - "narendasan"
 
 "component: tests":
-  - "andi4191"
   - "narendasan"
 
 "component: torchtrtc":
-  - "andi4191"
   - "narendasan"
 
 "component: dependencies":
@@ -74,24 +69,20 @@
   - "tanayvarshney"
 
 "infrastructre":
-  - "andi4191"
   - "narendasan"
 
 "component: packaging":
   - "narendasan"
-  - "andi4191"
   - "peri044"
 
 "channel: NGC":
-  - "andi4191"
   - "peri044"
 
 "channel: linux-x86":
   - "narendasan"
   - "peri044"
 
 "channel: linux-sbsa":
-  - "andi4191"
   - "bowang007"
 
 "channel: windows":
@@ -102,16 +93,13 @@
   - "bowang007"
 
 "component: tooling":
-  - "andi4191"
   - "narendasan"
 
 "performance":
-  - "andi4191"
   - "peri044"
   - "bowang007"
 
 "channel: docker":
-  - "andi4191"
   - "narendasan"
 
 "ux":

README.md

@@ -122,6 +122,10 @@ These are the following dependencies used to verify the testcases. Torch-TensorR
 
 Releases: https://github.com/pytorch/TensorRT/releases
 
+```
+pip install torch-tensorrt==1.2.0 --find-links https://github.com/pytorch/TensorRT/releases/expanded_assets/v1.2.0
+```
+
 ## Compiling Torch-TensorRT
 
 ### Installing Dependencies

core/conversion/converters/impl/einsum.cpp

@@ -18,6 +18,13 @@ auto einsum_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pat
        auto equation = args[0].unwrapToString();
        auto in = args[1].IValue()->toListRef();
 
+       TORCHTRT_CHECK(
+           in.size() <= 2,
+           "TensorRT currently supports up to 2 input tensors "
+               << "to einsum but operation had " << in.size()
+               << " input tensors, please specify torch_executed_ops=[\"aten::einsum\"] "
+               << "at compilation time to avoid this error.");
+
        std::vector<nvinfer1::ITensor*> tensors;
 
        // Populate vector of ITensor pointers

core/conversion/converters/impl/element_wise.cpp

@@ -166,11 +166,11 @@ auto element_wise_registrations TORCHTRT_UNUSED =
              [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                // Should implement self - alpha * other
                auto self = args[0].ITensorOrFreeze(ctx);
-               auto scalar = args[2].unwrapToScalar().to<float>();
                auto other = args[1].ITensorOrFreeze(ctx);
+               auto scalar = args[2].unwrapToScalar();
 
-               if (1 != scalar) {
-                 auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+               if (1 != scalar.to<float>()) {
+                 auto alphaTensor = scalar_to_tensor(ctx, scalar);
                  auto scaleLayer = add_elementwise(
                      ctx,
                      nvinfer1::ElementWiseOperation::kPROD,
@@ -214,11 +214,11 @@ auto element_wise_registrations TORCHTRT_UNUSED =
              [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                // Should implement self - alpha * other
                auto self = args[0].ITensorOrFreeze(ctx);
-               auto scalar = args[2].unwrapToScalar().to<float>();
                auto other = args[1].ITensorOrFreeze(ctx);
+               auto scalar = args[2].unwrapToScalar();
 
-               if (1 != scalar) {
-                 auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+               if (1 != scalar.to<float>()) {
+                 auto alphaTensor = scalar_to_tensor(ctx, scalar);
                  auto scaleLayer = add_elementwise(
                      ctx,
                      nvinfer1::ElementWiseOperation::kPROD,
@@ -351,8 +351,7 @@ auto element_wise_registrations TORCHTRT_UNUSED =
             {"aten::div.Scalar(Tensor self, Scalar other) -> (Tensor)",
              [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                auto self = args[0].ITensorOrFreeze(ctx);
-               auto otherScalar = args[1].unwrapToScalar().to<float>();
-               auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+               auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar());
                auto div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n));
                TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n);
 
@@ -381,8 +380,7 @@ auto element_wise_registrations TORCHTRT_UNUSED =
             {"aten::div_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)",
              [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                auto self = args[0].ITensorOrFreeze(ctx);
-               auto otherScalar = args[1].unwrapToScalar().to<float>();
-               auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+               auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar());
                auto div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n));
                TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n);
 
@@ -481,18 +479,12 @@ auto element_wise_registrations TORCHTRT_UNUSED =
             {"aten::ne.Scalar(Tensor self, Scalar other) -> (Tensor)",
              [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                auto self = args[0].ITensorOrFreeze(ctx);
-               auto scalar = args[1].unwrapToScalar();
-               nvinfer1::ITensor* scalar_tensor;
-               if (self->getType() == nvinfer1::DataType::kFLOAT || self->getType() == nvinfer1::DataType::kHALF) {
-                 scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar.to<float>()}));
-               } else {
-                 scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar.to<int>()}));
-               }
+               auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar());
                auto equal = add_elementwise(
                    ctx,
                    nvinfer1::ElementWiseOperation::kEQUAL,
                    self,
-                   scalar_tensor,
+                   other,
                    util::node_info(n) + std::string("is_equal"));
                TORCHTRT_CHECK(equal, "Unable to create elementwise equal layer from node: " << *n);
                // XOR with ones negates and produces not_equal result
@@ -534,8 +526,7 @@ auto element_wise_registrations TORCHTRT_UNUSED =
             {"aten::pow.Tensor_Scalar(Tensor self, Scalar exponent) -> (Tensor)",
              [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                auto self = args[0].ITensorOrFreeze(ctx);
-               auto exponentScalar = args[1].unwrapToScalar().to<float>();
-               auto exponent = tensor_to_const(ctx, torch::tensor({exponentScalar}));
+               auto exponent = scalar_to_tensor(ctx, args[1].unwrapToScalar());
                auto pow =
                    add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPOW, self, exponent, util::node_info(n));
                TORCHTRT_CHECK(pow, "Unable to create Power layer from node: " << *n);
@@ -681,9 +672,9 @@ auto element_wise_registrations TORCHTRT_UNUSED =
             {"aten::eq.Scalar(Tensor self, Scalar other) -> (Tensor)",
              [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                auto self = args[0].ITensorOrFreeze(ctx);
-               auto otherScalar = args[1].unwrapToScalar().to<float>();
-               auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+               auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar());
                if (self->getType() == nvinfer1::DataType::kBOOL) {
+                 auto otherScalar = args[1].unwrapToScalar().to<float>();
                  if (otherScalar == 0 || otherScalar == 1) {
                    LOG_DEBUG("Since input tensor is type bool, casting input tensor and scalar to int32");
                    other = castITensor(ctx, other, nvinfer1::DataType::kINT32);

core/conversion/converters/impl/normalize.cpp

@@ -53,23 +53,118 @@ void create_plugin(
   LOG_DEBUG("Normalize layer output tensor shape: " << layer_output->getDimensions());
 }
 
-auto normalize_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(
-    {"aten::norm.ScalarOpt_dim(Tensor self, Scalar? p, int[1] dim, bool keepdim=False) -> (Tensor)",
-     [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
-       auto in = args[0].ITensor();
-       auto in_shape = util::toVec(in->getDimensions());
-       auto order = args[1].unwrapToScalar().to<int32_t>();
-       auto axes_values = args[2].unwrapToIntList().vec();
-       std::vector<int32_t> axes(axes_values.begin(), axes_values.end());
-       auto keep_dims = (int32_t)args[3].unwrapToBool();
-       LOG_DEBUG("Order of normalize_plugin: " << order);
-       LOG_DEBUG("Axis: " << axes);
-       LOG_DEBUG("keep_dims: " << keep_dims);
-       create_plugin(ctx, n, in, order, axes, keep_dims, "NormalizePluginTorchTRT");
-       return true;
-     }
-
-    });
+int32_t axes_mask_from_axes_values(
+    const torch::jit::Node* n,
+    int32_t nb_dims,
+    const std::vector<int64_t>& axes_values) {
+  int32_t axes_mask = 0;
+  for (size_t i = 0UL; i < axes_values.size(); ++i) {
+    auto axis = axes_values[i];
+    if (axis < 0) {
+      axis += nb_dims;
+    }
+    TORCHTRT_CHECK(
+        axis < nb_dims, util::node_info(n) << " axis " << i << " with value: " << axis << " exceeds input rank");
+    axes_mask += 1 << axis;
+  }
+  return axes_mask;
+}
+
+nvinfer1::ITensor* frobenius_norm(
+    ConversionCtx* ctx,
+    const torch::jit::Node* n,
+    nvinfer1::ITensor* self,
+    int32_t axes_mask,
+    bool keep_dims) {
+  auto squared_layer =
+      add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, self, util::node_info(n) + "_squared");
+  TORCHTRT_CHECK(squared_layer, "Unabled to create square layer from node: " << *n);
+  auto squared_output = squared_layer->getOutput(0);
+
+  auto sum_layer = ctx->net->addReduce(*squared_output, nvinfer1::ReduceOperation::kSUM, axes_mask, keep_dims);
+  TORCHTRT_CHECK(sum_layer, "Unable to create sum layer from node: " << *n);
+  sum_layer->setName((util::node_info(n) + "_sum").c_str());
+  auto sum_output = sum_layer->getOutput(0);
+  LOG_DEBUG("SUM SHAPE: " << sum_output->getDimensions());
+
+  auto sqrt_layer = ctx->net->addUnary(*sum_output, nvinfer1::UnaryOperation::kSQRT);
+  TORCHTRT_CHECK(sqrt_layer, "Unable to create sqrt layer from node: " << *n);
+  sqrt_layer->setName((util::node_info(n) + "_sqrt").c_str());
+  auto sqrt_output = sqrt_layer->getOutput(0);
+  return sqrt_output;
+}
+
+auto normalize_registrations TORCHTRT_UNUSED =
+    RegisterNodeConversionPatterns()
+        .pattern(
+            {"aten::norm.ScalarOpt_dim(Tensor self, Scalar? p, int[1] dim, bool keepdim=False) -> (Tensor)",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               auto in = args[0].ITensorOrFreeze(ctx);
+               auto in_shape = util::toVec(in->getDimensions());
+               auto order = args[1].unwrapToScalar().to<int32_t>();
+               auto axes_values = args[2].unwrapToIntList().vec();
+               std::vector<int32_t> axes(axes_values.begin(), axes_values.end());
+               auto keep_dims = (int32_t)args[3].unwrapToBool();
+               LOG_DEBUG("Order of normalize_plugin: " << order);
+               LOG_DEBUG("Axis: " << axes);
+               LOG_DEBUG("keep_dims: " << keep_dims);
+               create_plugin(ctx, n, in, order, axes, keep_dims, "NormalizePluginTorchTRT");
+               return true;
+             }
+
+            })
+        .pattern(
+            {"aten::frobenius_norm.dim(Tensor self, int[1] dim, bool keepdim=False) -> (Tensor)",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               auto self = args[0].ITensorOrFreeze(ctx);
+               auto axes_values = args[1].unwrapToIntList().vec();
+               auto keep_dims = args[2].unwrapToBool();
+
+               auto axes_mask = axes_mask_from_axes_values(n, self->getDimensions().nbDims, axes_values);
+
+               auto norm = frobenius_norm(ctx, n, self, axes_mask, keep_dims);
+               auto out = ctx->AssociateValueAndTensor(n->outputs()[0], norm);
+               LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+               return true;
+             }})
+        .pattern(
+            {"aten::linalg_norm(Tensor self, Scalar? ord=None, int[1]? dim=None, bool keepdim=False, *, int? dtype=None) -> (Tensor)",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               // https://pytorch.org/docs/stable/generated/torch.linalg.norm.html
+               auto self = args[0].ITensorOrFreeze(ctx);
+               TORCHTRT_CHECK(
+                   args[1].IValue()->isNone(),
+                   "aten::linalg_norm converter does not yet support non-None 'ord' arguments. Add aten::linalg_norm to torch_executed_ops to force it to fallback.");
+               auto keep_dims = args[3].unwrapToBool();
+               auto self_nb_dims = self->getDimensions().nbDims;
+
+               if (!args.back().IValue()->isNone()) {
+                 // If specified, the input tensor is cast to dtype before performing the operation, and the returned
+                 // tensor’s type will be dtype
+                 auto dtype = args.back().unwrapToScalar().to<int64_t>();
+                 auto trt_dtype = util::ScalarTypeToTRTDataType(static_cast<at::ScalarType>(dtype));
+                 self = castITensor(ctx, self, trt_dtype);
+               }
+
+               int32_t axes_mask = 0;
+               if (args[2].IValue()->isNone()) {
+                 // If dim= None and ord= None, self will be flattened to 1D and the 2-norm of the resulting vector will
+                 // be computed.
+                 axes_mask = 1;
+                 keep_dims = true; // the single output dim is always preserved
+                 auto flatten_layer = ctx->net->addShuffle(*self);
+                 TORCHTRT_CHECK(flatten_layer, "Unable to create shuffle layer from node: " << *n);
+                 flatten_layer->setReshapeDimensions(util::toDims(std::vector<int64_t>({-1})));
+                 flatten_layer->setName((util::node_info(n) + "_flatten").c_str());
+                 self = flatten_layer->getOutput(0);
+               } else {
+                 axes_mask = axes_mask_from_axes_values(n, self_nb_dims, args[2].unwrapToIntList().vec());
+               }
+               auto norm = frobenius_norm(ctx, n, self, axes_mask, keep_dims);
+               auto out = ctx->AssociateValueAndTensor(n->outputs()[0], norm);
+               LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+               return true;
+             }});
 
 } // namespace
 } // namespace impl