Merge pull request #382 from NVIDIA/aten_to

narendasan · web-flow · commit 1d55db2bf6b6 · 2021-08-10T19:42:08.000-06:00
feat(aten::to): Add support for cast layer conversion
diff --git a/core/conversion/conversion_ignorelist.cpp b/core/conversion/conversion_ignorelist.cpp
@@ -16,7 +16,6 @@ const std::unordered_set<std::string>& get_non_convertable_nodes() {
     "aten::backward",
     "aten::save",
     "aten::contiguous",
-    "aten::to",
     "prim::RaiseException",
     "prim::Print",
     "prim::device",
diff --git a/core/conversion/converters/BUILD b/core/conversion/converters/BUILD
@@ -54,6 +54,7 @@ cc_library(
         "NodeConverterRegistry.cpp",
         "impl/activation.cpp",
         "impl/batch_norm.cpp",
+        "impl/cast.cpp",
         "impl/concat.cpp",
         "impl/constant.cpp",
         "impl/constant_pad.cpp",
diff --git a/core/conversion/converters/converter_util.cpp b/core/conversion/converters/converter_util.cpp
@@ -142,7 +142,7 @@ nvinfer1::ITensor* castITensor(ConversionCtx* ctx, nvinfer1::ITensor* tensor, nv
   }
 }
 
-nvinfer1::ITensor* tensor_to_const(ConversionCtx* ctx, at::Tensor t) {
+nvinfer1::ITensor* tensor_to_const(ConversionCtx* ctx, at::Tensor t, const std::string& name) {
   bool post_freeze_cast = false;
   nvinfer1::DataType post_freeze_cast_type = nvinfer1::DataType::kFLOAT;
   // Other "unsupported weights types" can be added to this check here
@@ -175,9 +175,15 @@ nvinfer1::ITensor* tensor_to_const(ConversionCtx* ctx, at::Tensor t) {
 
   std::ostringstream tensor_id;
   tensor_id << reinterpret_cast<int*>(out);
+  std::string tensor_name;
 
-  LOG_DEBUG(ctx->logger, "Freezing tensor " << tensor_id.str() << " as an IConstantLayer");
-  const_layer->setName(("[Freeze Tensor " + tensor_id.str() + " ]").c_str());
+  if (!name.empty()) {
+    tensor_name = name;
+  } else {
+    tensor_name = tensor_id.str();
+  }
+  LOG_DEBUG(ctx->logger, "Freezing tensor " << tensor_name << " as an IConstantLayer");
+  const_layer->setName(("[Freeze Tensor " + tensor_name + " ]").c_str());
 
   if (post_freeze_cast) {
     out = castITensor(ctx, out, post_freeze_cast_type);
diff --git a/core/conversion/converters/converter_util.h b/core/conversion/converters/converter_util.h
@@ -45,7 +45,7 @@ nvinfer1::ILayer* add_elementwise(
 nvinfer1::ITensor* castITensor(ConversionCtx* ctx, nvinfer1::ITensor* tensor, nvinfer1::DataType dtype);
 
 // Freeze an at::Tensor in a IConstant layer
-nvinfer1::ITensor* tensor_to_const(ConversionCtx* ctx, at::Tensor t);
+nvinfer1::ITensor* tensor_to_const(ConversionCtx* ctx, at::Tensor t, const std::string& name = std::string());
 
 } // namespace converters
 } // namespace conversion
diff --git a/core/conversion/converters/impl/cast.cpp b/core/conversion/converters/impl/cast.cpp
@@ -0,0 +1,63 @@
+#include <torch/torch.h>
+#include "core/conversion/converters/converter_util.h"
+#include "core/conversion/converters/converters.h"
+#include "core/util/prelude.h"
+#include "core/util/trt_util.h"
+
+namespace trtorch {
+namespace core {
+namespace conversion {
+namespace converters {
+namespace impl {
+namespace {
+
+auto cast_registrations TRTORCH_UNUSED =
+    RegisterNodeConversionPatterns()
+        .pattern(
+            {"aten::to.dtype(Tensor self, int dtype, bool non_blocking=False, bool copy=False, int? memory_format=None) -> (Tensor)",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               auto self = args[0].ITensorOrFreeze(ctx);
+               auto output_dtype = args[1].unwrapToScalar().to<int64_t>();
+               auto trt_dtype = util::ScalarTypeToTRTDataType(static_cast<at::ScalarType>(output_dtype));
+               auto casted_itensor = castITensor(ctx, self, trt_dtype);
+               auto output = ctx->AssociateValueAndTensor(n->outputs()[0], casted_itensor);
+               LOG_DEBUG("[aten::to.dtype] Output tensor shape: " << output->getDimensions());
+
+               return true;
+             }})
+        .pattern(
+            {"aten::to.other(Tensor self, Tensor other, bool non_blocking=False, bool copy=False, int? memory_format=None) -> (Tensor)",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               auto self = args[0].ITensorOrFreeze(ctx);
+               nvinfer1::DataType other_dtype = args[1].ITensorOrFreeze(ctx)->getType();
+               auto casted_itensor = castITensor(ctx, self, other_dtype);
+               auto output = ctx->AssociateValueAndTensor(n->outputs()[0], casted_itensor);
+               LOG_DEBUG("[aten::to.other] Output tensor shape: " << output->getDimensions());
+
+               return true;
+             }})
+        .pattern(
+            {"aten::to.prim_Device(Tensor(a) self, Device? device, int? dtype=None, bool non_blocking=False, bool copy=False) -> (Tensor(b|a))",
+             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+               auto self = args[0].ITensorOrFreeze(ctx);
+               if (args[2].isIValue() && !args[2].IValue()->isScalar()) {
+                 auto output = ctx->AssociateValueAndTensor(n->outputs()[0], self);
+                 LOG_DEBUG("[aten::to.prim_Device] Output tensor shape: " << output->getDimensions());
+                 return true;
+               }
+
+               auto output_dtype = args[2].unwrapToScalar().to<int64_t>();
+               auto trt_dtype = util::ScalarTypeToTRTDataType(static_cast<at::ScalarType>(output_dtype));
+               auto casted_itensor = castITensor(ctx, self, trt_dtype);
+               auto output = ctx->AssociateValueAndTensor(n->outputs()[0], casted_itensor);
+               LOG_DEBUG("[aten::to.prim_Device] Output tensor shape: " << output->getDimensions());
+
+               return true;
+             }});
+// clang-format on
+} // namespace
+} // namespace impl
+} // namespace converters
+} // namespace conversion
+} // namespace core
+} // namespace trtorch
diff --git a/core/conversion/converters/impl/constant.cpp b/core/conversion/converters/impl/constant.cpp
@@ -16,12 +16,16 @@ auto constant_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
               // used for Fundimentally this is because of the differing
               // philosophies between TensorRT and PyTorch, i.e. Variables contain
               // Tensors vs. just Tensors
-
-              auto t = args[0].unwrapToTensor();
-              auto const_out = ctx->AssociateValueAndTensor(n->outputs()[0], tensor_to_const(ctx, t));
-
-              LOG_DEBUG("Output tensor shape: " << const_out->getDimensions());
-
+              nvinfer1::ITensor* output;
+              if (args[0].isITensor()){
+                output = ctx->AssociateValueAndTensor(n->outputs()[0], args[0].ITensor());
+              } else{
+                auto t = args[0].unwrapToTensor();
+                auto const_out = tensor_to_const(ctx, t, util::node_info(n).c_str());
+                output = ctx->AssociateValueAndTensor(n->outputs()[0], const_out);
+              }
+              LOG_DEBUG("Output tensor shape: " << output->getDimensions());
+              
               return true;
             }});
 // clang-format on
diff --git a/core/conversion/converters/impl/shuffle.cpp b/core/conversion/converters/impl/shuffle.cpp
@@ -125,6 +125,33 @@ static auto shuffle_registrations TRTORCH_UNUSED =
                     auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0));
                     LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
 
+                    return true;
+                  }})
+        .pattern({"aten::t(Tensor self) -> Tensor",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    auto in = args[0].ITensorOrFreeze(ctx);
+                    auto input_dims = in->getDimensions();
+                    // For input tensors < 2D, return them as is
+                    // For a 2D input tensor, return transpose(input, 0, 1) which is a general 2d matrix transpose.
+                    if (input_dims.nbDims < 2) {
+                      auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], in);
+                      LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+                      return true;
+                    }
+
+                    auto shuffle_layer = ctx->net->addShuffle(*in);
+                    TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
+                    nvinfer1::Permutation firstPerm;
+                    firstPerm.order[0] = 1;
+                    firstPerm.order[1] = 0;
+
+                    shuffle_layer->setFirstTranspose(firstPerm);
+                    shuffle_layer->setZeroIsPlaceholder(false);
+                    shuffle_layer->setName(util::node_info(n).c_str());
+
+                    auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle_layer->getOutput(0));
+                    LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+
                     return true;
                   }})
         .pattern({"aten::pixel_shuffle(Tensor self, int upscale_factor) -> (Tensor)",
diff --git a/core/conversion/evaluators/aten.cpp b/core/conversion/evaluators/aten.cpp
@@ -468,20 +468,6 @@ auto aten_registrations TRTORCH_UNUSED =
                     EvalOptions().validSchemas({
                         "aten::numel(Tensor self) -> int",
                     })})
-        .evaluator({c10::Symbol::fromQualString("aten::t"),
-                    [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
-                      auto tensor_var = args.at(n->input(0));
-                      if (tensor_var.IValue()->isTensor()) {
-                        auto tensor = tensor_var.unwrapToTensor();
-                        return tensor.t();
-                      } else {
-                        TRTORCH_THROW_ERROR("Unimplemented data type for aten::t evaluator: ITensor");
-                        return {};
-                      }
-                    },
-                    EvalOptions().validSchemas({
-                        "aten::t(Tensor self) -> Tensor",
-                    })})
         .evaluator({c10::Symbol::fromQualString("aten::dim"),
                     [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
                       auto tensor_var = args.at(n->input(0));
diff --git a/core/conversion/evaluators/prim.cpp b/core/conversion/evaluators/prim.cpp
@@ -12,6 +12,7 @@
 #include "core/conversion/evaluators/eval_macros.h"
 #include "core/conversion/evaluators/eval_util.h"
 #include "core/conversion/evaluators/evaluators.h"
+#include "core/util/trt_util.h"
 
 namespace trtorch {
 namespace core {
@@ -101,6 +102,28 @@ auto prim_registrations =
                         return c10::optional<torch::jit::IValue>(std::move(torch::jit::IValue(list)));
                       }
                     }})
+        .evaluator({c10::Symbol::fromQualString("prim::dtype"),
+                    [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+                      auto input = args.at(n->input(0));
+                      if (input.isITensor()) {
+                        auto trt_dtype = input.ITensor()->getType();
+                        return static_cast<int>(util::TRTDataTypeToScalarType(trt_dtype));
+                      } else if (input.isIValue()) {
+                        if (input.IValue()->isTensor()) {
+                          auto pyt_input = input.IValue()->toTensor();
+                          return static_cast<int>(pyt_input.scalar_type());
+                        } else {
+                          TRTORCH_THROW_ERROR("Unsupported input type in prim::dtype operator");
+                          return {};
+                        }
+                      } else {
+                        TRTORCH_THROW_ERROR("Unsupported input type in prim::dtype operator");
+                        return {};
+                      }
+                    },
+                    EvalOptions().validSchemas({
+                        "prim::dtype(Tensor a) -> (int)",
+                    })})
         .evaluator({c10::Symbol::fromQualString("prim::min"),
                     [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
                       if (n->inputs().size() == 1) {
diff --git a/core/lowering/lowering.cpp b/core/lowering/lowering.cpp
@@ -35,6 +35,7 @@ void LowerGraph(std::shared_ptr<torch::jit::Graph>& g) {
   passes::EliminateExceptionOrPassPattern(g);
   torch::jit::FuseLinear(g);
   torch::jit::LowerAllTuples(g);
+  passes::ReduceToOperation(g);
   passes::RemoveContiguous(g);
   passes::RemoveDropout(g);
   passes::LinearToAddMM(g);
diff --git a/core/lowering/passes/BUILD b/core/lowering/passes/BUILD
@@ -16,6 +16,7 @@ cc_library(
         "fuse_addmm_branches.cpp",
         "linear_to_addmm.cpp",
         "op_aliasing.cpp",
+        "reduce_to.cpp",
         "remove_bn_dim_check.cpp",
         "remove_contiguous.cpp",
         "remove_dropout.cpp",
diff --git a/core/lowering/passes/passes.h b/core/lowering/passes/passes.h
@@ -12,6 +12,7 @@ void Conv3DToConvolution(std::shared_ptr<torch::jit::Graph>& graph);
 void FuseAddMMBranches(std::shared_ptr<torch::jit::Graph> graph);
 void LinearToAddMM(std::shared_ptr<torch::jit::Graph>& graph);
 void EliminateExceptionOrPassPattern(std::shared_ptr<torch::jit::Graph> graph);
+void ReduceToOperation(std::shared_ptr<torch::jit::Graph>& graph);
 void RemoveBNDimCheck(std::shared_ptr<torch::jit::Graph> graph);
 void RemoveContiguous(std::shared_ptr<torch::jit::Graph>& graph);
 void RemoveDropout(std::shared_ptr<torch::jit::Graph>& graph);
diff --git a/core/lowering/passes/reduce_to.cpp b/core/lowering/passes/reduce_to.cpp
@@ -0,0 +1,48 @@
+#include <torch/csrc/jit/passes/subgraph_rewrite.h>
+
+#include "core/util/prelude.h"
+
+namespace trtorch {
+namespace core {
+namespace lowering {
+namespace passes {
+
+void ReduceToOperation(std::shared_ptr<torch::jit::Graph>& graph) {
+  std::string to_device_pattern = R"IR(
+        graph(%x, %device, %dtype, %nb, %copy, %format):
+            %out : Tensor = aten::to(%x, %device, %dtype, %nb, %copy, %format)
+            return (%out))IR";
+  std::string to_dtype_pattern = R"IR(
+        graph(%x, %device, %dtype, %nb, %copy, %format):
+            %out : Tensor = aten::to(%x, %dtype, %nb, %copy, %format)
+            return (%out))IR";
+
+  std::string to_type_as_pattern = R"IR(
+        graph(%input, %other):
+            %out : Tensor = aten::type_as(%input, %other)
+            return (%out))IR";
+
+  std::string to_other_pattern = R"IR(
+        graph(%input, %other):
+            %5 : bool = prim::Constant[value=0]()
+            %6 : None = prim::Constant()
+            %out : Tensor = aten::to(%input, %other, %5, %5, %6)
+            return (%out))IR";
+
+  // replace aten::to.device with aten::to.dtype
+  torch::jit::SubgraphRewriter map_aten_device_to_dtype;
+  map_aten_device_to_dtype.RegisterRewritePattern(to_device_pattern, to_dtype_pattern);
+  map_aten_device_to_dtype.runOnGraph(graph);
+
+  // replace aten::type_as with aten::to.other
+  torch::jit::SubgraphRewriter map_aten_type_as_to_other;
+  map_aten_type_as_to_other.RegisterRewritePattern(to_type_as_pattern, to_other_pattern);
+  map_aten_type_as_to_other.runOnGraph(graph);
+
+  LOG_GRAPH("Post lowering of [aten::to.device|aten::type_as] -> " << *graph);
+}
+
+} // namespace passes
+} // namespace lowering
+} // namespace core
+} // namespace trtorch
diff --git a/core/lowering/passes/remove_nops.cpp b/core/lowering/passes/remove_nops.cpp
@@ -20,12 +20,10 @@ struct NOPRemoval {
   NOPRemoval(std::shared_ptr<Graph> graph) : graph_(std::move(graph)) {}
 
   void run() {
-    removeNode(graph_->block(), "aten::to");
     removeNode(graph_->block(), "aten::detach");
     torch::jit::EliminateDeadCode(graph_);
-    LOG_DEBUG(
-        "RemoveNOPs - Note: Removing remaining aten::to operators (in addition to other ops that have no meaning in TRT), if type casts need to be preserved, add a pass before this pass is run");
-    LOG_GRAPH("Post aten::to removal: " << *graph_);
+    LOG_DEBUG("RemoveNOPs - Note: Removing operators that have no meaning in TRT");
+    LOG_GRAPH("Post aten::detach removal: " << *graph_);
   }
 
  private:
@@ -52,4 +50,4 @@ void RemoveNOPs(std::shared_ptr<Graph> graph) {
 } // namespace passes
 } // namespace lowering
 } // namespace core
-} // namespace trtorch
+} // namespace trtorch
diff --git a/tests/core/conversion/converters/BUILD b/tests/core/conversion/converters/BUILD
@@ -15,6 +15,10 @@ converter_test(
     name = "test_batch_norm",
 )
 
+converter_test(
+    name = "test_cast",
+)
+
 converter_test(
     name = "test_clone",
 )
@@ -120,6 +124,7 @@ test_suite(
     tests = [
         ":test_activation",
         ":test_batch_norm",
+        ":test_cast",
         ":test_clone",
         ":test_concat",
         ":test_constant_pad",
diff --git a/tests/core/conversion/converters/test_cast.cpp b/tests/core/conversion/converters/test_cast.cpp
diff --git a/tests/core/lowering/BUILD b/tests/core/lowering/BUILD
diff --git a/tests/core/lowering/test_reduce_to_pass.cpp b/tests/core/lowering/test_reduce_to_pass.cpp
diff --git a/tests/core/lowering/test_remove_to.cpp b/tests/core/lowering/test_remove_to.cpp
