Re-organize the runtime python directory to be more extensible. (iree…

…-org#9071)

* Native library goes from iree/runtime/_binding.so -> iree/_runtime.so
* Make iree/runtime/_binding.py a trampoline to the correct runtime module. In the future, this will allow us to use a different native implementation based on environment and other packages installed (i.e. would allow an iree-runtime-instrumented and iree-runtime-debug to co-exist and be selected at runtime).
* Move tests to a top-level tests directory (common convention).
* Move native code to the top-level.
* Adds a package_test which builds/installs the package to a venv and does some sanity checking of the modules and scripts (takes about ~12s to run on my machine, so should be hidden in the noise of bigger tests).

See iree-org#9080 for updating docker images and actually enabling package_test.

runtime/CMakeLists.txt

@@ -12,5 +12,5 @@ if(IREE_BUILD_PYTHON_BINDINGS)
   configure_file(pyproject.toml pyproject.toml COPYONLY)
   configure_file(setup.py setup.py @ONLY)
 
-  add_subdirectory(bindings/python/iree/runtime)
+  add_subdirectory(bindings/python)
 endif()

runtime/bindings/python/iree/runtime/CMakeLists.txt → runtime/bindings/python/CMakeLists.txt

@@ -15,7 +15,7 @@ set(_TRACY_ENABLED OFF)
 if(TARGET IREETracyCaptureServer)
   message(STATUS "Bundline Tracy CLI tools with Python API")
   set(_TRACY_ENABLED ON)
-  list(APPEND _PYTHON_EXTRA_SRCS "scripts/iree-tracy-capture")
+  list(APPEND _PYTHON_EXTRA_SRCS "iree/runtime/scripts/iree-tracy-capture")
   list(APPEND _EXTRA_INSTALL_TOOL_TARGETS "IREETracyCaptureServer")
 endif()
 
@@ -26,7 +26,7 @@ endif()
 iree_pyext_module(
   NAME
     PyExtRt
-  MODULE_NAME binding
+  MODULE_NAME iree/_runtime
   SRCS
     "binding.h"
     "initialize_module.cc"
@@ -60,15 +60,16 @@ iree_py_library(
   NAME
     runtime
   SRCS
-    "__init__.py"
-    "array_interop.py"
-    "flags.py"
-    "function.py"
-    "system_api.py"
-    "tracing.py"
-    "scripts/iree_benchmark_trace/__main__.py"
-    "scripts/iree_run_trace/__main__.py"
-    "scripts/iree_run_module/__main__.py"
+    "iree/runtime/__init__.py"
+    "iree/runtime/_binding.py"
+    "iree/runtime/array_interop.py"
+    "iree/runtime/flags.py"
+    "iree/runtime/function.py"
+    "iree/runtime/system_api.py"
+    "iree/runtime/tracing.py"
+    "iree/runtime/scripts/iree_benchmark_trace/__main__.py"
+    "iree/runtime/scripts/iree_run_trace/__main__.py"
+    "iree/runtime/scripts/iree_run_module/__main__.py"
     ${_PYTHON_EXTRA_SRCS}
   PYEXT_DEPS
     ::PyExtRt
@@ -114,44 +115,55 @@ iree_py_test(
   NAME
     array_interop_test
   SRCS
-    "array_interop_test.py"
+    "tests/array_interop_test.py"
 )
 
 iree_py_test(
   NAME
     flags_test
   SRCS
-    "flags_test.py"
+    "tests/flags_test.py"
 )
 
 iree_py_test(
   NAME
     function_test
   SRCS
-    "function_test.py"
+    "tests/function_test.py"
 )
 
 iree_py_test(
   NAME
     hal_test
   SRCS
-    "hal_test.py"
+    "tests/hal_test.py"
 )
 
 iree_py_test(
   NAME
     system_api_test
   SRCS
-    "system_api_test.py"
+    "tests/system_api_test.py"
 )
 
 iree_py_test(
   NAME
     vm_test
   SRCS
-    "vm_test.py"
+    "tests/vm_test.py"
 )
 
+# TODO: Enable this once the CI bots are updated to install the python3-venv
+# apt package. https://github.com/google/iree/issues/9080
+# iree_py_test(
+#   NAME
+#     package_test
+#   SRCS
+#     "tests/package_test.py"
+#   ARGS
+#     "${IREE_BINARY_DIR}/runtime"
+# )
+
 ################################################################################
 # Install
 ################################################################################
@@ -163,7 +175,7 @@ iree_py_install_package(
   DEPS
     # TODO: Update CMake target/path mangling rules to make this syntactically
     # rooted on "iree" in some way.
-    runtime_bindings_python_iree_runtime_PyExtRt
+    runtime_bindings_python_PyExtRt
     iree_tools_iree-benchmark-module
     iree_tools_iree-benchmark-trace
     iree_tools_iree-run-module
@@ -176,7 +188,7 @@ iree_py_install_package(
 install(
   # TODO: Update CMake target/path mangling rules to make this syntactically
   # rooted on "iree" in some way.
-  TARGETS runtime_bindings_python_iree_runtime_PyExtRt
+  TARGETS runtime_bindings_python_PyExtRt
   COMPONENT ${PY_INSTALL_COMPONENT}
   DESTINATION "${PY_INSTALL_MODULE_DIR}"
 )

runtime/bindings/python/iree/runtime/initialize_module.cc → runtime/bindings/python/initialize_module.cc

@@ -16,7 +16,7 @@
 namespace iree {
 namespace python {
 
-PYBIND11_MODULE(binding, m) {
+PYBIND11_MODULE(_runtime, m) {
   IREE_CHECK_OK(iree_hal_register_all_available_drivers(
       iree_hal_driver_registry_default()));
 

runtime/bindings/python/iree/runtime/__init__.py

@@ -10,11 +10,11 @@
 # pylint: disable=g-bad-import-order
 # pylint: disable=wildcard-import
 
-from . import binding
+from . import _binding
 
 # Pull some of the native symbols into the public API.
 # Hal imports
-from .binding import (
+from ._binding import (
     BufferCompatibility,
     BufferUsage,
     HalAllocator,
@@ -29,7 +29,7 @@
 )
 
 # Vm imports
-from .binding import (
+from ._binding import (
     create_hal_module,
     Linkage,
     VmVariantList,

runtime/bindings/python/iree/runtime/_binding.py

@@ -0,0 +1,18 @@
+# Copyright 2022 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+"""Imports the native runtime library.
+
+All code in the runtime should use runtime imports via this module, which
+locates the actual _runtime module based on environment configuration.
+Currently, we only bundle a single runtime library, but in the future, this
+will let us dynamically switch between instrumented, debug, etc by changing
+the way this trampoline functions.
+"""
+
+import sys
+
+from iree import _runtime
+sys.modules[__name__] = _runtime

runtime/bindings/python/iree/runtime/array_interop.py

@@ -10,7 +10,7 @@
 import numpy as np
 import numpy.lib.mixins
 
-from .binding import (
+from ._binding import (
     BufferUsage,
     HalBufferView,
     HalDevice,

runtime/bindings/python/iree/runtime/flags.py

@@ -4,7 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from .binding import parse_flags
+from ._binding import parse_flags
 
 # When enabled, performs additional function input validation checks. In the
 # event of errors, this will yield nicer error messages but comes with a

runtime/bindings/python/iree/runtime/function.py

@@ -11,7 +11,7 @@
 
 import numpy as np
 
-from .binding import (
+from ._binding import (
     _invoke_statics,
     ArgumentPacker,
     BufferUsage,

runtime/bindings/python/iree/runtime/system_api.py

@@ -33,7 +33,7 @@
 
 from typing import Any, List, Mapping, Optional, Sequence, Tuple, Union
 
-from . import binding as _binding
+from . import _binding
 from .function import FunctionInvoker
 from . import tracing
 

runtime/bindings/python/iree/runtime/tracing.py

@@ -13,7 +13,7 @@
 import os
 import sys
 
-from . import binding as _binding
+from . import _binding
 
 try:
   import yaml

runtime/bindings/python/iree/runtime/function_test.py → runtime/bindings/python/tests/function_test.py

@@ -14,7 +14,7 @@
     IMPLICIT_BUFFER_ARG_MEMORY_TYPE,
     IMPLICIT_BUFFER_ARG_USAGE,
 )
-from iree.runtime.binding import VmVariantList
+from iree.runtime._binding import VmVariantList
 
 
 class MockVmContext:

runtime/bindings/python/tests/package_test.py

@@ -0,0 +1,73 @@
+# Copyright 2022 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+"""Creates a venv and installs the runtime package, doing some smoke tests.
+
+This is not a perfect approximation of how the runtime package is built
+for real since it is run from the build dir and does not do a recursive
+invocation of CMake. However, it can detect gross errors in the installation
+process, including missing modules, scripts, etc.
+"""
+
+import os
+import subprocess
+import sys
+import tempfile
+from typing import List
+
+SETUP_PY_DIR = sys.argv[1]
+if not os.path.exists(os.path.join(SETUP_PY_DIR, "setup.py")):
+  print("ERROR: Expected directory containing setup.py as argument")
+print(f"Using setup.py directory: {SETUP_PY_DIR}")
+
+# Figure out where to stage output.
+TEMP_DIR = os.getenv("TEST_TMPDIR")
+if not TEMP_DIR:
+  TEMP_DIR = tempfile.gettempdir()
+
+# Create the venv.
+VENV_DIR = os.path.join(TEMP_DIR, "iree_runtime_venv")
+print(f"Using venv directory: {VENV_DIR}")
+subprocess.check_call([sys.executable, "-m", "venv", VENV_DIR])
+venv_python = None
+for venv_bin in [
+    os.path.join(VENV_DIR, "bin"),  # Posix.
+    os.path.join(VENV_DIR, "Scripts")  # Windows.
+]:
+  if os.path.exists(os.path.join(venv_bin, "activate")):
+    venv_python = os.path.join(venv_bin, "python")
+if not venv_python:
+  print("ERROR: Could not find venv python")
+venv_bin = os.path.dirname(venv_python)
+print(f"Running with python: {venv_python}")
+
+# Install the package.
+subprocess.check_call([
+    venv_python, "-m", "pip", "install", "--force-reinstall", SETUP_PY_DIR + "/"
+])
+
+# Run some sanity checks.
+if "PYTHONPATH" in os.environ:
+  del os.environ["PYTHONPATH"]
+
+print("***** Sanity checking that module loads...")
+subprocess.check_call(
+    [venv_python, "-c", "import iree.runtime; print('Runtime loaded')"],
+    cwd=venv_bin)
+
+
+# Check tools.
+def check_tool(tool_name: str, args: List[str]):
+  print(f"**** Checking tool {tool_name} with args {args}")
+  subprocess.check_call([os.path.join(venv_bin, tool_name)] + args,
+                        cwd=venv_bin)
+
+
+check_tool("iree-benchmark-module", ["--help"])
+check_tool("iree-benchmark-trace", ["--help"])
+check_tool("iree-run-module", ["--help"])
+check_tool("iree-run-trace", ["--help"])
+
+print("***** All done *****")

runtime/bindings/python/iree/runtime/unix_version.lds → runtime/bindings/python/unix_version.lds

@@ -6,6 +6,6 @@
  */
 
 {
-  global: PyInit_binding;
+  global: PyInit__runtime;
   local: *;
 };

runtime/setup.py

@@ -260,16 +260,14 @@ def prepare_installation():
       print(f"Not re-configuring (already configured)", file=sys.stderr)
 
     # Build.
-    subprocess.check_call([
-        "cmake", "--build", ".", "--target",
-        "runtime/bindings/python/iree/runtime/all"
-    ],
-                          cwd=IREE_BINARY_DIR)
+    subprocess.check_call(
+        ["cmake", "--build", ".", "--target", "runtime/bindings/python/all"],
+        cwd=IREE_BINARY_DIR)
     print("Build complete.", file=sys.stderr)
 
   # Install the directory we care about.
   install_subdirectory = os.path.join(IREE_BINARY_DIR, "runtime", "bindings",
-                                      "python", "iree", "runtime")
+                                      "python")
   install_args = [
       "-DCMAKE_INSTALL_DO_STRIP=ON",
       f"-DCMAKE_INSTALL_PREFIX={CMAKE_INSTALL_DIR_ABS}/",
@@ -348,6 +346,7 @@ def generate_version_py():
                                                       "python_packages",
                                                       "iree_runtime"),
                                    include=[
+                                       "iree._runtime",
                                        "iree.runtime",
                                        "iree.runtime.*",
                                    ])
@@ -381,7 +380,7 @@ def generate_version_py():
     url="https://github.com/google/iree",
     python_requires=">=3.7",
     ext_modules=[
-        CMakeExtension("iree.runtime.binding"),
+        CMakeExtension("iree._runtime"),
     ],
     cmdclass={
         "build": CustomBuild,