feat: Implement lowering for aten::to.dtype schema

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

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, %nb, %copy, %format, %other):
+            %out : Tensor = aten::to.dtype_layout(%x, %device, %dtype, %layout, %nb, %copy, %format, %other)
+            return (%out))IR";
+
+  std::string to_dtype_multi_input_pattern = R"IR(
+        graph(%x, %device, %dtype, %layout, %nb, %copy, %format, %other):
+            %out : Tensor = aten::to(%x, %device, %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);