From 629c93e4e9dfd44cec933d2320e64e786a992cfb Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 4 Dec 2024 19:14:52 -0800 Subject: [PATCH 1/2] [et][dim order] Makes DimOrderOpsMap as an operator to operator mapping Pull Request resolved: https://github.com/pytorch/executorch/pull/7187 This diff updates DimOrderOpsMap from name-to-operator mapping to operator-to-operator mapping, which has multiple benefits: 1. Reduce dialect ambiguity. Different dialects op may map to same name (e.g. `aten.to_copy` and `exir_ops.edge.aten._to_copy.default`). Directly using op can diminish the ambiguity. 2. Auto-maintain MemoryFormatOpsMap by reverting DimOrderOpsMap ghstack-source-id: 256633613 @exported-using-ghexport Differential Revision: [D66773612](https://our.internmc.facebook.com/intern/diff/D66773612/) --- exir/passes/dim_order_ops_registry.py | 18 +++++------------- exir/passes/memory_format_ops_pass.py | 12 ++++++------ exir/verification/verifier.py | 1 + 3 files changed, 12 insertions(+), 19 deletions(-) diff --git a/exir/passes/dim_order_ops_registry.py b/exir/passes/dim_order_ops_registry.py index c4436aaa910..f3fc009f109 100644 --- a/exir/passes/dim_order_ops_registry.py +++ b/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()} diff --git a/exir/passes/memory_format_ops_pass.py b/exir/passes/memory_format_ops_pass.py index ba89a510a71..db597b2b233 100644 --- a/exir/passes/memory_format_ops_pass.py +++ b/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, diff --git a/exir/verification/verifier.py b/exir/verification/verifier.py index 2c45929bf23..f906623ca25 100644 --- a/exir/verification/verifier.py +++ b/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 " ) From 735a09746d4455841a13a4c53ec80c18cc8ea95d Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 4 Dec 2024 20:53:47 -0800 Subject: [PATCH 2/2] [et][dim order] Make edge verifier support empty operator as title. Differential Revision: [D66801315](https://our.internmc.facebook.com/intern/diff/D66801315/) ghstack-source-id: 256642462 Pull Request resolved: https://github.com/pytorch/executorch/pull/7188 --- exir/verification/test/test_verifier.py | 5 +++-- exir/verification/verifier.py | 6 ++++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/exir/verification/test/test_verifier.py b/exir/verification/test/test_verifier.py index b2e31dbc599..1ee48ef4d43 100644 --- a/exir/verification/test/test_verifier.py +++ b/exir/verification/test/test_verifier.py @@ -117,8 +117,9 @@ def __init__(self): def forward(self, x: torch.Tensor) -> torch.Tensor: t1 = x.to(dtype=torch.double, memory_format=torch.channels_last) - t2 = t1 + t1 - return t1 * t2 + t2 = torch.empty(t1.size(), memory_format=torch.channels_last) + t2.copy_(t1) + return t2 m = Model().eval() diff --git a/exir/verification/verifier.py b/exir/verification/verifier.py index f906623ca25..4a37b457e30 100644 --- a/exir/verification/verifier.py +++ b/exir/verification/verifier.py @@ -15,10 +15,12 @@ from executorch.exir.dialects.edge._ops import EdgeOpOverload from executorch.exir.error import ExportError, ExportErrorType from executorch.exir.lowered_backend_module import LoweredBackendModule +from executorch.exir.passes.dim_order_ops_registry import DimOrderOpsMap from executorch.exir.verification.arg_validator import ( EdgeOpArgValidator, RunHigherOrderOperatorError, ) + from torch._dispatch.python import enable_python_dispatcher from torch._export.utils import _detect_fake_mode_from_gm @@ -44,7 +46,7 @@ def _check_tensors_are_contiguous(gm: GraphModule) -> None: def _check_valid_dim_order_ops(op, use_dim_order) -> None: if use_dim_order: - if op in (torch.ops.aten._to_copy.default,): + if op in DimOrderOpsMap: raise SpecViolationError(f"{op} should not be used in dim_order mode") else: # not using dim_order if op.namespace in ("dim_order_ops",): @@ -249,7 +251,7 @@ def check_valid_edge_op(self, op): ) ) if isinstance(op, EdgeOpOverload): - _check_valid_dim_order_ops(op._op, self.use_dim_order) + _check_valid_dim_order_ops(op, self.use_dim_order) self.check_valid_aten_op(op._op) if isinstance(op, types.FunctionType):