Introduce dim_order edge dialect op for .to()

Summary:
These ops takes dim_order instead of memory_format.

These are new edge dialect only ops for now. In the future we are planning to introduce these into aten (and the aten cope opset).

Reviewed By: larryliu0820

Differential Revision: D48195056

fbshipit-source-id: 2ba97024647553ae99517bfbc42f1aba78395c13

Author: digantdesai

exir/passes/TARGETS

@@ -255,3 +255,12 @@ python_library(
         "//executorch/exir/dialects:lib",
     ],
 )
+
+python_library(
+    name = "dim_order_ops_registry",
+    srcs = ["dim_order_ops_registry.py"],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/exir:dim_order_utils",
+    ],
+)

exir/passes/dim_order_ops_registry.py

@@ -0,0 +1,47 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.dim_order_utils import get_memory_format
+
+from torch.library import impl, Library
+
+lib = Library("dim_order_ops", "DEF")
+lib.define(
+    "_to_dim_order_copy(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, bool non_blocking=False, int[]? dim_order=None) -> Tensor"
+)
+
+# Out variant drops TensorOptions
+lib.define(
+    "_to_dim_order_copy.out(Tensor self, *, bool non_blocking=False, int[]? dim_order=None, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+def _op_impl(target, *args, **kwargs):
+    kwargs["memory_format"] = get_memory_format(kwargs["dim_order"])
+    _ = kwargs.pop("dim_order")
+    res = target(*args, **kwargs)
+    # assert list(res.dim_order()) == dim_order
+    return res
+
+
+@impl(lib, "_to_dim_order_copy", "CompositeImplicitAutograd")
+def _to_dim_order_copy_impl(*args, **kwargs):
+    return _op_impl(torch.ops.aten._to_copy, *args, **kwargs)
+
+
+@impl(lib, "_to_dim_order_copy.out", "CompositeImplicitAutograd")
+def _to_dim_order_copy_out_impl(*args, **kwargs):
+    return _op_impl(torch.ops.aten._to_copy.out, *args, **kwargs)
+
+
+"""
+Defines a map of aten or edge ops to the corresponding dim_order ops for quick lookup
+"""
+DimOrderOpsMap = {
+    "aten._to_copy.default": exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+}