move prim ops into kernels/

Summary:
https://docs.google.com/spreadsheets/d/17mEkBHN0XwXN21pG2IxpANX-jd4hNMrX4x9FthYB5Ok/edit#gid=0

this may not be the final spot these ops end up. They are in a weird spot where they are essentially custom ops that get added by the emitter to resolve specific issues around control flow and symbolic shapes

Reviewed By: mergennachin

Differential Revision: D47378903

fbshipit-source-id: f4bf23817438727b62aabf852615c37976a458a5

Author: JacobSzwejbka

codegen/codegen.bzl

@@ -375,7 +375,7 @@ def executorch_generated_lib(
             compiler_flags = ["-Wno-global-constructors"] + compiler_flags,
             deps = [
                 "//executorch/runtime/kernel:operator_registry",
-                "//executorch/core/prim_ops:prim_ops_registry" + aten_suffix,
+                "//executorch/kernels/prim_ops:prim_ops_registry" + aten_suffix,
                 "//executorch/runtime/core:evalue" + aten_suffix,
                 "//executorch/profiler:profiler",
                 "//executorch/codegen:macros",

core/kernel_types/testing/targets.bzl

@@ -26,7 +26,7 @@ def define_common_targets():
                 # list.
                 "//executorch/core/kernel_types/util/test/...",
                 "//executorch/core/values/test/...",
-                "//executorch/core/prim_ops/test/...",
+                "//executorch/kernels/prim_ops/test/...",
                 "//executorch/kernels/portable/test/...",
                 "//executorch/kernels/portable/cpu/util/test/...",
                 "//executorch/kernels/quantized/test/...",

exir/emit/TARGETS

@@ -20,7 +20,6 @@ python_library(
         "fbsource//third-party/pypi/typing-extensions:typing-extensions",
         "//caffe2:torch",
         "//caffe2/functorch:functorch_src",
-        "//executorch/core/prim_ops:prim_to_executorch_ops",
         "//executorch/exir:common",
         "//executorch/exir:delegate",
         "//executorch/exir:error",
@@ -34,6 +33,7 @@ python_library(
         "//executorch/exir/dialects/backend:lib",
         "//executorch/exir/dialects/edge:lib",
         "//executorch/exir/operator:convert",
+        "//executorch/kernels/prim_ops:prim_to_executorch_ops",
         "//executorch/pytree:pylib",
     ],
 )

exir/emit/_emitter.py

@@ -33,7 +33,6 @@
 import executorch.pytree as ex_pytree
 import torch
 import torch.fx
-from executorch.core.prim_ops.prim_to_executorch_ops import is_sym_op
 from executorch.exir import delegate
 from executorch.exir.common import add_cursor_to_graph
 from executorch.exir.delegate import LoweredBackendModule
@@ -85,6 +84,7 @@
     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.utils import _pytree as pytree
 

exir/passes/TARGETS

@@ -24,7 +24,6 @@ python_library(
         ":spec_prop_pass",
         ":sym_shape_eval_pass",
         "//caffe2:torch",
-        "//executorch/core/prim_ops:prim_to_executorch_ops",
         "//executorch/exir:common",
         "//executorch/exir:control_flow",
         "//executorch/exir:delegate",
@@ -37,6 +36,7 @@ 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,9 +203,9 @@ python_library(
     ],
     deps = [
         "//caffe2:torch",
-        "//executorch/core/prim_ops:prim_to_executorch_ops",
         "//executorch/exir:pass_base",
         "//executorch/exir/dialects:lib",
+        "//executorch/kernels/prim_ops:prim_to_executorch_ops",
     ],
 )
 
@@ -227,10 +227,10 @@ python_library(
     ],
     deps = [
         "//caffe2:torch",
-        "//executorch/core/prim_ops:prim_to_executorch_ops",
         "//executorch/exir:pass_base",
         "//executorch/exir/dialects:lib",
         "//executorch/exir/dialects/edge:lib",
+        "//executorch/kernels/prim_ops:prim_to_executorch_ops",
     ],
 )
 

exir/passes/__init__.py

@@ -10,8 +10,6 @@
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
 import torch
-
-from executorch.core.prim_ops.prim_to_executorch_ops import _EXECUTORCH_SYM_OPS
 from executorch.exir import control_flow, delegate, memory, memory_planning
 from executorch.exir.common import override_logger
 from executorch.exir.dialects.backend._ops import BackendOpOverload
@@ -44,6 +42,8 @@
 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

@@ -6,7 +6,6 @@
 
 import torch
 import torch.utils._pytree as pytree
-from executorch.core.prim_ops.prim_to_executorch_ops import _EXECUTORCH_SYM_OPS
 from executorch.exir.delegate import LoweredBackendModule
 from executorch.exir.dynamic_shape import (
     calculate_dynamic_shape_spec,
@@ -18,6 +17,7 @@
 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
 

exir/passes/replace_aten_with_edge_pass.py

@@ -1,10 +1,10 @@
 # (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
 
 import torch
-from executorch.core.prim_ops.prim_to_executorch_ops import _EXECUTORCH_SYM_OPS
 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 torch.fx.node import Target
 
 

exir/passes/replace_edge_with_backend_pass.py

@@ -1,11 +1,11 @@
 # pyre-strict
 
 import torch
-from executorch.core.prim_ops.prim_to_executorch_ops import (
-    _PYTHON_SYM_OPS_TO_EXECUTORCH_SYM_OPS,
-)
 from executorch.exir.dialects._ops import ops
 from executorch.exir.pass_base import ExportPass
+from executorch.kernels.prim_ops.prim_to_executorch_ops import (
+    _PYTHON_SYM_OPS_TO_EXECUTORCH_SYM_OPS,
+)
 
 
 class EdgeToBackendOpsPass(ExportPass):

core/prim_ops/TARGETS renamed to kernels/prim_ops/TARGETS

@@ -1,6 +1,6 @@
 // (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
 
-#include <executorch/core/prim_ops/et_copy_index.h>
+#include <executorch/kernels/prim_ops/et_copy_index.h>
 
 #include <cstring>
 

core/prim_ops/et_copy_index.cpp renamed to kernels/prim_ops/et_copy_index.cpp

@@ -2,7 +2,7 @@
 from typing import Dict, Set
 
 # necessary to ensure the ops are registered
-import executorch.core.prim_ops.executorch_prim_ops_registry  # noqa
+import executorch.kernels.prim_ops.executorch_prim_ops_registry  # noqa
 import torch
 from executorch.exir.dialects._ops import ops
 from torch._ops import OpOverload

core/prim_ops/et_copy_index.h renamed to kernels/prim_ops/et_copy_index.h

@@ -1,4 +1,4 @@
-#include <executorch/core/prim_ops/et_copy_index.h>
+#include <executorch/kernels/prim_ops/et_copy_index.h>
 #include <executorch/runtime/core/evalue.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 #include <executorch/runtime/kernel/operator_registry.h>

core/prim_ops/executorch_prim_ops_registry.py renamed to kernels/prim_ops/executorch_prim_ops_registry.py

@@ -26,7 +26,7 @@ def define_common_targets():
 
         runtime.cxx_library(
             name = "prim_ops_registry" + aten_suffix,
-            srcs = ["RegisterPrimOps.cpp"],
+            srcs = ["register_prim_ops.cpp"],
             visibility = [
                 "//executorch/...",
                 "@EXECUTORCH_CLIENTS",

core/prim_ops/prim_to_executorch_ops.py renamed to kernels/prim_ops/prim_to_executorch_ops.py

@@ -15,6 +15,6 @@ python_unittest(
     supports_static_listing = True,
     deps = [
         "//caffe2:torch",
-        "//executorch/core/prim_ops:prim_ops_py_registry",
+        "//executorch/kernels/prim_ops:prim_ops_py_registry",
     ],
 )

core/prim_ops/RegisterPrimOps.cpp renamed to kernels/prim_ops/register_prim_ops.cpp

@@ -9,11 +9,11 @@ def define_common_targets():
     runtime.cxx_test(
         name = "register_prim_ops_test",
         srcs = [
-            "RegisterPrimOpsTest.cpp",
+            "register_prim_ops_test.cpp",
         ],
         deps = [
             "//executorch/core/kernel_types/testing:tensor_util",
-            "//executorch/core/prim_ops:prim_ops_registry",
+            "//executorch/kernels/prim_ops:prim_ops_registry",
             "//executorch/runtime/kernel:operator_registry",
             "//executorch/runtime/kernel:kernel_runtime_context",
         ],

core/prim_ops/targets.bzl renamed to kernels/prim_ops/targets.bzl

@@ -1,7 +1,7 @@
 import unittest
 
 # necessary to ensure the ops are registered
-import executorch.core.prim_ops.executorch_prim_ops_registry
+import executorch.kernels.prim_ops.executorch_prim_ops_registry
 
 import torch
 

core/prim_ops/test/TARGETS renamed to kernels/prim_ops/test/TARGETS

@@ -70,7 +70,7 @@ def define_common_targets():
             deps = [
                 ":memory_manager",
                 "//executorch/runtime/backend:backend_registry",
-                "//executorch/core/prim_ops:prim_ops_registry" + aten_suffix,
+                "//executorch/kernels/prim_ops:prim_ops_registry" + aten_suffix,
                 "//executorch/runtime/kernel:kernel_runtime_context" + aten_suffix,
                 "//executorch/profiler:profiler",
                 "//executorch/schema:schema",

core/prim_ops/test/RegisterPrimOpsTest.cpp renamed to kernels/prim_ops/test/register_prim_ops_test.cpp

@@ -30,7 +30,7 @@ def define_common_targets():
                 "kernel_runtime_context.h",
             ],
             visibility = [
-                "//executorch/core/prim_ops/...",  # Contains kernels
+                "//executorch/kernels/prim_ops/...",  # Contains kernels
                 "//executorch/runtime/kernel/...",
                 "//executorch/kernels/...",
                 "//executorch/runtime/executor/...",
@@ -50,7 +50,7 @@ def define_common_targets():
             visibility = [
                 "//executorch/runtime/kernel/...",
                 "//executorch/kernels/...",
-                "//executorch/core/prim_ops/...",  # Prim kernels
+                "//executorch/kernels/prim_ops/...",  # Prim kernels
                 "@EXECUTORCH_CLIENTS",
             ],
             exported_deps = [