[pkgs] Add Python wheel packaging infrastructure for aster

.gitignore

@@ -14,3 +14,7 @@
 /examples/trace_*
 /examples/*zip
 /llvm/llvm-project/
+/wheelhouse/
+*.whl
+/.llvm-extract/
+/.llvm-pkg.tar.gz

CMakeLists.txt

@@ -49,6 +49,30 @@ add_subdirectory(llvm)
 
 configure_llvm()
 
+# When building a Python wheel via scikit-build-core, redirect install prefixes
+# so Python modules and tools land in the correct wheel locations.  For normal
+# builds the cache variables default to the standard development layout.
+if(DEFINED SKBUILD_PLATLIB_DIR)
+  # Compute the relative path from CMAKE_INSTALL_PREFIX to SKBUILD_PLATLIB_DIR
+  # so Python modules and tools land under the aster/ package inside the wheel.
+  file(RELATIVE_PATH _skbuild_rel_platlib
+    "${CMAKE_INSTALL_PREFIX}" "${SKBUILD_PLATLIB_DIR}")
+  if(_skbuild_rel_platlib)
+    set(_aster_prefix "${_skbuild_rel_platlib}/aster")
+  else()
+    set(_aster_prefix "aster")
+  endif()
+  set(ASTER_PYTHON_INSTALL_PREFIX "${_aster_prefix}"
+    CACHE STRING "Python install prefix relative to CMAKE_INSTALL_PREFIX" FORCE)
+  set(ASTER_TOOLS_INSTALL_PREFIX "${_aster_prefix}/_mlir_libs"
+    CACHE STRING "Tools install prefix relative to CMAKE_INSTALL_PREFIX" FORCE)
+else()
+  set(ASTER_PYTHON_INSTALL_PREFIX "python_packages/aster"
+    CACHE STRING "Python install prefix relative to CMAKE_INSTALL_PREFIX")
+  set(ASTER_TOOLS_INSTALL_PREFIX "bin"
+    CACHE STRING "Tools install prefix relative to CMAKE_INSTALL_PREFIX")
+endif()
+
 ################################################################################
 # Configure aster
 ################################################################################

pyproject.toml

@@ -1,3 +1,81 @@
+# pyproject.toml - Build and tool configuration for the aster project.
+#
+# Editable install (local development):
+#   CMAKE_PREFIX_PATH=<llvm-install>/lib/cmake pip install -e .
+#   CMAKE_PREFIX_PATH=<llvm-install>/lib/cmake pip install -e .[dev]
+#
+# Wheel build (via cibuildwheel, requires LLVM tarball):
+#   bash tools/pkgs/aster/build.sh --llvm-tarball=<path-to-llvm.tar.gz>
+
+[build-system]
+requires = ["scikit-build-core>=0.9", "nanobind>=2.9,<3.0", "cmake>=3.31.6", "ninja"]
+build-backend = "scikit_build_core.build"
+
+[project]
+name = "aster"
+version = "0.42.0"
+requires-python = ">=3.11"
+description = "ASTER MLIR-based compiler Python bindings"
+license = { text = "LLVM Apache-2.0" }
+dependencies = [
+  "nanobind>=2.9,<3.0",
+  "PyYAML>=5.4.0,<7.0",
+  "typing_extensions>=4.12.2",
+  "numpy>=2.1.0,<2.2.0",
+  "ml_dtypes>=0.5.0,<0.7.0",
+  "tqdm",
+]
+
+[project.optional-dependencies]
+dev = [
+  "lit>=18.1.7",
+  "black",
+  "pre-commit",
+  "pytest",
+  "pytest-xdist",
+  "pyright",
+]
+
+[tool.scikit-build]
+cmake.build-type = "Release"
+build.verbose = true
+# No pure-Python packages to collect; all content is placed by cmake --install.
+wheel.packages = []
+
+[tool.scikit-build.cmake.define]
+ASTER_ENABLE_PYTHON      = "ON"
+ASTER_ENABLE_UNIT_TESTS  = "OFF"
+MLIR_BINDINGS_PYTHON_NB_DOMAIN = "aster"
+
+[tool.cibuildwheel]
+build = "cp311-* cp312-* cp313-*"
+skip  = "*-musllinux*"
+# Runs once per container: extracts the LLVM tarball and installs build tools.
+before-all = "bash {project}/tools/pkgs/aster/before-all.sh"
+
+[tool.cibuildwheel.linux]
+archs = ["x86_64", "aarch64"]
+manylinux-x86_64-image  = "quay.io/pypa/manylinux_2_28_x86_64"
+manylinux-aarch64-image = "quay.io/pypa/manylinux_2_28_aarch64"
+repair-wheel-command = "auditwheel repair -w {dest_dir} {wheel}"
+
+# before-all.sh symlinks /opt/llvm -> <extracted>/lib/cmake so CMake's
+# find_package() resolves /opt/llvm/<PackageName>/Config.cmake.
+[tool.cibuildwheel.linux.environment]
+LLVM_PKG_PATH     = ".llvm-pkg.tar.gz"
+CMAKE_PREFIX_PATH = "/opt/llvm"
+CC                = "clang"
+CXX               = "clang++"
+
+[tool.cibuildwheel.macos]
+archs = ["x86_64", "arm64"]
+repair-wheel-command = "delocate-wheel --require-archs {delocate_archs} -w {dest_dir} {wheel}"
+
+# before-all.sh symlinks ${HOME}/.llvm-pkg -> <extracted>/lib/cmake.
+[tool.cibuildwheel.macos.environment]
+LLVM_PKG_PATH     = ".llvm-pkg.tar.gz"
+CMAKE_PREFIX_PATH = "{homedir}/.llvm-pkg"
+
 [tool.pytest.ini_options]
 testpaths = ["test", "mlir_kernels", "contrib", "python"]
 addopts = "--import-mode=importlib"

python/CMakeLists.txt

@@ -15,7 +15,7 @@ set(CMAKE_VISIBILITY_INLINES_HIDDEN "ON")
 add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=aster.")
 
 set(_PYTHON_BUILD_PREFIX "${ASTER_BINARY_DIR}/python_packages/")
-set(_PYTHON_INSTALL_PREFIX "python_packages/aster")
+set(_PYTHON_INSTALL_PREFIX "${ASTER_PYTHON_INSTALL_PREFIX}")
 
 set(CMAKE_PLATFORM_NO_VERSIONED_SONAME 1)
 
@@ -86,7 +86,7 @@ mlir_generate_type_stubs(
   OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs"
   OUTPUTS "${_core_type_stub_sources}"
   DEPENDS_TARGET_SRC_DEPS "${_core_extension_srcs}"
-  IMPORT_PATHS "${ASTER_BINARY_DIR}/${_PYTHON_INSTALL_PREFIX}/_mlir_libs"
+  IMPORT_PATHS "${ASTER_BINARY_DIR}/python_packages/aster/_mlir_libs"
 )
 set(_mlir_typestub_gen_target "${NB_STUBGEN_CUSTOM_TARGET}")
 
@@ -196,7 +196,7 @@ set(_SOURCE_COMPONENTS
 
 add_mlir_python_modules(ASTERPythonModules
   ROOT_PREFIX "${ASTER_BINARY_DIR}/python_packages/aster"
-  INSTALL_PREFIX "python_packages/aster"
+  INSTALL_PREFIX "${_PYTHON_INSTALL_PREFIX}"
   DECLARED_SOURCES
   ${_SOURCE_COMPONENTS}
   COMMON_CAPI_LINK_LIBS
@@ -209,10 +209,10 @@ add_mlir_python_modules(ASTERPythonModules
 # This uses the ASTER target created in tools/aster-shlib
 install(TARGETS ASTER
   LIBRARY
-    DESTINATION python_packages/aster/_mlir_libs
+    DESTINATION "${_PYTHON_INSTALL_PREFIX}/_mlir_libs"
     COMPONENT ASTERPythonModules
   RUNTIME
-    DESTINATION python_packages/aster/_mlir_libs
+    DESTINATION "${_PYTHON_INSTALL_PREFIX}/_mlir_libs"
     COMPONENT ASTERPythonModules
 )
 

tools/aster-opt/CMakeLists.txt

@@ -13,4 +13,12 @@ add_llvm_executable(aster-opt
 llvm_update_compile_flags(aster-opt)
 target_link_libraries(aster-opt PRIVATE ${LIBS})
 mlir_check_all_link_libraries(aster-opt)
-install(TARGETS aster-opt)
+if(DEFINED SKBUILD_PLATLIB_DIR)
+  install(TARGETS aster-opt
+    RUNTIME
+      DESTINATION "${ASTER_TOOLS_INSTALL_PREFIX}"
+      COMPONENT ASTERPythonModules
+  )
+else()
+  install(TARGETS aster-opt)
+endif()

tools/aster-translate/CMakeLists.txt

@@ -15,4 +15,12 @@ add_llvm_executable(aster-translate
 llvm_update_compile_flags(aster-translate)
 target_link_libraries(aster-translate PRIVATE ${LIBS})
 mlir_check_all_link_libraries(aster-translate)
-install(TARGETS aster-translate)
+if(DEFINED SKBUILD_PLATLIB_DIR)
+  install(TARGETS aster-translate
+    RUNTIME
+      DESTINATION "${ASTER_TOOLS_INSTALL_PREFIX}"
+      COMPONENT ASTERPythonModules
+  )
+else()
+  install(TARGETS aster-translate)
+endif()

tools/pkgs/aster/README.md

@@ -0,0 +1,111 @@
+# aster Python wheel packaging
+
+This directory contains the wheel build infrastructure for the `aster` Python package.
+
+| File | Purpose |
+|---|---|
+| `build.sh` | Orchestrates a `cibuildwheel` run; stages the LLVM tarball |
+| `before-all.sh` | `cibuildwheel` hook; extracts LLVM and installs build tools inside the container |
+
+Build system and package metadata live in the root [`pyproject.toml`](../../../pyproject.toml).
+
+---
+
+## Editable install (local development)
+
+Requires `CMAKE_PREFIX_PATH` to point to the LLVM cmake directory so that `find_package(MLIR)` succeeds.
+
+```bash
+# Activate the project sandbox first
+source <worktree>/sandbox/bin/activate_sandbox
+
+# Install (builds C++ extensions via scikit-build-core + CMake)
+CMAKE_PREFIX_PATH=~/shared-llvm/lib/cmake pip install -e .
+
+# With dev extras (lit, pytest, black, pyright, …)
+CMAKE_PREFIX_PATH=~/shared-llvm/lib/cmake pip install -e .[dev]
+```
+
+To pass extra CMake flags (e.g. override build type or set a custom install prefix):
+
+```bash
+CMAKE_PREFIX_PATH=~/shared-llvm/lib/cmake \
+  pip install -e . \
+  --config-settings=cmake.build-type=Debug \
+  --config-settings="cmake.args=-DASTER_ENABLE_UNIT_TESTS=ON"
+```
+
+After installation `import aster` works from any directory without modifying `PYTHONPATH`.
+
+---
+
+## Release wheels (cibuildwheel)
+
+`build.sh` stages the LLVM installation as a tarball in the project root and runs `cibuildwheel` from there. Docker is required on Linux.
+
+### Building the LLVM tarball
+
+A pre-built LLVM tarball is required. See [`tools/pkgs/llvm/README.md`](../llvm/README.md) for full details. Quick start:
+
+```bash
+# Linux (requires Docker); produces llvm-linux-x86_64-<commit>.tar.gz
+bash tools/pkgs/llvm/build.sh --output=/tmp/llvm-out
+
+# macOS (native); produces llvm-macos-arm64-<commit>.tar.gz
+bash tools/pkgs/llvm/build.sh --output=/tmp/llvm-out
+```
+
+### From a pre-built LLVM tarball
+
+```bash
+bash tools/pkgs/aster/build.sh --llvm-tarball=~/llvm-out/llvm-linux-x86_64-<commit>.tar.gz
+```
+
+### From an unpacked LLVM directory
+
+`build.sh` re-archives it into the expected `llvm/` layout automatically.
+
+```bash
+bash tools/pkgs/aster/build.sh --llvm-dir=~/shared-llvm
+```
+
+Wheels land in `./wheelhouse/` by default.
+
+### Common options
+
+| Option | Default | Description |
+|---|---|---|
+| `--llvm-tarball=PATH` | — | Pre-built LLVM `.tar.gz` |
+| `--llvm-dir=PATH` | — | Unpacked LLVM installation directory |
+| `--output=DIR` | `./wheelhouse` | Destination for produced wheels |
+| `--arch=ARCH` | host arch | `x86_64`, `aarch64`, or `arm64` |
+| `--jobs=N` | — | Parallelism forwarded to the build |
+
+### Cross-compilation (aarch64 on x86_64)
+
+Requires Docker with QEMU binfmt support (`docker run --platform linux/arm64` must work).
+
+```bash
+bash tools/pkgs/aster/build.sh \
+  --llvm-tarball=~/llvm-out/llvm-linux-aarch64-<commit>.tar.gz \
+  --arch=aarch64
+```
+
+---
+
+## Tarball layout
+
+`build.sh` and `before-all.sh` expect the LLVM tarball to have a top-level `llvm/` directory produced by `tools/pkgs/llvm/build.sh`:
+
+```
+llvm/
+  bin/
+  include/
+  lib/
+    cmake/
+      llvm/
+      mlir/
+      lld/
+```
+
+`before-all.sh` locates `lib/cmake` inside the extracted tree and symlinks it to `/opt/llvm` (Linux) or `~/.llvm-pkg` (macOS) so CMake's `find_package()` can resolve it via `CMAKE_PREFIX_PATH`.

tools/pkgs/aster/before-all.sh

@@ -0,0 +1,70 @@
+#!/usr/bin/env bash
+# before-all.sh - cibuildwheel before-all hook for the aster Python wheel.
+#
+# Extracts the LLVM tarball and installs build tools.  Runs once per build
+# environment before any Python version is built.
+#
+# Required env var:
+#   LLVM_PKG_PATH  Path to the LLVM tar.gz, relative to the project root.
+set -euo pipefail
+
+OS="$(uname -s)"
+
+PROJECT_DIR="$(pwd)"
+LLVM_TARBALL="${PROJECT_DIR}/${LLVM_PKG_PATH:?LLVM_PKG_PATH must be set}"
+
+if [[ ! -f "${LLVM_TARBALL}" ]]; then
+  echo "error: LLVM tarball not found: ${LLVM_TARBALL}" >&2
+  exit 1
+fi
+
+if [[ "${OS}" == "Linux" ]]; then
+  LLVM_INSTALL="/opt/llvm"
+else
+  LLVM_INSTALL="${HOME}/.llvm-pkg"
+fi
+
+# Extract to a stable directory and locate lib/cmake, then symlink
+# LLVM_INSTALL -> lib/cmake so that CMAKE_PREFIX_PATH resolves
+# <PackageName>/Config.cmake regardless of the tarball's internal layout.
+# The extract directory must persist for the lifetime of the build since
+# LLVM_INSTALL is a symlink into it.
+_EXTRACT_DIR="${PROJECT_DIR}/.llvm-extract"
+rm -rf "${_EXTRACT_DIR}"
+mkdir -p "${_EXTRACT_DIR}"
+
+echo "Extracting ${LLVM_TARBALL} ..."
+tar -xzf "${LLVM_TARBALL}" -C "${_EXTRACT_DIR}"
+
+_CMAKE_DIR="$(set +o pipefail; find "${_EXTRACT_DIR}" -maxdepth 8 -type d -name "cmake" -path "*/lib/cmake" | head -n1)"
+if [[ -z "${_CMAKE_DIR}" ]]; then
+  echo "error: lib/cmake directory not found in tarball" >&2
+  exit 1
+fi
+echo "LLVM cmake dir: ${_CMAKE_DIR}"
+
+ln -sfn "${_CMAKE_DIR}" "${LLVM_INSTALL}"
+echo "LLVM cmake dir linked at ${LLVM_INSTALL}."
+
+if [[ "${OS}" == "Linux" ]]; then
+  echo "Installing build tools ..."
+  if command -v dnf >/dev/null 2>&1; then
+    dnf install -y clang lld cmake ninja-build
+  elif command -v apk >/dev/null 2>&1; then
+    apk add --no-cache clang lld cmake ninja
+  fi
+  echo "Installing nanobind ..."
+  pip install "nanobind>=2.9,<3.0"
+elif [[ "${OS}" == "Darwin" ]]; then
+  for _tool in cmake ninja clang; do
+    if ! command -v "${_tool}" >/dev/null 2>&1; then
+      echo "error: required tool not found: ${_tool}" >&2
+      echo "Install it via Homebrew: brew install ${_tool}" >&2
+      exit 1
+    fi
+  done
+  echo "Installing nanobind ..."
+  pip install "nanobind>=2.9,<3.0"
+fi
+
+echo "before-all.sh complete."