Enable IR checks

Differential Revision: D59638386 Pull Request resolved: pytorch#7408
kirklandsign · Dec 20, 2024 · 5190106 · 5190106
1 parent 59a2594
commit 5190106
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/backends/cadence/aot/compiler.py b/backends/cadence/aot/compiler.py
@@ -201,7 +201,15 @@ def export_to_edge(
     edge_prog_manager = to_edge(
         expo_program,
         compile_config=EdgeCompileConfig(
-            _check_ir_validity=False, _skip_dim_order=True
+            _skip_dim_order=True,
+            # Allow specific non-core aten ops in the IR.
+            _core_aten_ops_exception_list=[
+                torch.ops.aten.linear.default,
+                torch.ops.aten.native_batch_norm.default,
+                torch.ops.aten.linalg_vector_norm.default,
+                torch.ops.aten.unfold.default,
+                torch.ops.aten.angle.default,
+            ],
         ),
     )
 

diff --git a/backends/cadence/aot/replace_ops.py b/backends/cadence/aot/replace_ops.py
@@ -579,7 +579,7 @@ def call_operator(self, op, args, kwargs, meta):
                 arg_shape[:dim] + (left_padding_size,) + arg_shape[dim + 1 :]
             )
             left_padding_node = super().call_operator(
-                torch.ops.aten.full.default,
+                exir_ops.edge.aten.full.default,
                 (
                     left_padding_shape,
                     value,
@@ -596,7 +596,7 @@ def call_operator(self, op, args, kwargs, meta):
                 arg_shape[:dim] + (right_padding_size,) + arg_shape[dim + 1 :]
             )
             right_padding_node = super().call_operator(
-                torch.ops.aten.full.default,
+                exir_ops.edge.aten.full.default,
                 (
                     right_padding_shape,
                     value,
@@ -726,7 +726,7 @@ def call_operator(self, op, args, kwargs, meta):
 
             flipped_weight = (
                 super().call_operator(
-                    torch.ops.aten.flip.default,
+                    exir_ops.edge.aten.flip.default,
                     (
                         transposed_weight,
                         [-1] if transposed_weight.to_tensor().dim() == 3 else [-1, -2],