[et][dim order] Makes DimOrderOpsMap as an operator to operator mapping

Differential Revision: D66773612

Pull Request resolved: #7187

Author: Gasoonjia

exir/passes/dim_order_ops_registry.py

@@ -58,22 +58,14 @@ def _empty_dim_order_out_impl(*args, **kwargs):
 
 
 """
-Defines a map of aten or edge ops to the corresponding dim_order ops for quick lookup
+Defines a map of 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,
-    "aten.empty.memory_format": exir_ops.edge.dim_order_ops._empty_dim_order.default,
+    exir_ops.edge.aten._to_copy.default: exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+    exir_ops.edge.aten.empty.memory_format: exir_ops.edge.dim_order_ops._empty_dim_order.default,
 }
 
 """
-Defines a map of aten or edge ops to the corresponding memory format ops for quick lookup
+Defines a map of edge ops to the corresponding memory format ops for quick lookup, which is the revert of DimOrderOpsMap
 """
-MemoryFormatOpsMap = {
-    "dim_order_ops._to_dim_order_copy.default": exir_ops.edge.aten._to_copy.default,
-    "dim_order_ops._empty_dim_order.default": exir_ops.edge.aten.empty.memory_format,
-}
-
-# If we are replacing an aten op with a dim_order op, we must have a 1:1 mapping through these dicts.
-assert len(DimOrderOpsMap) == len(MemoryFormatOpsMap)
-
-# TODO stricter check for 1:1 mapping
+MemoryFormatOpsMap = {v: k for k, v in DimOrderOpsMap.items()}

exir/passes/memory_format_ops_pass.py

@@ -32,7 +32,7 @@ class MemoryFormatOpsPass(ExportPass):
     """
 
     def call_operator(self, op, args, kwargs, meta):
-        if not (isinstance(op, EdgeOpOverload) and op.__name__ in DimOrderOpsMap):
+        if not (isinstance(op, EdgeOpOverload) and op in DimOrderOpsMap):
             return super().call_operator(
                 op,
                 args,
@@ -61,10 +61,10 @@ def call_operator(self, op, args, kwargs, meta):
         nkwargs["dim_order"] = get_dim_order(mem_format, ndim)
         logger.debug(
             f"{op.__name__} = rank: {ndim}, memory_format: {mem_format}."
-            f" {DimOrderOpsMap[op.__name__].__name__} = dim_order: {nkwargs['dim_order']}"
+            f" {DimOrderOpsMap[op].__name__} = dim_order: {nkwargs['dim_order']}"
         )
 
-        t = DimOrderOpsMap[op.__name__]
+        t = DimOrderOpsMap[op]
 
         return super().call_operator(
             t,
@@ -80,7 +80,7 @@ class DimOrderOpsRevertPass(ExportPass):
     """
 
     def call_operator(self, op, args, kwargs, meta):
-        if not (isinstance(op, EdgeOpOverload) and op.__name__ in MemoryFormatOpsMap):
+        if not (isinstance(op, EdgeOpOverload) and op in MemoryFormatOpsMap):
             return super().call_operator(
                 op,
                 args,
@@ -109,10 +109,10 @@ def call_operator(self, op, args, kwargs, meta):
 
         logger.debug(
             f" {op.__name__} = dim_order: {dim_order}."
-            f" {MemoryFormatOpsMap[op.__name__].__name__} = rank: {ndim}, memory_format: {nkwargs['memory_format']}."
+            f" {MemoryFormatOpsMap[op].__name__} = rank: {ndim}, memory_format: {nkwargs['memory_format']}."
         )
 
-        t = MemoryFormatOpsMap[op.__name__]
+        t = MemoryFormatOpsMap[op]
 
         return super().call_operator(
             t,

exir/verification/verifier.py

@@ -179,6 +179,7 @@ def _check_tensor_args_matching_op_allowed_dtype(gm: GraphModule) -> None:
     if validator.violating_ops:
         raise SpecViolationError(
             f"These operators are taking Tensor inputs with mismatched dtypes: {validator.violating_ops}"
+            "Please make sure the dtypes of the Tensor inputs are the same as the dtypes of the corresponding "
         )