Release pre-release 2.3.0rc20240410. (#605)

Includes some scripting and release documentation updates that have been
in my queue for a while.

In this version, we start executing on the change to deploy the
`iree-turbine` package and the `iree.turbine` namespace. This is just a
copy of the shark-turbine package currently and the top level namespace
is just a redirect of the public API. The next release will complete the
switchover.

Since this is pushed as an rc, it is a pre-release requiring
installation with the `--pre` flag. When torch 2.3.0 drops, we should
push a final.

build_tools/build_release.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python
+# Copyright 2024 Advanced Micro Devices, Inc.
+#
+# 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
+
+"""Fetches dependent release artifacts and builds wheels.
+
+See docs/releasing.md for usage.
+"""
+
+import argparse
+from datetime import date
+import json
+import os
+from pathlib import Path
+import shlex
+import subprocess
+
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+VERSION_INFO_FILE = REPO_ROOT / "version_info.json"
+CORE_DIR = REPO_ROOT / "core"
+WHEEL_DIR = REPO_ROOT / "wheelhouse"
+
+# The platform flags that we will download IREE wheels for. This must match
+# the platforms and Python versions we build. If it mismatches or something
+# is wrong, this will error. Note that the platform and python-version
+# indicates "fetch me a wheel that will install on this combo" vs "fetch me
+# a specific wheel".
+IREE_PLATFORM_ARGS = [
+    # Linux aarch64
+    ["--platform", "manylinux_2_28_aarch64", "--python-version", "3.9"],
+    ["--platform", "manylinux_2_28_aarch64", "--python-version", "3.10"],
+    ["--platform", "manylinux_2_28_aarch64", "--python-version", "3.11"],
+    ["--platform", "manylinux_2_28_aarch64", "--python-version", "3.12"],
+    # Linux x86_64
+    ["--platform", "manylinux_2_28_x86_64", "--python-version", "3.9"],
+    ["--platform", "manylinux_2_28_x86_64", "--python-version", "3.10"],
+    ["--platform", "manylinux_2_28_x86_64", "--python-version", "3.11"],
+    ["--platform", "manylinux_2_28_x86_64", "--python-version", "3.12"],
+    # MacOS
+    ["--platform", "macosx_13_0_universal2", "--python-version", "3.11"],
+    # Windows
+    ["--platform", "win_amd64", "--python-version", "3.11"],
+]
+
+
+def eval_version(version_spec: str):
+    date_stamp = date.today().strftime("%Y%m%d")
+    return version_spec.replace("YYYYMMDD", date_stamp)
+
+
+def write_version_info(args):
+    with open(VERSION_INFO_FILE, "rt") as f:
+        info_dict = json.load(f)
+
+    # Compute core-version.
+    core_version = eval_version(args.core_version)
+    if args.core_pre_version:
+        core_version += eval_version(args.core_pre_version)
+    if args.core_post_version:
+        core_version += f".{eval_version(args.core_post_version)}"
+    info_dict["core-version"] = core_version
+
+    with open(VERSION_INFO_FILE, "wt") as f:
+        json.dump(info_dict, f)
+
+    print(f"Updated version_info.json:\n{json.dumps(info_dict, indent=2)}")
+
+
+def exec(args, env=None):
+    args = [str(s) for s in args]
+    print(f": Exec: {shlex.join(args)}")
+    if env is not None:
+        full_env = dict(os.environ)
+        full_env.update(env)
+    else:
+        full_env = None
+    subprocess.check_call(args, env=full_env)
+
+
+def download_requirements(requirements_file, platforms=()):
+    args = [
+        "pip",
+        "download",
+        "-d",
+        WHEEL_DIR,
+    ]
+    if platforms:
+        args.append("--no-deps")
+        for p in platforms:
+            args.extend(["--platform", p])
+    args += [
+        "-f",
+        WHEEL_DIR,
+        "-r",
+        requirements_file,
+    ]
+    exec(args)
+
+
+def download_iree_binaries():
+    for platform_args in IREE_PLATFORM_ARGS:
+        print("Downloading for platform:", platform_args)
+        args = [
+            "pip",
+            "download",
+            "-d",
+            WHEEL_DIR,
+            "--no-deps",
+        ]
+        args.extend(platform_args)
+        args += [
+            "-f",
+            "https://openxla.github.io/iree/pip-release-links.html",
+            "-f",
+            WHEEL_DIR,
+            "-r",
+            CORE_DIR / "iree-requirements.txt",
+        ]
+        exec(args)
+
+
+def build_wheel(path, env=None):
+    exec(
+        ["pip", "wheel", "--no-index", "-f", WHEEL_DIR, "-w", WHEEL_DIR, path], env=env
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--core-version", help="Version for the core component", required=True
+    )
+    parser.add_argument(
+        "--core-pre-version",
+        help="Pre-release version segment or (YYYYMMDD)",
+        default="",
+    )
+    parser.add_argument(
+        "--core-post-version",
+        help="Post-release version segment or (YYYYMMDD)",
+        default="",
+    )
+    parser.add_argument(
+        "--no-download", help="Disable dep download", action="store_true"
+    )
+    args = parser.parse_args()
+
+    write_version_info(args)
+    WHEEL_DIR.mkdir(parents=True, exist_ok=True)
+
+    if not args.no_download:
+        print("Prefetching all IREE binaries")
+        download_iree_binaries()
+        print("Prefetching torch CPU")
+        download_requirements(CORE_DIR / "pytorch-cpu-requirements.txt")
+        print("Downloading remaining requirements")
+        download_requirements(CORE_DIR / "requirements.txt")
+
+    print("Building shark-turbine")
+    build_wheel(CORE_DIR)
+    print("Building iree-turbine")
+    build_wheel(CORE_DIR, env={"TURBINE_PACKAGE_NAME": "iree-turbine"})
+
+
+if __name__ == "__main__":
+    main()

core/iree/turbine/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2024 Advanced Micro Devices, Inc.
+#
+# 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
+
+# TODO: This redirection layer exists while we are migrating from the
+# shark_turbine top-level package name to iree.turbine. It exports the
+# public API but not the internal details. In a future switch, all code
+# will be directly located here and the redirect will be done in the
+# shark_turbine namespace.
+
+from shark_turbine import aot
+from shark_turbine import dynamo
+from shark_turbine import kernel
+from shark_turbine import ops
+from shark_turbine import runtime

core/requirements.txt

@@ -6,3 +6,7 @@
 
 -r pytorch-requirements.txt
 -r iree-requirements.txt
+
+# From pyproject.toml.
+setuptools
+wheel

core/setup.py

@@ -15,6 +15,8 @@
 REPO_DIR = os.path.dirname(THIS_DIR)
 VERSION_INFO_FILE = os.path.join(REPO_DIR, "version_info.json")
 
+# Transitional as we migrate from shark-turbine -> iree-turbine.
+TURBINE_PACKAGE_NAME = os.getenv("TURBINE_PACKAGE_NAME", "shark-turbine")
 
 with open(
     os.path.join(
@@ -32,10 +34,12 @@ def load_version_info():
 
 
 version_info = load_version_info()
-PACKAGE_VERSION = version_info["package-version"]
+PACKAGE_VERSION = version_info["core-version"]
 
 packages = find_namespace_packages(
     include=[
+        "iree.turbine",
+        "iree.turbine.*",
         "shark_turbine",
         "shark_turbine.*",
     ],
@@ -77,7 +81,7 @@ def initialize_options(self):
 
 
 setup(
-    name=f"shark-turbine",
+    name=f"{TURBINE_PACKAGE_NAME}",
     version=f"{PACKAGE_VERSION}",
     author="SHARK Authors",
     author_email="[email protected]",
@@ -87,7 +91,7 @@ def initialize_options(self):
     url="https://github.com/nod-ai/SHARK-Turbine",
     license="Apache-2.0",
     classifiers=[
-        "Development Status :: 3 - Alpha",
+        "Development Status :: 5 - Production/Stable",
         "License :: OSI Approved :: Apache Software License",
         "Programming Language :: Python :: 3",
     ],
@@ -102,6 +106,9 @@ def initialize_options(self):
         f"iree-compiler{get_version_spec('iree-compiler')}",
         f"iree-runtime{get_version_spec('iree-runtime')}",
         # Use the [torch-cpu-nightly] spec to get a more recent/specific version.
+        # Note that during the transition to torch 2.3.0 we technically support
+        # back to torch 2.1, which is why we pin here in this way. However,
+        # the CI tests on 2.3.
         "torch>=2.1.0",
     ],
     extras_require={

core/tests/top_level_package_test.py

@@ -0,0 +1,20 @@
+# Copyright 2024 Advanced Micro Devices, Inc.
+#
+# 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
+
+import logging
+import unittest
+
+
+class TopLevelPackageTest(unittest.TestCase):
+    def testIreeTurbineRedirect(self):
+        # We have a temporary redirect of the top-level API to the
+        # iree.turbine namespace.
+        from iree.turbine import aot, dynamo, kernel, ops, runtime
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.DEBUG)
+    unittest.main()

docs/releasing.md

@@ -0,0 +1,75 @@
+# Releasing SHARK-Turbine/core
+
+There are multiple release artifacts that are deployed from this project:
+
+* shark-turbine wheel (transitional while switching to iree-turbine)
+* iree-turbine wheel
+* iree-compiler wheels
+* iree-runtime wheels
+
+Typically we deploy IREE compiler and runtime wheels along with a turbine
+release, effectively promoting a nightly.
+
+## Building Artifacts
+
+Build a pre-release:
+
+```
+./build_tools/build_release.py --core-version 2.3.0 --core-pre-version=rcYYYYMMDD
+```
+
+Build an official release:
+
+```
+./build_tools/build_release.py --core-version 2.3.0
+```
+
+This will download all deps, including wheels for all supported platforms and
+Python versions for iree-compiler and iree-runtime. All wheels will be placed
+in the `wheelhouse/` directory.
+
+
+## Testing
+
+TODO: Write a script for this.
+
+```
+python -m venv wheelhouse/test.venv
+source wheelhouse/test.venv/bin/activate
+pip install -f wheelhouse iree-turbine[testing]
+# Temp: tests require torchvision.
+pip install -f wheelhouse torchvision
+pytest core/tests
+```
+
+## Push
+
+From the testing venv, verify that everything is sane:
+
+```
+pip freeze
+```
+
+Push IREE deps (if needed/updated):
+
+```
+twine upload wheelhouse/iree_compiler-* wheelhouse/iree_runtime-*
+```
+
+Push built wheels:
+
+```
+twine upload wheelhouse/iree_turbine-* wheelhouse/shark_turbine-*
+```
+
+## Install from PyPI and Sanity Check
+
+TODO: Script this
+
+From the testing venv:
+
+```
+pip uninstall -y shark-turbine iree-turbine iree-compiler iree-runtime
+pip install iree-turbine
+pytest core/tests
+```

version_info.json

@@ -1,3 +1 @@
-{
-    "package-version": "0.9.7.dev1"
-}
+{"core-version": "2.3.0rc20240410", "package-version": "0.9.7.dev1"}