Merge branch 'main' into einsum_wip

.github/workflows/build_packages.yml

@@ -0,0 +1,111 @@
+# 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
+
+name: Build packages
+
+on:
+  workflow_dispatch:
+  schedule:
+    # Runs at 11:00 AM UTC, which is 3:00 AM PST (UTC-8)
+    - cron: '0 11 * * *'
+
+jobs:
+  # Note: metadata generation could happen in a separate trigger/schedule
+  # workflow. For cross platform builds, it's useful to just generate the
+  # metadata on Linux and pass that to later jobs using artifacts.
+  setup_metadata:
+    runs-on: ubuntu-24.04
+    outputs:
+      shark_package_version: ${{ steps.version.outputs.shark_package_version }}
+    steps:
+      # For now the version is just a calendar date + an automatically
+      # incrementing value. We may want different versions for nightly/dev
+      # builds and stable releases published to official places like pypi.
+      - name: Compute version
+        id: version
+        run: |
+          shark_package_version="$(printf '%(%Y%m%d)T.${{ github.run_number }}')"
+          echo "shark_package_version=${shark_package_version}" >> $GITHUB_OUTPUT
+          cat << EOF > ./version_info.json
+          {
+            "package-version": "${shark_package_version}"
+          }
+          EOF
+          cat ./version_info.json
+      - name: Upload version_info.json
+        uses: actions/upload-artifact@50769540e7f4bd5e21e526ee35c689e35e0d6874 # v4.4.0
+        with:
+          name: version_info
+          path: version_info.json
+
+  build_packages:
+    name: "${{ matrix.package }} :: ${{ matrix.platform }} :: ${{ matrix.python-version }}"
+    runs-on: ${{ matrix.runs-on }}
+    needs: [setup_metadata]
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          # Ubuntu packages.
+          - runs-on: ubuntu-24.04
+            platform: linux-x86_64
+            package: shortfin
+            python-version: cp312-cp312
+          - runs-on: ubuntu-24.04
+            platform: linux-x86_64
+            package: shortfin
+            python-version: cp313-cp313
+          - runs-on: ubuntu-24.04
+            platform: linux-x86_64
+            package: shortfin
+            python-version: cp313-cp313t
+
+          # TODO(#130): macOS platform
+          # TODO(#130): Windows platform
+          # TODO(#130): sharktank packages
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4.1.7
+      with:
+        path: "c" # Windows can hit path length limits, so use a short path.
+        submodules: false
+
+    - name: Download version_info.json
+      uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8
+      with:
+        name: version_info
+        path: ./c/shortfin/
+        merge-multiple: true
+
+    - name: Build shortfin (Linux x86_64, ${{ matrix.python-version }})
+      if: "matrix.package == 'shortfin' && matrix.platform == 'linux-x86_64'"
+      env:
+        OUTPUT_DIR: "${{ github.workspace }}/bindist"
+        OVERRIDE_PYTHON_VERSIONS: "${{ matrix.python-version }}"
+      run: |
+        [ -e ./bindist/* ] && rm ./bindist/*
+        ./c/shortfin/build_tools/build_linux_package.sh
+
+    - name: Upload python wheels
+      uses: actions/upload-artifact@50769540e7f4bd5e21e526ee35c689e35e0d6874 # v4.4.0
+      with:
+        if-no-files-found: error
+        name: snapshot-${{ matrix.package }}-${{ matrix.platform }}-${{ matrix.python-version }}
+        path: bindist
+
+    - name: Release python wheels
+      uses: ncipollo/release-action@2c591bcc8ecdcd2db72b97d6147f871fcd833ba5 # v1.14.0
+      with:
+        artifacts: bindist/*.whl
+        token: "${{ secrets.GITHUB_TOKEN }}"
+        tag: "dev-wheels"
+        name: "dev-wheels"
+        body: "Automatic snapshot release of SHARK-Platform python wheels."
+        removeArtifacts: false
+        allowUpdates: true
+        replacesArtifacts: true
+        makeLatest: false

shortfin/.gitignore

@@ -0,0 +1,2 @@
+# Local-only config options
+version_info.json

shortfin/README.md

@@ -33,16 +33,22 @@ pip install -v -e .
 
 ## Package Python Release Builds
 
-```bash
-# Build shortfin.*.whl into the dist/ directory
-#   e.g. `shortfin-0.9-cp312-cp312-linux_x86_64.whl`
-python3 -m pip wheel -v -w dist .
+* To build wheels for Linux using a manylinux Docker container:
 
-# Install the built wheel.
-python3 -m pip install dist/*.whl
-```
+    ```bash
+    sudo ./build_tools/build_linux_package.sh
+    ```
+
+* To build a wheel for your host OS/arch manually:
+
+    ```bash
+    # Build shortfin.*.whl into the dist/ directory
+    #   e.g. `shortfin-0.9-cp312-cp312-linux_x86_64.whl`
+    python3 -m pip wheel -v -w dist .
 
-TODO: helper script to build under manylinux using Docker
+    # Install the built wheel.
+    python3 -m pip install dist/*.whl
+    ```
 
 ## Python Dev Builds
 

shortfin/build_tools/build_linux_package.sh

@@ -0,0 +1,122 @@
+#!/bin/bash
+# Copyright 2024 Advanced Micro Devices, Inc.
+# 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
+
+# build_linux_package.sh
+# One stop build of shortfin Python package for Linux. The Linux build is
+# complicated because it has to be done via a manylinux docker container.
+#
+# Usage:
+# Build everything (all python versions):
+#   sudo ./build_tools/build_linux_package.sh
+#
+# Build specific Python versions to custom directory:
+#   OVERRIDE_PYTHON_VERSIONS="cp312-cp312 cp313-cp313" \
+#   OUTPUT_DIR="/tmp/wheelhouse" \
+#   sudo -E ./build_tools/build_linux_package.sh
+#
+# Valid Python versions match a subdirectory under /opt/python in the docker
+# image. Typically:
+#   cp312-cp312 cp313-cp313
+#
+# Note that this script is meant to be run on CI and it will pollute both the
+# output directory and in-tree build/ directories with docker created, root
+# owned builds. Sorry - there is no good way around it.
+#
+# It can be run on a workstation but recommend using a git worktree dedicated
+# to packaging to avoid stomping on development artifacts.
+set -xeu -o errtrace
+
+THIS_DIR="$(cd $(dirname $0) && pwd)"
+REPO_ROOT="$(cd "$THIS_DIR"/../../ && pwd)"
+SCRIPT_NAME="$(basename $0)"
+ARCH="$(uname -m)"
+
+# TODO(#130): Update to manylinux_2_28, upstream or a fork
+#   * upstream uses a version of gcc that has build warnings/errors
+#   * https://github.com/nod-ai/base-docker-images is a bit out of date but can include a recent clang
+# MANYLINUX_DOCKER_IMAGE="${MANYLINUX_DOCKER_IMAGE:-quay.io/pypa/manylinux_2_28_${ARCH}:latest}"
+MANYLINUX_DOCKER_IMAGE="${MANYLINUX_DOCKER_IMAGE:-quay.io/pypa/manylinux2014_${ARCH}:latest}"
+PYTHON_VERSIONS="${OVERRIDE_PYTHON_VERSIONS:-cp312-cp312 cp313-cp313}"
+OUTPUT_DIR="${OUTPUT_DIR:-${THIS_DIR}/wheelhouse}"
+
+function run_on_host() {
+  echo "Running on host"
+  echo "Launching docker image ${MANYLINUX_DOCKER_IMAGE}"
+
+  # Canonicalize paths.
+  mkdir -p "${OUTPUT_DIR}"
+  OUTPUT_DIR="$(cd "${OUTPUT_DIR}" && pwd)"
+  echo "Outputting to ${OUTPUT_DIR}"
+  mkdir -p "${OUTPUT_DIR}"
+  docker run --rm \
+    -v "${REPO_ROOT}:${REPO_ROOT}" \
+    -v "${OUTPUT_DIR}:${OUTPUT_DIR}" \
+    -e __MANYLINUX_BUILD_WHEELS_IN_DOCKER=1 \
+    -e "OVERRIDE_PYTHON_VERSIONS=${PYTHON_VERSIONS}" \
+    -e "OUTPUT_DIR=${OUTPUT_DIR}" \
+    "${MANYLINUX_DOCKER_IMAGE}" \
+    -- ${THIS_DIR}/${SCRIPT_NAME}
+
+  echo "******************** BUILD COMPLETE ********************"
+  echo "Generated binaries:"
+  ls -l "${OUTPUT_DIR}"
+}
+
+function run_in_docker() {
+  echo "Running in docker"
+  echo "Marking git safe.directory"
+  git config --global --add safe.directory '*'
+
+  echo "Using python versions: ${PYTHON_VERSIONS}"
+  local orig_path="${PATH}"
+
+  # Build phase.
+  echo "******************** BUILDING PACKAGE ********************"
+  for python_version in ${PYTHON_VERSIONS}; do
+    python_dir="/opt/python/${python_version}"
+    if ! [ -x "${python_dir}/bin/python" ]; then
+      echo "ERROR: Could not find python: ${python_dir} (skipping)"
+      continue
+    fi
+    export PATH="${python_dir}/bin:${orig_path}"
+    echo ":::: Python version $(python --version)"
+    clean_wheels "shortfin" "${python_version}"
+    build_shortfin
+    run_audit_wheel "shortfin" "${python_version}"
+  done
+}
+
+function build_shortfin() {
+  export SHORTFIN_ENABLE_TRACING=ON
+  python -m pip wheel --disable-pip-version-check -v -w "${OUTPUT_DIR}" "${REPO_ROOT}/shortfin"
+}
+
+function run_audit_wheel() {
+  local wheel_basename="$1"
+  local python_version="$2"
+  # Force wildcard expansion here
+  generic_wheel="$(echo "${OUTPUT_DIR}/${wheel_basename}-"*"-${python_version}-linux_${ARCH}.whl")"
+  ls "${generic_wheel}"
+  echo ":::: Auditwheel ${generic_wheel}"
+  auditwheel repair -w "${OUTPUT_DIR}" "${generic_wheel}"
+  rm -v "${generic_wheel}"
+}
+
+function clean_wheels() {
+  local wheel_basename="$1"
+  local python_version="$2"
+  echo ":::: Clean wheels ${wheel_basename} ${python_version}"
+  rm -f -v "${OUTPUT_DIR}/${wheel_basename}-"*"-${python_version}-"*".whl"
+}
+
+# Trampoline to the docker container if running on the host.
+if [ -z "${__MANYLINUX_BUILD_WHEELS_IN_DOCKER-}" ]; then
+  run_on_host "$@"
+else
+  run_in_docker "$@"
+fi

shortfin/setup.py

@@ -4,16 +4,17 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from distutils.core import setup, Extension
-import sys
+import json
+import os
 import shutil
 import subprocess
-import os
-from pathlib import Path
+import sys
 from distutils.command.build import build as _build
+from distutils.core import setup, Extension
+from pathlib import Path
 from setuptools import find_namespace_packages
-from setuptools.command.build_ext import build_ext as _build_ext
 from setuptools.command.build_py import build_py as _build_py
+from setuptools.command.build_ext import build_ext as _build_ext
 
 
 def get_env_boolean(name: str, default_value: bool = False) -> bool:
@@ -133,6 +134,40 @@ def copy_extensions_to_source(self, *args, **kwargs):
         ...
 
 
+# Setup and get version information.
+VERSION_INFO_FILE = os.path.join(SOURCE_DIR, "version_info.json")
+
+
+def load_version_info():
+    with open(VERSION_INFO_FILE, "rt") as f:
+        return json.load(f)
+
+
+def find_git_version():
+    try:
+        return (
+            subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=SOURCE_DIR)
+            .decode("utf-8")
+            .strip()
+        )
+    except subprocess.SubprocessError as e:
+        print(f"ERROR: Could not get git revision: {e}", file=sys.stderr)
+    return None
+
+
+try:
+    version_info = load_version_info()
+except FileNotFoundError:
+    print("version_info.json not found. Using defaults", file=sys.stderr)
+    version_info = {}
+git_version = find_git_version()
+
+PACKAGE_VERSION = version_info.get("package-version")
+if not PACKAGE_VERSION:
+    PACKAGE_VERSION = f"0.dev0+{git_version or '0'}"
+print(f"Using PACKAGE_VERSION: '{PACKAGE_VERSION}'")
+
+
 def maybe_nuke_cmake_cache(cmake_build_dir):
     # From run to run under pip, we can end up with different paths to ninja,
     # which isn't great and will confuse cmake. Detect if the location of
@@ -324,7 +359,7 @@ def populate_built_package(abs_dir):
 
 setup(
     name="shortfin",
-    version="0.9",
+    version=f"{PACKAGE_VERSION}",
     description="Shortfin native library implementation",
     author="SHARK Authors",
     packages=packages,