Use relative import for build from source case

Author: larryliu0820

codegen/api/et_cpp.py

@@ -2,15 +2,6 @@
 
 from typing import TYPE_CHECKING
 
-from executorch.codegen.api.types import (
-    ArrayRefCType,
-    BaseTypeToCppMapping,
-    OptionalCType,
-    scalarT,
-    tensorListT,
-    tensorT,
-)
-
 from torchgen import local
 from torchgen.api.types import (
     ArgName,
@@ -40,6 +31,15 @@
 )
 from typing_extensions import assert_never
 
+from .types import (
+    ArrayRefCType,
+    BaseTypeToCppMapping,
+    OptionalCType,
+    scalarT,
+    tensorListT,
+    tensorT,
+)
+
 
 if TYPE_CHECKING:
     from collections.abc import Sequence

codegen/api/types/__init__.py

@@ -1,5 +1,5 @@
 # flake8: noqa: F403, F401
-from executorch.codegen.api.types.types import *
+from .types import *
 
 # flake8: noqa: F403, F401
-from executorch.codegen.api.types.signatures import *  # usort: skip
+from .signatures import *  # usort: skip

codegen/api/types/signatures.py

@@ -4,7 +4,8 @@
 from typing import TYPE_CHECKING
 
 import torchgen.api.cpp as aten_cpp
-from executorch.codegen.api.types.types import contextArg
+
+from .types import contextArg
 
 
 if TYPE_CHECKING:
@@ -73,4 +74,4 @@ def from_native_function(
         )
 
 
-from executorch.codegen.api import et_cpp
+from .. import et_cpp

codegen/gen.py

@@ -8,15 +8,32 @@
 from typing import Any, Callable, TextIO, TYPE_CHECKING
 
 import yaml
-from executorch.codegen.api import et_cpp
-from executorch.codegen.api.custom_ops import (
-    ComputeNativeFunctionStub,
-    gen_custom_ops_registration,
-)
-from executorch.codegen.api.types import contextArg, ExecutorchCppSignature
-from executorch.codegen.api.unboxing import Unboxing
-from executorch.codegen.model import ETKernelIndex, ETKernelKey, ETParsedYaml
-from executorch.codegen.parse import ET_FIELDS, parse_et_yaml, parse_et_yaml_struct
+
+try:
+    from executorch.codegen.api import et_cpp
+    from executorch.codegen.api.custom_ops import (
+        ComputeNativeFunctionStub,
+        gen_custom_ops_registration,
+    )
+    from executorch.codegen.api.types import contextArg, ExecutorchCppSignature
+    from executorch.codegen.api.unboxing import Unboxing
+    from executorch.codegen.model import ETKernelIndex, ETKernelKey, ETParsedYaml
+    from executorch.codegen.parse import ET_FIELDS, parse_et_yaml, parse_et_yaml_struct
+except ImportError:
+    # If we build from source, executorch.codegen is not available.
+    from .api import et_cpp  # type: ignore[no-redef]
+    from .api.custom_ops import (  # type: ignore
+        ComputeNativeFunctionStub,
+        gen_custom_ops_registration,
+    )
+    from .api.types import contextArg, ExecutorchCppSignature  # type: ignore
+    from .api.unboxing import Unboxing  # type: ignore
+    from .model import ETKernelIndex, ETKernelKey, ETParsedYaml  # type: ignore
+    from .parse import (  # type: ignore[no-redef]
+        ET_FIELDS,
+        parse_et_yaml,
+        parse_et_yaml_struct,
+    )
 
 # Parse native_functions.yaml into a sequence of NativeFunctions and Backend Indices.
 from torchgen import dest

codegen/tools/gen_oplist.py

@@ -17,8 +17,16 @@
     from executorch.codegen.parse import strip_et_fields
 except ImportError:
     # If we build from source, executorch.codegen is not available.
-    # We can use relative import instead.
-    from ..parse import strip_et_fields
+    ET_FIELDS = ["kernels", "type_alias", "dim_order_alias"]
+
+    def strip_et_fields(es: object) -> None:
+        """Given a loaded yaml representing a list of operators,
+        remove ET specific fields from every entries for BC compatibility
+        """
+        for entry in es:  # type: ignore[attr-defined]
+            for field in ET_FIELDS:
+                entry.pop(field, None)
+
 
 from torchgen.gen import LineLoader, parse_native_yaml_struct
 from torchgen.selective_build.operator import SelectiveBuildOperator

tools/cmake/Codegen.cmake

@@ -93,7 +93,7 @@ function(generate_bindings_for_kernels)
   file(GLOB_RECURSE _torchgen_srcs "${torchgen-out}/*.py")
   # Not using module executorch.codegen.gen because it's not installed yet.
   set(_gen_command
-      "${PYTHON_EXECUTABLE}" ${EXECUTORCH_ROOT}/codegen/gen.py
+      "${PYTHON_EXECUTABLE}" -m codegen.gen
       --source-path=${EXECUTORCH_ROOT}/codegen --install-dir=${_out_dir}
       --tags-path=${torchgen-out}/packaged/ATen/native/tags.yaml
       --aten-yaml-path=${torchgen-out}/packaged/ATen/native/native_functions.yaml