[Python] Enable building Python bindings as editable wheels, document it (iree-org#19716)

In order to not need to constantly source PYTHONPATH, and to run the
rink of potentially having it set wrong when dealing with multiple IREE
builds, and to allow packages like the kernel bunchmarking suite to use
a local build without needing to edit requirements.txt, add the ability
to build these packages as editableble wheels.

This method, newly added to the build documentation, tells CMake to use
symbolic links when "installing" the Python packages from the build
directroy into a different build directory. In combination with telling
copytree to preserve symlinks, this creates Python packages that link
back to the build or source directory when the `-e` flag is used on pip.

This means that an automated virtual environment switcher, like `pyenv`,
will pick up the correct Python bindings automatically and will direct
`iree-compiler` and `iree-runtime` shims to the correct build.

Author: krzysz00

compiler/setup.py

@@ -21,6 +21,10 @@
 #
 # Select CMake options are available from environment variables:
 #   IREE_ENABLE_CPUINFO
+#
+# If building from a development tree and aiming to get an "editable" install,
+# use the environment option CMAKE_INSTALL_MODE=ABS_SYMLINK on your
+# `pip install -e .` invocation.
 
 from gettext import install
 import json
@@ -337,7 +341,7 @@ def run(self):
         shutil.copytree(
             os.path.join(CMAKE_INSTALL_DIR_ABS, "python_packages", "iree_compiler"),
             target_dir,
-            symlinks=False,
+            symlinks=self.editable_mode,
         )
         print("Target populated.", file=sys.stderr)
 

docs/website/docs/building-from-source/getting-started.md

@@ -396,6 +396,43 @@ cmake --build ../iree-build/
 
 ### Using the Python bindings
 
+There are two available methods for installing the Python bindings, either
+through creating an editable wheel or through extending `PYTHONPATH`.
+
+#### Option A: Installing the bindings as editable wheels
+
+This method links the files in your build tree into your Python package directory
+as an editable wheel.
+
+=== ":fontawesome-brands-linux: Linux"
+
+    ``` shell
+    CMAKE_INSTALL_METHOD=ABS_SYMLINK python -m pip install -e ../iree-build/compiler
+    CMAKE_INSTALL_METHOD=ABS_SYMLINK python -m pip install -e ../iree-build/runtime
+    ```
+
+=== ":fontawesome-brands-apple: macOS"
+
+    ``` shell
+    CMAKE_INSTALL_METHOD=ABS_SYMLINK python -m pip install -e ../iree-build/compiler
+    CMAKE_INSTALL_METHOD=ABS_SYMLINK python -m pip install -e ../iree-build/runtime
+    ```
+
+=== ":fontawesome-brands-windows: Windows"
+
+    ``` powershell
+    $env:CMAKE_INSTALL_MODE="ABS_SYMLINK"
+    python -m pip install -e ..\iree-build\compiler
+    python -m pip install -e ..\iree-build\runtime
+    $env:CMAKE_INSTALL_MODE=null
+    ```
+
+#### Option B: Extending PYTHONPATH
+
+This method more effectively captures the state of your build directory,
+but is prone to errors arising from forgetting to source the environment
+variables.
+
 Extend your `PYTHONPATH` with IREE's `bindings/python` paths and try importing:
 
 === ":fontawesome-brands-linux: Linux"
@@ -431,13 +468,16 @@ Extend your `PYTHONPATH` with IREE's `bindings/python` paths and try importing:
     python -c "import iree.runtime; help(iree.runtime)"
     ```
 
+#### Tensorflow/TFLite bindings
+
 Using IREE's TensorFlow/TFLite importers requires a few extra steps:
 
 ``` shell
 # Install test requirements
 python -m pip install -r integrations/tensorflow/test/requirements.txt
 
 # Install pure Python packages (no build required)
+# You may use `pip install -e` here to create an editable wheel.
 python -m pip install integrations/tensorflow/python_projects/iree_tf
 python -m pip install integrations/tensorflow/python_projects/iree_tflite
 

runtime/CMakeLists.txt

@@ -34,5 +34,7 @@ iree_include_cmake_plugin_dirs(
 add_subdirectory(src)
 
 if(IREE_BUILD_PYTHON_BINDINGS)
+  configure_file(pyproject.toml pyproject.toml COPYONLY)
+  configure_file(setup.py setup.py @ONLY)
   add_subdirectory(bindings/python)
 endif()

runtime/setup.py

@@ -23,6 +23,10 @@
 # Select CMake options are available from environment variables:
 #   IREE_HAL_DRIVER_VULKAN
 #   IREE_ENABLE_CPUINFO
+#
+# If building from a development tree and aiming to get an "editable" install,
+# use the environment option CMAKE_INSTALL_MODE=ABS_SYMLINK on your
+# `pip install -e .` invocation.
 
 import json
 import os
@@ -41,8 +45,11 @@
 from setuptools.command.build_py import build_py as _build_py
 
 
-def getenv_bool(key, default_value="OFF"):
-    value = os.getenv(key, default_value)
+def getenv_bool(key: str, cmake_arg: str, default_value="OFF"):
+    if cmake_arg == "" or cmake_arg[0] != "@":
+        value = cmake_arg
+    else:
+        value = os.getenv(key, default_value)
     return value.upper() in ["ON", "1", "TRUE"]
 
 
@@ -53,7 +60,17 @@ def combine_dicts(*ds):
     return result
 
 
-ENABLE_TRACY = getenv_bool("IREE_RUNTIME_BUILD_TRACY", "ON")
+# This file can be run directly from the source tree or it can be CMake
+# configured so it can run from the build tree with an already existing
+# build tree. We detect the difference based on whether the following
+# are expanded by CMake.
+CONFIGURED_SOURCE_DIR = "@IREE_SOURCE_DIR@"
+CONFIGURED_BINARY_DIR = "@IREE_BINARY_DIR@"
+
+ENABLE_TRACY = getenv_bool(
+    "IREE_RUNTIME_BUILD_TRACY", "@IREE_RUNTIME_BUILD_TRACY@", "ON"
+)
+
 if ENABLE_TRACY:
     print(
         "*** Enabling Tracy instrumented runtime (disable with IREE_RUNTIME_BUILD_TRACY=OFF)",
@@ -64,7 +81,9 @@ def combine_dicts(*ds):
         "*** Tracy instrumented runtime not enabled (enable with IREE_RUNTIME_BUILD_TRACY=ON)",
         file=sys.stderr,
     )
-ENABLE_TRACY_TOOLS = getenv_bool("IREE_RUNTIME_BUILD_TRACY_TOOLS")
+ENABLE_TRACY_TOOLS = getenv_bool(
+    "IREE_RUNTIME_BUILD_TRACY_TOOLS", "@IREE_RUNTIME_BUILD_TRACY_TOOLS@"
+)
 if ENABLE_TRACY_TOOLS:
     print("*** Enabling Tracy tools (may error if missing deps)", file=sys.stderr)
 else:
@@ -111,21 +130,33 @@ def check_pip_version():
 CMAKE_TRACY_INSTALL_DIR_REL = os.path.join("build", "i", "t")
 CMAKE_TRACY_INSTALL_DIR_ABS = os.path.join(SETUPPY_DIR, CMAKE_TRACY_INSTALL_DIR_REL)
 
-IREE_SOURCE_DIR = os.path.join(SETUPPY_DIR, "..")
-# Note that setuptools always builds into a "build" directory that
-# is a sibling of setup.py, so we just colonize a sub-directory of that
-# by default.
-BASE_BINARY_DIR = os.getenv(
-    "IREE_RUNTIME_API_CMAKE_BUILD_DIR", os.path.join(SETUPPY_DIR, "build", "b")
-)
-IREE_BINARY_DIR = os.path.join(BASE_BINARY_DIR, "d")
-IREE_TRACY_BINARY_DIR = os.path.join(BASE_BINARY_DIR, "t")
-print(
-    f"Running setup.py from source tree: "
-    f"SOURCE_DIR = {IREE_SOURCE_DIR} "
-    f"BINARY_DIR = {IREE_BINARY_DIR}",
-    file=sys.stderr,
-)
+IS_CONFIGURED = CONFIGURED_SOURCE_DIR[0] != "@"
+if IS_CONFIGURED:
+    IREE_SOURCE_DIR = CONFIGURED_SOURCE_DIR
+    IREE_BINARY_DIR = CONFIGURED_BINARY_DIR
+    IREE_TRACY_BINARY_DIR = CONFIGURED_BINARY_DIR
+    print(
+        f"Running setup.py from build tree: "
+        f"SOURCE_DIR = {IREE_SOURCE_DIR} "
+        f"BINARY_DIR = {IREE_BINARY_DIR}",
+        file=sys.stderr,
+    )
+else:
+    IREE_SOURCE_DIR = os.path.join(SETUPPY_DIR, "..")
+    # Note that setuptools always builds into a "build" directory that
+    # is a sibling of setup.py, so we just colonize a sub-directory of that
+    # by default.
+    BASE_BINARY_DIR = os.getenv(
+        "IREE_RUNTIME_API_CMAKE_BUILD_DIR", os.path.join(SETUPPY_DIR, "build", "b")
+    )
+    IREE_BINARY_DIR = os.path.join(BASE_BINARY_DIR, "d")
+    IREE_TRACY_BINARY_DIR = os.path.join(BASE_BINARY_DIR, "t")
+    print(
+        f"Running setup.py from source tree: "
+        f"SOURCE_DIR = {IREE_SOURCE_DIR} "
+        f"BINARY_DIR = {IREE_BINARY_DIR}",
+        file=sys.stderr,
+    )
 
 # Setup and get version information.
 VERSION_FILE = os.path.join(IREE_SOURCE_DIR, "runtime/version.json")
@@ -263,61 +294,62 @@ def build_configuration(cmake_build_dir, cmake_install_dir, extra_cmake_args=())
     cfg = os.getenv("IREE_CMAKE_BUILD_TYPE", "Release")
     strip_install = cfg == "Release"
 
-    # Build from source tree.
-    os.makedirs(cmake_build_dir, exist_ok=True)
-    maybe_nuke_cmake_cache(cmake_build_dir, cmake_install_dir)
-    print(f"CMake build dir: {cmake_build_dir}", file=sys.stderr)
-    print(f"CMake install dir: {cmake_install_dir}", file=sys.stderr)
-    cmake_args = [
-        "-GNinja",
-        "--log-level=VERBOSE",
-        f"-DIREE_RUNTIME_OPTIMIZATION_PROFILE={IREE_RUNTIME_OPTIMIZATION_PROFILE}",
-        "-DIREE_BUILD_PYTHON_BINDINGS=ON",
-        "-DIREE_BUILD_COMPILER=OFF",
-        "-DIREE_BUILD_SAMPLES=OFF",
-        "-DIREE_BUILD_TESTS=OFF",
-        "-DPython3_EXECUTABLE={}".format(sys.executable),
-        "-DCMAKE_BUILD_TYPE={}".format(cfg),
-        get_env_cmake_option(
-            "IREE_HAL_DRIVER_VULKAN",
-            "OFF" if platform.system() == "Darwin" else "ON",
-        ),
-        get_env_cmake_option(
-            "IREE_HAL_DRIVER_CUDA",
-            "OFF",
-        ),
-        get_env_cmake_option(
-            "IREE_HAL_DRIVER_HIP",
-            "OFF",
-        ),
-        get_env_cmake_list("IREE_EXTERNAL_HAL_DRIVERS", ""),
-        get_env_cmake_option("IREE_ENABLE_CPUINFO", "ON"),
-    ] + list(extra_cmake_args)
-    add_env_cmake_setting(cmake_args, "IREE_TRACING_PROVIDER")
-    add_env_cmake_setting(cmake_args, "IREE_TRACING_PROVIDER_H")
-
-    # These usually flow through the environment, but we add them explicitly
-    # so that they show clearly in logs (getting them wrong can have bad
-    # outcomes).
-    add_env_cmake_setting(cmake_args, "CMAKE_OSX_ARCHITECTURES")
-    add_env_cmake_setting(
-        cmake_args, "MACOSX_DEPLOYMENT_TARGET", "CMAKE_OSX_DEPLOYMENT_TARGET"
-    )
-
-    # Only do a from-scratch configure if not already configured.
-    cmake_cache_file = os.path.join(cmake_build_dir, "CMakeCache.txt")
-    if not os.path.exists(cmake_cache_file):
-        print(f"Configuring with: {cmake_args}", file=sys.stderr)
-        subprocess.check_call(
-            ["cmake", IREE_SOURCE_DIR] + cmake_args, cwd=cmake_build_dir
+    if not IS_CONFIGURED:
+        # Build from source tree.
+        os.makedirs(cmake_build_dir, exist_ok=True)
+        maybe_nuke_cmake_cache(cmake_build_dir, cmake_install_dir)
+        print(f"CMake build dir: {cmake_build_dir}", file=sys.stderr)
+        print(f"CMake install dir: {cmake_install_dir}", file=sys.stderr)
+        cmake_args = [
+            "-GNinja",
+            "--log-level=VERBOSE",
+            f"-DIREE_RUNTIME_OPTIMIZATION_PROFILE={IREE_RUNTIME_OPTIMIZATION_PROFILE}",
+            "-DIREE_BUILD_PYTHON_BINDINGS=ON",
+            "-DIREE_BUILD_COMPILER=OFF",
+            "-DIREE_BUILD_SAMPLES=OFF",
+            "-DIREE_BUILD_TESTS=OFF",
+            "-DPython3_EXECUTABLE={}".format(sys.executable),
+            "-DCMAKE_BUILD_TYPE={}".format(cfg),
+            get_env_cmake_option(
+                "IREE_HAL_DRIVER_VULKAN",
+                "OFF" if platform.system() == "Darwin" else "ON",
+            ),
+            get_env_cmake_option(
+                "IREE_HAL_DRIVER_CUDA",
+                "OFF",
+            ),
+            get_env_cmake_option(
+                "IREE_HAL_DRIVER_HIP",
+                "OFF",
+            ),
+            get_env_cmake_list("IREE_EXTERNAL_HAL_DRIVERS", ""),
+            get_env_cmake_option("IREE_ENABLE_CPUINFO", "ON"),
+        ] + list(extra_cmake_args)
+        add_env_cmake_setting(cmake_args, "IREE_TRACING_PROVIDER")
+        add_env_cmake_setting(cmake_args, "IREE_TRACING_PROVIDER_H")
+
+        # These usually flow through the environment, but we add them explicitly
+        # so that they show clearly in logs (getting them wrong can have bad
+        # outcomes).
+        add_env_cmake_setting(cmake_args, "CMAKE_OSX_ARCHITECTURES")
+        add_env_cmake_setting(
+            cmake_args, "MACOSX_DEPLOYMENT_TARGET", "CMAKE_OSX_DEPLOYMENT_TARGET"
         )
-    else:
-        print(f"Not re-configuring (already configured)", file=sys.stderr)
 
-    # Build. Since we have restricted to just the runtime, build everything
-    # so as to avoid fragility with more targeted selection criteria.
-    subprocess.check_call(["cmake", "--build", "."], cwd=cmake_build_dir)
-    print("Build complete.", file=sys.stderr)
+        # Only do a from-scratch configure if not already configured.
+        cmake_cache_file = os.path.join(cmake_build_dir, "CMakeCache.txt")
+        if not os.path.exists(cmake_cache_file):
+            print(f"Configuring with: {cmake_args}", file=sys.stderr)
+            subprocess.check_call(
+                ["cmake", IREE_SOURCE_DIR] + cmake_args, cwd=cmake_build_dir
+            )
+        else:
+            print(f"Not re-configuring (already configured)", file=sys.stderr)
+
+        # Build. Since we have restricted to just the runtime, build everything
+        # so as to avoid fragility with more targeted selection criteria.
+        subprocess.check_call(["cmake", "--build", "."], cwd=cmake_build_dir)
+        print("Build complete.", file=sys.stderr)
 
     # Install the component we care about.
     install_args = [
@@ -361,7 +393,9 @@ class CMakeBuildPy(_build_py):
     def run(self):
         # The super-class handles the pure python build.
         super().run()
-        self.build_default_configuration()
+        # If we're in an existing build with tracy enabled, don't build the default config.
+        if not (IS_CONFIGURED and ENABLE_TRACY):
+            self.build_default_configuration()
         if ENABLE_TRACY:
             self.build_tracy_configuration()
 
@@ -388,7 +422,7 @@ def build_default_configuration(self):
                 "_runtime_libs",
             ),
             target_dir,
-            symlinks=False,
+            symlinks=self.editable_mode,
         )
         print("Target populated.", file=sys.stderr)
 
@@ -424,7 +458,7 @@ def build_tracy_configuration(self):
                 "_runtime_libs",
             ),
             target_dir,
-            symlinks=False,
+            symlinks=self.editable_mode,
         )
         print("Target populated.", file=sys.stderr)