Refactor aot_util build (#2631)

Summary:
Refactor aot_util build to be independent of runtime build, clean up build scripts. The intent of this change is to make the aot_util build independent of the runtime build, as well as to clean up the build scripts and setup.py logic. This change also makes aot_util build by default from setup.py, whereas it previously was tied to the EXECUTORCH_BUILD_PYBIND flag (which requires a full runtime build).

Bypassing known failures in external CI.
bypass-github-export-checks
bypass-github-executorch-ci-checks

Pull Request resolved: #2631

Test Plan:
pip uninstall executorch
BUCK=buck2-et2 ./install_requirements.sh

Confirmed that libaot_util.so is present under site-packages/executorch/extension/aot_util/.

Added logging when loading aot_util lib to log on successful invocation. Exported Stories110M and confirmed aotlib was loaded successfully.

python -m examples.models.llama2.export_llama -c /mnt/c/Users/gregory/Downloads/Stories110M.pt -p /mnt/c/Users/gregory/Downloads/Stories110M_params.json -kv -qmode 8da4w -X -d fp32

Reviewed By: mcr229

Differential Revision: D55292709

Pulled By: GregoryComer

fbshipit-source-id: 7f38cdfb0bcb326897a67c1fe99557ca3dc58a77

Author: GregoryComer

CMakeLists.txt

@@ -140,7 +140,7 @@ option(EXECUTORCH_BUILD_ARM_BAREMETAL
 
 option(EXECUTORCH_BUILD_COREML "Build the Core ML backend" OFF)
 
-option(EXECUTORCH_BUILD_EXTENSION_AOT_UTIL "Build the AOT Util extension" OFF)
+option(EXECUTORCH_BUILD_EXTENSION_AOT_UTIL "Build the AOT util library" OFF)
 
 option(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER "Build the Data Loader extension"
        OFF)
@@ -362,6 +362,10 @@ if(EXECUTORCH_BUILD_EXECUTOR_RUNNER)
   target_compile_options(executor_runner PUBLIC ${_common_compile_options})
 endif()
 
+if(EXECUTORCH_BUILD_EXTENSION_AOT_UTIL)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/aot_util)
+ endif()
+
 # Add googletest if any test targets should be built
 if(EXECUTORCH_BUILD_GTESTS)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third-party/googletest)
@@ -390,10 +394,6 @@ if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/runner_util)
 endif()
 
-if(EXECUTORCH_BUILD_EXTENSION_AOT_UTIL)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/aot_util)
-endif()
-
 if(EXECUTORCH_BUILD_XNNPACK)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/xnnpack)
 endif()

backends/arm/third-party/serialization_lib

@@ -451,8 +451,8 @@ def define_tensor(
 
     @staticmethod
     def find_aot_util_path() -> str:
-        # Look for .so installed by wheel (OSS).
-        rel_path = "executorch/extension/pybindings/libaot_util.so"
+        # Look for .so installed by wheel (OSS). TODO(gjcomer) Improve this.
+        rel_path = "executorch/extension/aot_util/libaot_util.so"
         for sys_path in sys.path:
             so_path = Path(sys_path) / rel_path
             if so_path.exists():

backends/xnnpack/operators/node_visitor.py

@@ -88,6 +88,16 @@ filters = [
 # ---------------------------------- core end ----------------------------------
 # ---------------------------------- extension start ----------------------------------
 
+[targets.extension_aot_util]
+buck_targets = [
+  "//extension/aot_util:aot_util",
+]
+filters = [
+  ".cpp$",
+]
+deps = [
+  "executorch",
+]
 [targets.extension_data_loader]
 buck_targets = [
   "//extension/data_loader:buffer_data_loader",
@@ -125,17 +135,6 @@ deps = [
   "executorch",
 ]
 
-[targets.extension_aot_util]
-buck_targets = [
-  "//extension/aot_util:aot_util",
-]
-filters = [
-  ".cpp$",
-]
-deps = [
-  "executorch",
-]
-
 # ---------------------------------- extension end ----------------------------------
 # ---------------------------------- binary start ----------------------------------
 

build/cmake_deps.toml

@@ -9,16 +9,60 @@
 # cmake-format --first-comment-is-literal=True CMakeLists.txt
 # ~~~
 
+cmake_minimum_required(VERSION 3.19)
+project(aot_util)
+include(../../build/Utils.cmake)
+
+if(NOT CMAKE_CXX_STANDARD)
+  set(CMAKE_CXX_STANDARD 17)
+endif()
+
+if(NOT CMAKE_BUILD_TYPE)
+  set(CMAKE_BUILD_TYPE Debug)
+endif()
+
+if(NOT EXECUTORCH_ROOT)
+    set(EXECUTORCH_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/../../")
+endif()
+
+if(NOT BUCK2)
+  set(BUCK2 buck2)
+endif()
+
+if(NOT PYTHON_EXECUTABLE)
+  set(PYTHON_EXECUTABLE python3)
+endif()
+
+if(NOT EXECUTORCH_SRCS_FILE)
+  # A file wasn't provided. Run a script to extract the source lists from the
+  # buck2 build system and write them to a file we can include.
+  #
+  # NOTE: This will only happen once during cmake setup, so it will not re-run
+  # if the buck2 targets change.
+  message(STATUS "executorch: Generating source lists")
+  set(EXECUTORCH_SRCS_FILE "${CMAKE_CURRENT_BINARY_DIR}/executorch_srcs.cmake")
+  extract_sources(${EXECUTORCH_SRCS_FILE})
+endif()
+
+# This file defines the `_<target>__srcs` variables used below.
+message(STATUS "executorch: Using sources file ${EXECUTORCH_SRCS_FILE}")
+include(${EXECUTORCH_SRCS_FILE})
+
+
+# Ahead-of-time (AOT) utility library. Contains native code used by the
+# AOT lowering and delegation logic. Note that this library should build
+# independently of the runtime code, and as such, should not have
+# dependencies on runtime targets.
 find_package(Torch CONFIG REQUIRED)
 find_library(TORCH_PYTHON_LIBRARY torch_python
                PATHS "${TORCH_INSTALL_PREFIX}/lib")
 
-# Override compiler flags set in upper scope. ExecuTorch builds with
-# -fno-exceptions and -fno-rtti, but we need these for ATen.
-set(CMAKE_CXX_FLAGS_RELEASE "")
+# Override compiler flags set in upper scope when included from the top-level
+# CMakeLists. ExecuTorch builds with -fno-exceptions and -fno-rtti, but we 
+# need these for ATen.
+unset(CMAKE_CXX_FLAGS_RELEASE)
 
-#list(TRANSFORM _xnnpack_convert_to_qc4w__srcs PREPEND "${EXECUTORCH_ROOT}/")
-# add_library(aot_util ${_xnnpack_convert_to_qc4w__srcs})
-add_library(aot_util "convert_to_qc4w.cpp")
+list(TRANSFORM _extension_aot_util__srcs PREPEND "${EXECUTORCH_ROOT}/")
+add_library(aot_util ${_extension_aot_util__srcs})
 target_include_directories(aot_util PUBLIC ${TORCH_INCLUDE_DIRS})
 target_link_libraries(aot_util torch)

extension/aot_util/CMakeLists.txt

@@ -0,0 +1,9 @@
+# AOT Util
+
+Ahead-of-time (AOT) utility library. Contains native code used by the AOT lowering and delegation logic. Note 
+that this library should build independently of the runtime code, and as such, should not have dependencies 
+on runtime targets.
+
+This library is intended to be built and distributed as part of the Python pip package, such that it can be
+loaded by AOT Python code.
+

extension/aot_util/README.md

@@ -96,11 +96,10 @@ def build_extension(self, ext: CMakeExtension) -> None:
         # Set Python_EXECUTABLE instead if you use PYBIND11_FINDPYTHON
         # EXAMPLE_VERSION_INFO shows you how to pass a value into the C++ code
         # from Python.
-        buck = os.environ.get("BUCK", "buck2")
+        buck = os.environ.get("BUCK", "")
         cmake_prefix_path = os.environ.get("CMAKE_PREFIX_PATH", get_python_lib())
         cmake_args = [
             "-DEXECUTORCH_BUILD_PYBIND=ON",
-            "-DEXECUTORCH_BUILD_EXTENSION_AOT_UTIL=ON",
             "-DBUILD_SHARED_LIBS=ON",  # For flatcc
             f"-DBUCK2={buck}",
             f"-DCMAKE_PREFIX_PATH={cmake_prefix_path}",
@@ -225,14 +224,20 @@ def run(self):
 
 
 cmdclass = {
+    "build_ext": CMakeBuild,
     "install": CustomInstallCommand,
     "develop": CustomDevelopCommand,
     "egg_info": CustomEggInfoCommand,
 }
-ext_modules = None
+ext_modules = []
+
+if os.environ.get("EXECUTORCH_BUILD_AOT_UTIL", "ON") == "ON":
+    ext_modules.append(
+        CMakeExtension("executorch.extension.aot_util.aot_util", "extension/aot_util")
+    )
+
 if os.environ.get("EXECUTORCH_BUILD_PYBIND", "OFF") == "ON":
-    cmdclass["build_ext"] = CMakeBuild
-    ext_modules = [CMakeExtension("executorch.extension.pybindings.portable_lib")]
+    ext_modules.append(CMakeExtension("executorch.extension.pybindings.portable_lib"))
 
 setup(
     package_dir={