Merge pull request #827 from NVIDIA/peri044/dtype_layout

feat: Implement aten::to.dtype_layout pass

Author: narendasan

core/lowering/passes/reduce_to.cpp

@@ -16,6 +16,15 @@ void ReduceToOperation(std::shared_ptr<torch::jit::Graph>& graph) {
         graph(%x, %device, %dtype, %nb, %copy, %format):
             %out : Tensor = aten::to(%x, %dtype, %nb, %copy, %format)
             return (%out))IR";
+  std::string to_dtype_layout_pattern = R"IR(
+        graph(%x, %device, %dtype, %layout, %pm, %nb, %copy, %format):
+            %out : Tensor = aten::to(%x, %device, %dtype, %layout, %pm, %nb, %copy, %format)
+            return (%out))IR";
+
+  std::string to_dtype_multi_input_pattern = R"IR(
+        graph(%x, %device, %dtype, %layout, %pm, %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):
@@ -34,6 +43,11 @@ void ReduceToOperation(std::shared_ptr<torch::jit::Graph>& graph) {
   map_aten_device_to_dtype.RegisterRewritePattern(to_device_pattern, to_dtype_pattern);
   map_aten_device_to_dtype.runOnGraph(graph);
 
+  // replace aten::to.dtype_layout with aten::to.dtype
+  torch::jit::SubgraphRewriter map_aten_dtype_layout;
+  map_aten_dtype_layout.RegisterRewritePattern(to_dtype_layout_pattern, to_dtype_multi_input_pattern);
+  map_aten_dtype_layout.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);

tests/core/lowering/test_reduce_to_pass.cpp

@@ -28,6 +28,28 @@ TEST(LoweringPasses, ReduceToCorrectly) {
   ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
 }
 
+TEST(LoweringPasses, ReduceToDtypeLayoutCorrectly) {
+  std::string source_graph = R"IR(
+    graph(%x, %device, %dtype, %layout, %pm, %nb, %copy, %format):
+        %out : Tensor = aten::to(%x, %device, %dtype, %layout, %pm, %nb, %copy, %format)
+        return (%out))IR";
+  std::string target_graph = R"IR(
+    graph(%x, %device, %dtype, %layout, %pm, %nb, %copy, %format):
+        %out : Tensor = aten::to(%x, %dtype, %nb, %copy, %format)
+        return (%out))IR";
+
+  torch_tensorrt::core::util::logging::get_logger().set_reportable_log_level(
+      torch_tensorrt::core::util::logging::LogLevel::kGRAPH);
+  auto sg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(source_graph, &*sg);
+  torch_tensorrt::core::lowering::passes::ReduceToOperation(sg);
+
+  auto tg = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(target_graph, &*tg);
+
+  ASSERT_TRUE(!torch::jit::findPatternMatches(*tg, *sg).empty());
+}
+
 TEST(LoweringPasses, ReduceAtenTypeAsCorrectly) {
   std::string source_graph = R"IR(
     graph(%input, %other):