pypi: update setup.py and pyproject.toml for cibuildwheel (#4639)

Corresponds to llvm/torch-mlir-release#32

Cf. #4603

This can be tested locally with:

```
python3 -m venv venv
source venv/bin/activate
pip install cibuildwheel uv
CIBW_BUILD="cp312-manylinux_x86_64" cibuildwheel --platform linux
```

I also went ahead and added abi3 support for better cross-platform
compatibility (in particular, the gLinux box I was testing it on could
not install the wheel without abi3, though I didn't dig into why).
However, that decision adds a constraint: LLVM pins the stable ABI to
Python3.12, so that would limit this wheel to python3.12 or later. Is
that OK? Python 3.11 is EOL on October 31, 2027, so there is an argument
to be made that it would be useful to keep support for it around until
then. But if LLVM as a whole agreed on python3.12 (I wouldn't know why)
then maybe it's better to be consistent with that.

---------

Co-authored-by: Jeremy Kun <j2kun@users.noreply.github.com>

Author: j2kun

CMakeLists.txt

@@ -21,6 +21,10 @@ if(POLICY CMP0116)
   cmake_policy(SET CMP0116 OLD)
 endif()
 
+if(POLICY CMP0079)
+  cmake_policy(SET CMP0079 NEW)
+endif()
+
 project(torch-mlir LANGUAGES CXX C)
 set(CMAKE_C_STANDARD 11)
 set(CMAKE_CXX_STANDARD 17)

build_tools/smoke_test.py

@@ -0,0 +1,33 @@
+# Part of the LLVM Project, 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
+# Also available under a BSD-style license. See LICENSE.
+
+import torch
+import torch_mlir
+import torch_mlir.fx
+
+
+class MLPModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = torch.nn.Linear(4, 4)
+        self.relu = torch.nn.ReLU()
+
+    def forward(self, x):
+        return self.relu(self.linear(x))
+
+
+module_to_compile = MLPModule()
+backends = ["linalg-on-tensors", "tosa", "stablehlo"]
+
+for backend in backends:
+    print(f"Compiling MLPModule via FX to {backend}...")
+    try:
+        torch_mlir.fx.export_and_import(
+            module_to_compile, torch.ones(2, 4), output_type=backend
+        )
+        print(f"Compilation to {backend} successful!")
+    except Exception as e:
+        print(f"Compilation to {backend} FAILED!")
+        raise e

pyproject.toml

@@ -1,7 +1,60 @@
 [build-system]
-requires = ["setuptools", "wheel"]
+# These must be enough to drive setup.py's CMake build of the MLIR Python
+# bindings under build isolation (cibuildwheel). Kept in sync with the
+# [dependency-groups].build table below and
+# externals/llvm-project/mlir/python/requirements.txt (re-sync on LLVM bumps).
+requires = [
+    "setuptools",
+    "wheel",
+    "cmake",
+    "ninja<1.13.0; sys_platform == 'win32'",
+    "ninja; sys_platform != 'win32'",
+    "packaging",
+    "nanobind>=2.9, <3.0",
+    "PyYAML>=5.4.0, <=6.0.1",
+    "typing_extensions>=4.12.2",
+    "numpy>=2.1.0, <=2.1.2",
+    "ml_dtypes>=0.5.0, <=0.6.0",
+]
 build-backend = "setuptools.build_meta"
 
 [tool.black]
 line-length = 88
 target-version = ['py38']
+
+# NOTE: torch-mlir's Python package is produced by the CMake build (see the
+# custom `build_py` in setup.py), not by setuptools' source-tree discovery.
+# Auto-discovery is disabled via `packages=[]` passed directly in setup.py.
+# We intentionally do NOT define a `[tool.setuptools]` (or `[project]`) table
+# here: doing so makes modern setuptools treat pyproject.toml as the metadata
+# source and reject the setup.py-driven build with
+# "`project` must contain ['version'] properties".
+
+[tool.cibuildwheel]
+build = "cp310-* cp311-* cp312-*"
+skip = "*-musllinux* *_i686 *-win32"
+build-frontend = "build[uv]"
+test-requires = ["-r pytorch-requirements.txt"]
+test-command = "python {project}/build_tools/smoke_test.py"
+
+# Build with the manylinux default gcc toolchain. ccache reuses the LLVM/MLIR
+# object cache across the per-Python-version builds; CMake reads the
+# *_COMPILER_LAUNCHER cache vars from these env vars. CCACHE_DIR lives inside
+# the mounted /project so the host's actions/cache can persist it.
+[tool.cibuildwheel.linux]
+archs = "auto64"
+manylinux-x86_64-image = "manylinux_2_28"
+manylinux-aarch64-image = "manylinux_2_28"
+before-all = "dnf install -y --quiet epel-release || true && dnf install -y --quiet ccache || dnf install -y ccache"
+environment-pass = ["TORCH_MLIR_PYTHON_PACKAGE_VERSION"]
+environment = { TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS = "1", TORCH_MLIR_ENABLE_JIT_IR_IMPORTER = "0", TORCH_MLIR_ENABLE_LTC = "0", CMAKE_GENERATOR = "Ninja", CMAKE_C_COMPILER_LAUNCHER = "ccache", CMAKE_CXX_COMPILER_LAUNCHER = "ccache", CCACHE_DIR = "/project/.ccache", CCACHE_MAXSIZE = "2G" }
+
+[tool.cibuildwheel.macos]
+archs = "auto64"
+before-all = "brew install ccache ninja"
+environment = { MACOSX_DEPLOYMENT_TARGET = "11.0", TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS = "1", TORCH_MLIR_ENABLE_JIT_IR_IMPORTER = "0", TORCH_MLIR_ENABLE_LTC = "0", CMAKE_GENERATOR = "Ninja", CMAKE_C_COMPILER_LAUNCHER = "ccache", CMAKE_CXX_COMPILER_LAUNCHER = "ccache" }
+
+[tool.cibuildwheel.windows]
+before-build = "pip install delvewheel"
+repair-wheel-command = "delvewheel repair --add-path {project}/build/cmake_build/tools/torch-mlir/python_packages/torch_mlir/torch_mlir/_mlir_libs --add-dll TorchMLIRAggregateCAPI.dll --no-dll \"c10.dll;torch_python.dll;torch_cpu.dll\" -w {dest_dir} {wheel}"
+environment = { TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS = "1", TORCH_MLIR_ENABLE_JIT_IR_IMPORTER = "0", TORCH_MLIR_ENABLE_LTC = "0", CMAKE_GENERATOR = "Ninja", TORCH_MLIR_CMAKE_BUILD_DIR = "build/cmake_build" }

setup.py

@@ -39,6 +39,14 @@
 # CMAKE_GENERATOR=Ninja CMAKE_C_COMPILER_LAUNCHER=ccache CMAKE_CXX_COMPILER_LAUNCHER=ccache
 # ```
 #
+# For release builds, use cibuildwheel. Except for
+# TORCH_MLIR_PYTHON_PACKAGE_VERSION, the environment variables needed for
+# cibuildwheel are set in pyproject.toml. For an example local run:
+#
+# ```
+# CIBW_BUILD="cp312-manylinux_x86_64" cibuildwheel --platform linux
+# ```
+#
 # Implementation notes:
 # The contents of the wheel is just the contents of the `python_packages`
 # directory that our CMake build produces. We go through quite a bit of effort
@@ -104,6 +112,9 @@ def run(self):
         self.run_command("build_scripts")
 
 
+use_stable_abi = sys.version_info >= (3, 12)
+
+
 class CMakeBuild(build_py):
     def cmake_build(self, cmake_build_dir):
         llvm_dir = str(SRC_DIR / "externals" / "llvm-project" / "llvm")
@@ -125,6 +136,7 @@ def cmake_build(self, cmake_build_dir):
             f"-DCMAKE_CXX_VISIBILITY_PRESET=hidden",
             f"-DTORCH_MLIR_ENABLE_LTC={'ON' if TORCH_MLIR_ENABLE_LTC else 'OFF'}",
             f"-DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS={'OFF' if TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS else 'ON'}",
+            f"-DMLIR_ENABLE_PYTHON_STABLE_ABI={'ON' if use_stable_abi else 'OFF'}",
         ]
         if LLVM_INSTALL_DIR:
             cmake_config_args += [
@@ -214,7 +226,7 @@ def run(self):
 
 class CMakeExtension(Extension):
     def __init__(self, name, sourcedir=""):
-        Extension.__init__(self, name, sources=[])
+        Extension.__init__(self, name, sources=[], py_limited_api=use_stable_abi)
         self.sourcedir = os.path.abspath(sourcedir)
 
 
@@ -255,6 +267,10 @@ def build_extension(self, ext):
     )
 
 
+setup_options = {}
+if use_stable_abi:
+    setup_options["bdist_wheel"] = {"py_limited_api": "cp312"}
+
 setup(
     name=NAME,
     version=f"{PACKAGE_VERSION}",
@@ -266,10 +282,16 @@ def build_extension(self, ext):
     include_package_data=True,
     cmdclass={
         "build": CustomBuild,
-        "built_ext": NoopBuildExtension,
+        "build_ext": NoopBuildExtension,
         "build_py": CMakeBuild,
     },
     ext_modules=EXT_MODULES,
+    # The Python package contents are placed into the build dir by the custom
+    # CMake `build_py` above; there are no source-tree packages to discover.
+    # An explicit empty list disables setuptools auto-discovery (which would
+    # otherwise trip over top-level non-Python dirs like lib/, include/,
+    # projects/, ...).
+    packages=[],
     python_requires=">=3.8",
     install_requires=INSTALL_REQUIRES,
     extras_require={
@@ -284,4 +306,5 @@ def build_extension(self, ext):
         ],
     },
     zip_safe=False,
+    options=setup_options,
 )