Move executorch_prim_ops_registry into exir

Summary: We need prim ops to be registered at exir level to support sym ops. For some reason this piece of code currently lives in executorch/kernels. With the oss plan we need to build package for exir and this target needs to be part of the package. This diff moves the target into exir.

Reviewed By: JacobSzwejbka

Differential Revision: D47646479

fbshipit-source-id: 378a1751675218920ddb9d87eb481c8cba9665cc

Author: larryliu0820

exir/emit/TARGETS

@@ -32,7 +32,7 @@ python_library(
         "//executorch/exir/dialects/backend:lib",
         "//executorch/exir/dialects/edge:lib",
         "//executorch/exir/operator:convert",
+        "//executorch/exir/passes:prim_ops_py_registry",
         "//executorch/extension/pytree:pylib",
-        "//executorch/kernels/prim_ops:prim_to_executorch_ops",
     ],
 )

exir/emit/_emitter.py

@@ -40,6 +40,7 @@
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.error import ExportError, ExportErrorType, InternalError
 from executorch.exir.operator.convert import is_out_variant
+from executorch.exir.passes.executorch_prim_ops_registry import is_sym_op
 from executorch.exir.print_program import pretty_print_stacktraces
 from executorch.exir.schema import (
     BackendDelegate,
@@ -83,12 +84,10 @@
     TensorSpec,
 )
 from executorch.exir.types import LeafValueSpec, ValueSpec
-from executorch.kernels.prim_ops.prim_to_executorch_ops import is_sym_op
 from functorch.experimental import control_flow
 from torch._export.exported_program import ExportedProgram
 from torch.utils import _pytree as pytree
 
-# @manual=fbsource//third-party/pypi/typing-extensions:typing-extensions
 from typing_extensions import TypeAlias
 
 

exir/passes/TARGETS

@@ -13,6 +13,7 @@ python_library(
         ":memory_planning_pass",
         ":normalize_transpose_pass",
         ":pass_registry",
+        ":prim_ops_py_registry",
         ":quant_fusion_pass",
         ":remove_mixed_type_operators",
         ":remove_noop_pass",
@@ -36,7 +37,6 @@ python_library(
         "//executorch/exir/dialects/backend:lib",
         "//executorch/exir/dialects/edge:lib",
         "//executorch/exir/operator:convert",
-        "//executorch/kernels/prim_ops:prim_to_executorch_ops",
     ],
 )
 
@@ -203,10 +203,10 @@ python_library(
         "replace_edge_with_backend_pass.py",
     ],
     deps = [
+        ":prim_ops_py_registry",
         "//caffe2:torch",
         "//executorch/exir:pass_base",
         "//executorch/exir/dialects:lib",
-        "//executorch/kernels/prim_ops:prim_to_executorch_ops",
     ],
 )
 
@@ -227,11 +227,11 @@ python_library(
         "replace_aten_with_edge_pass.py",
     ],
     deps = [
+        ":prim_ops_py_registry",
         "//caffe2:torch",
         "//executorch/exir:pass_base",
         "//executorch/exir/dialects:lib",
         "//executorch/exir/dialects/edge:lib",
-        "//executorch/kernels/prim_ops:prim_to_executorch_ops",
     ],
 )
 
@@ -245,3 +245,12 @@ python_library(
         "//executorch/exir:pass_base",
     ],
 )
+
+python_library(
+    name = "prim_ops_py_registry",
+    srcs = ["executorch_prim_ops_registry.py"],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/exir/dialects:lib",
+    ],
+)

exir/passes/__init__.py

@@ -27,6 +27,8 @@
 from executorch.exir.pass_manager import PassManager, PassType
 from executorch.exir.passes.const_prop_pass import ConstPropPass
 from executorch.exir.passes.debug_handle_generator_pass import DebugHandleGeneratorPass
+
+from executorch.exir.passes.executorch_prim_ops_registry import _EXECUTORCH_SYM_OPS
 from executorch.exir.passes.memory_planning_pass import MemoryPlanningPass
 from executorch.exir.passes.normalize_transpose_pass import NormalizeTransposePass
 from executorch.exir.passes.pass_registry import PassRegistry
@@ -42,8 +44,6 @@
 from executorch.exir.passes.scalar_to_tensor_pass import ScalarToTensorPass
 from executorch.exir.passes.spec_prop_pass import SpecPropPass
 from executorch.exir.passes.sym_shape_eval_pass import SymShapeEvalPass
-
-from executorch.kernels.prim_ops.prim_to_executorch_ops import _EXECUTORCH_SYM_OPS
 from torch import fx
 from torch._subclasses import FakeTensor
 from torch.fx.passes.infra.pass_base import PassBase, PassResult

exir/passes/dynamic_shape_prop_pass.py

@@ -14,10 +14,10 @@
 from executorch.exir.pass_base import Argument, BackendPass
 from executorch.exir.pass_infra.node_metadata import NodeMetadata
 from executorch.exir.pass_infra.proxy_value import ProxyValue
+from executorch.exir.passes.executorch_prim_ops_registry import _EXECUTORCH_SYM_OPS
 from executorch.exir.schema import TensorShapeDynamism
 from executorch.exir.sym_util import collect_free_symbols, eval_expr
 from executorch.exir.tensor import TensorSpec
-from executorch.kernels.prim_ops.prim_to_executorch_ops import _EXECUTORCH_SYM_OPS
 from torch._subclasses import FakeTensor
 from torch.fx import GraphModule
 

kernels/prim_ops/executorch_prim_ops_registry.py renamed to exir/passes/executorch_prim_ops_registry.py

@@ -1,7 +1,14 @@
-from executorch.exir.dialects._ops import bind_pattern_to_op
+import operator
+from typing import Dict, Set
+
+# necessary to ensure the ops are registered
+import torch
+from executorch.exir.dialects._ops import bind_pattern_to_op, ops
 from torch import SymInt
+from torch._ops import OpOverload
 from torch.library import Library
 
+
 executorch_prims_lib = Library("executorch_prim", "DEF")
 
 
@@ -50,3 +57,35 @@ def le(a: SymInt, b: SymInt) -> bool:
 @bind_pattern_to_op(executorch_prims_lib, "eq.int(SymInt a, SymInt b) -> bool")
 def eq(a: SymInt, b: SymInt) -> bool:
     return a == b
+
+
+_PYTHON_SYM_OPS_TO_EXECUTORCH_SYM_OPS: Dict[OpOverload, OpOverload] = {
+    operator.sub: ops.backend.executorch_prim.sub.int,
+    operator.mul: ops.backend.executorch_prim.mul.int,
+    operator.add: ops.backend.executorch_prim.add.int,
+    operator.floordiv: ops.backend.executorch_prim.floordiv.int,
+    operator.eq: ops.backend.executorch_prim.eq.int,
+    operator.gt: ops.backend.executorch_prim.gt.int,
+    operator.lt: ops.backend.executorch_prim.lt.int,
+    operator.ge: ops.backend.executorch_prim.ge.int,
+    operator.le: ops.backend.executorch_prim.le.int,
+}
+
+
+_EXECUTORCH_SYM_OPS: Set[OpOverload] = set(
+    _PYTHON_SYM_OPS_TO_EXECUTORCH_SYM_OPS.values()
+)
+_EXECUTORCH_SYM_OPS.update(
+    {
+        torch.ops.aten.sym_stride.int,
+        torch.ops.aten.sym_size.int,
+        torch.ops.aten.sym_numel.default,
+    }
+)
+
+
+def is_sym_op(target) -> bool:
+    return (
+        target in _PYTHON_SYM_OPS_TO_EXECUTORCH_SYM_OPS.keys()
+        or target in _EXECUTORCH_SYM_OPS
+    )

exir/passes/replace_aten_with_edge_pass.py

@@ -4,7 +4,7 @@
 from executorch.exir.dialects._ops import ops
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.pass_base import ExportPass
-from executorch.kernels.prim_ops.prim_to_executorch_ops import _EXECUTORCH_SYM_OPS
+from executorch.exir.passes.executorch_prim_ops_registry import _EXECUTORCH_SYM_OPS
 from torch.fx.node import Target
 
 

exir/passes/replace_edge_with_backend_pass.py

@@ -3,7 +3,7 @@
 import torch
 from executorch.exir.dialects._ops import ops
 from executorch.exir.pass_base import ExportPass
-from executorch.kernels.prim_ops.prim_to_executorch_ops import (
+from executorch.exir.passes.executorch_prim_ops_registry import (
     _PYTHON_SYM_OPS_TO_EXECUTORCH_SYM_OPS,
 )
 

kernels/prim_ops/TARGETS

@@ -1,27 +1,7 @@
 # Any targets that should be shared between fbcode and xplat must be defined in
 # targets.bzl. This file can contain fbcode-only targets.
-load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
 load(":targets.bzl", "define_common_targets")
 
 oncall("executorch")
 
 define_common_targets()
-
-python_library(
-    name = "prim_ops_py_registry",
-    srcs = ["executorch_prim_ops_registry.py"],
-    deps = [
-        "//caffe2:torch",
-        "//executorch/exir/dialects:lib",
-    ],
-)
-
-python_library(
-    name = "prim_to_executorch_ops",
-    srcs = ["prim_to_executorch_ops.py"],
-    deps = [
-        ":prim_ops_py_registry",
-        "//caffe2:torch",
-        "//executorch/exir/dialects:lib",
-    ],
-)

kernels/prim_ops/prim_to_executorch_ops.py

@@ -1,12 +1,11 @@
+load("@fbcode_macros//build_defs:cpp_unittest.bzl", "cpp_unittest")
+
 # Any targets that should be shared between fbcode and xplat must be defined in
 # targets.bzl. This file can contain fbcode-only targets.
 load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest")
-load(":targets.bzl", "define_common_targets")
 
 oncall("executorch")
 
-define_common_targets()
-
 python_unittest(
     name = "test_prim_ops",
     srcs = [
@@ -15,6 +14,24 @@ python_unittest(
     supports_static_listing = True,
     deps = [
         "//caffe2:torch",
-        "//executorch/kernels/prim_ops:prim_ops_py_registry",
+        "//executorch/exir/passes:prim_ops_py_registry",
+    ],
+)
+
+cpp_unittest(
+    name = "register_prim_ops_test",
+    srcs = [
+        "register_prim_ops_test.cpp",
+    ],
+    deps = [
+        "//executorch/kernels/prim_ops:prim_ops_registry",  # @manual
+        "//executorch/runtime/core:evalue",  # @manual
+        "//executorch/runtime/core/exec_aten:lib",  # @manual
+        "//executorch/runtime/core/exec_aten/testing_util:tensor_util",  # @manual
+        "//executorch/runtime/core/exec_aten/util:tensor_util",  # @manual
+        "//executorch/runtime/kernel:kernel_runtime_context",  # @manual
+        "//executorch/runtime/kernel:operator_registry",
+        "//executorch/runtime/platform:platform",
+        "//executorch/test/utils:utils",
     ],
 )

kernels/prim_ops/test/TARGETS

@@ -1,7 +1,7 @@
 import unittest
 
 # necessary to ensure the ops are registered
-import executorch.kernels.prim_ops.executorch_prim_ops_registry
+import executorch.exir.passes.executorch_prim_ops_registry  # noqa: F401
 
 import torch