From b313ace49795f028c45d2eefc45450743080c54f Mon Sep 17 00:00:00 2001
From: Songhao Jia <gasoonjia@meta.com>
Date: Fri, 24 Jan 2025 11:32:43 -0800
Subject: [PATCH] turn on dim order in cadence test (#7756)
Summary:
This diff turns on dim order in cadence test. Also we get around `to_copy` operator in verifier to keep the verifier check enable.
Differential Revision: D68246404
---
backends/cadence/aot/compiler.py | 6 +++++-
backends/cadence/aot/replace_ops.py | 32 +++++++++++++++++++++++++++++
2 files changed, 37 insertions(+), 1 deletion(-)
diff --git a/backends/cadence/aot/compiler.py b/backends/cadence/aot/compiler.py
index 834cd495aea..ff5be8fb605 100644
--- a/backends/cadence/aot/compiler.py
+++ b/backends/cadence/aot/compiler.py
@@ -33,6 +33,7 @@
ExecutorchProgramManager,
to_edge,
)
+from executorch.exir.dialects._ops import ops as exir_ops
from executorch.exir.pass_base import PassResult
from executorch.exir.passes import ToOutVarPass
from executorch.exir.passes.sym_shape_eval_pass import HintBasedSymShapeEvalPass
@@ -185,7 +186,6 @@ def export_to_edge(
edge_prog_manager = to_edge(
expo_program,
compile_config=EdgeCompileConfig(
- _skip_dim_order=True,
# Allow specific non-core aten ops in the IR.
_core_aten_ops_exception_list=[
torch.ops.aten._native_batch_norm_legit_functional.default,
@@ -193,6 +193,10 @@ def export_to_edge(
torch.ops.aten.linalg_vector_norm.default,
torch.ops.aten.unfold.default,
torch.ops.aten.angle.default,
+ # cadence replaced to_dim_order_copy with _to_copy for performance
+ # skip _to_copy op to get around of dim order check
+ # We should remove this op once cadence can support dim order
+ exir_ops.edge.aten._to_copy.default,
],
),
)
diff --git a/backends/cadence/aot/replace_ops.py b/backends/cadence/aot/replace_ops.py
index 89ef821c569..fd04abcf383 100644
--- a/backends/cadence/aot/replace_ops.py
+++ b/backends/cadence/aot/replace_ops.py
@@ -1799,6 +1799,37 @@ def call_operator(
)
+@register_cadence_pass(CadencePassAttribute(opt_level=0))
+class ReplaceToDimOrderCopyWithToCopyPass(ExportPass):
+ """
+ dim_order_ops::to_dim_order_copy is not supported, so this is an opt_level=0 pass.
+ If the dim order is sequential, we don't need the extra work with strides and
+ can just use to_copy.
+ """
+
+ def call_operator(
+ self,
+ op,
+ args: Tuple[Argument, ...],
+ kwargs: Dict[str, Argument],
+ meta: NodeMetadata,
+ ) -> ProxyValue:
+ if op != exir_ops.edge.dim_order_ops._to_dim_order_copy.default:
+ return super().call_operator(op, args, kwargs, meta)
+
+ # pyre-ignore[16]: `None` has no attribute `to_tensor`.
+ assert (args[0] == range(args[0].to_tensor().dim()), "Only sequential dims supported")
+
+ return super().call_operator(
+ exir_ops.edge.aten._to_copy.default,
+ (
+ args[0],
+ ),
+ {},
+ meta,
+ )
+
+
@register_cadence_pass(CadencePassAttribute(opt_level=0))
class ReplaceFullLikeWithFullPass(ExportPass):
"""
@@ -2108,4 +2139,5 @@ class CadenceReplaceOpsInGraph:
ReplaceSingleElementTensorArgumentsFromFullOpWithScalarPass,
ReplaceAtenAvgPoolWithJarvisAvgPoolPass,
ReplaceAtenLinalgVectorNormWithCadenceLinalgVectorNormPass,
+ ReplaceToDimOrderCopyWithToCopyPass,
]