[build] Support build python wheel with scikit-build-core

CMakeLists.txt

@@ -216,23 +216,6 @@ if(BUDDY_MLIR_ENABLE_RISCV_GNU_TOOLCHAIN)
   )
 endif()
 
-#-------------------------------------------------------------------------------
-# Initialize Python packages
-#-------------------------------------------------------------------------------
-if(BUDDY_MLIR_ENABLE_PYTHON_PACKAGES)
-  # Find the Python interpreter and development components,
-  # requiring a minimum version of 3.10
-  find_package(Python3 3.10 REQUIRED COMPONENTS Interpreter Development)
-  # Create directories for the BUDDY-MLIR Python packages
-  file(MAKE_DIRECTORY ${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy)
-  file(MAKE_DIRECTORY ${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy/compiler)
-  # Create empty __init__.py files to make these directories Python packages
-  file(WRITE ${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy/__init__.py "")
-  file(WRITE ${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy/compiler/__init__.py "")
-
-  install(DIRECTORY ${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy DESTINATION python_packages)
-endif()
-
 #-------------------------------------------------------------------------------
 # Directory setup
 #-------------------------------------------------------------------------------

README.md

@@ -81,6 +81,33 @@ $ cmake -G Ninja .. -DBUDDY_ENABLE_E2E_TESTS=ON
 $ ninja check-e2e
 ```
 
+## Build Python Package
+
+We use `scikit-build-core` + CMake to produce a wheel that contains both the
+`buddy` front-end Python package and the `buddy_mlir` bindings.
+
+```bash
+# From repo root
+export MLIR_PREFIX=$(pwd)/llvm/build
+export CMAKE_ARGS="-DMLIR_DIR=${MLIR_PREFIX}/lib/cmake/mlir \
+  -DLLVM_DIR=${MLIR_PREFIX}/lib/cmake/llvm \
+  -DCMAKE_PREFIX_PATH=${MLIR_PREFIX}/lib/cmake \
+  -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
+  -DBUDDY_MLIR_ENABLE_PYTHON_PACKAGES=ON \
+  -DBUDDY_MLIR_PYTHON_INSTALL_PREFIX=. \
+  -DPython3_EXECUTABLE=$(which python)"
+export SKBUILD_CONFIGURE_OPTIONS="$CMAKE_ARGS"
+python -m build -w -o build/dist
+```
+
+Install and test the wheel:
+
+```bash
+pip install torch --index-url https://download.pytorch.org/whl/cpu
+pip install build/dist/buddy-*.whl --no-deps
+python -c "import buddy; import buddy_mlir; print('ok')"
+```
+
 ## Examples
 
 We provide examples to demonstrate how to use the passes and interfaces in `buddy-mlir`, including IR-level transformations, domain-specific applications, and testing demonstrations.

frontend/Python/CMakeLists.txt

@@ -1,11 +1,40 @@
+# Find the Python interpreter and development components,
+# requiring a minimum version of 3.10
+find_package(Python3 3.10 REQUIRED COMPONENTS Interpreter Development)
+
+# Create directories for the BUDDY-MLIR Python packages
+file(MAKE_DIRECTORY ${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy)
+file(MAKE_DIRECTORY ${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy/compiler)
+# Create empty __init__.py files to make these directories Python packages
+file(WRITE ${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy/__init__.py "")
+file(WRITE ${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy/compiler/__init__.py "")
+
 # Recursively retrieve all python files from the current directory.
 file(GLOB_RECURSE ALL_PY_FILES RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}" "*.py")
 
+set(PYTHON_OUTPUT_FILES "")
 foreach(FILE ${ALL_PY_FILES})
-    # Get the directory of the current file.
-    get_filename_component(DIR "${FILE}" DIRECTORY)
-    # Set the destination directory for the target file.
-    set(DEST "${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy/compiler/${DIR}")
-    # Copy the file into the destination directory.
-    file(COPY ${FILE} DESTINATION ${DEST})
+    get_filename_component(REL_DIR "${FILE}" DIRECTORY)
+    set(SRC "${CMAKE_CURRENT_SOURCE_DIR}/${FILE}")
+    set(DEST_DIR "${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy/compiler/${REL_DIR}")
+    set(DEST_FILE "${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy/compiler/${FILE}")
+
+    add_custom_command(
+        OUTPUT "${DEST_FILE}"
+        COMMAND ${CMAKE_COMMAND} -E make_directory "${DEST_DIR}"
+        COMMAND ${CMAKE_COMMAND} -E copy_if_different "${SRC}" "${DEST_FILE}"
+        DEPENDS "${SRC}"
+        COMMENT "Syncing ${FILE}"
+        VERBATIM
+    )
+
+    list(APPEND PYTHON_OUTPUT_FILES "${DEST_FILE}")
 endforeach()
+
+add_custom_target(python-package
+    ALL DEPENDS ${PYTHON_OUTPUT_FILES}
+    COMMENT "Syncing python package files"
+)
+
+# Install the front-end Python package into the install prefix root.
+install(DIRECTORY ${BUDDY_MLIR_PYTHON_PACKAGES_DIR}/buddy DESTINATION .)

frontend/Python/frontend.py

@@ -28,11 +28,11 @@
 import ctypes
 import platform
 
-import mlir.ir as ir
-import mlir.dialects.func as func
-from mlir.passmanager import *
-from mlir.execution_engine import *
-from mlir import runtime as rt
+import buddy_mlir.ir as ir
+import buddy_mlir.dialects.func as func
+from buddy_mlir.passmanager import *
+from buddy_mlir.execution_engine import *
+from buddy_mlir import runtime as rt
 import torch
 import torch._dynamo as dynamo
 from torch._functorch.aot_autograd import aot_module_simplified

frontend/Python/graph/graph.py

@@ -24,11 +24,11 @@
 import functools
 import numpy as np
 
-import mlir.ir as ir
-import mlir.dialects.func as func
-from mlir.passmanager import *
-from mlir.execution_engine import *
-from mlir import runtime as rt
+import buddy_mlir.ir as ir
+import buddy_mlir.dialects.func as func
+from buddy_mlir.passmanager import *
+from buddy_mlir.execution_engine import *
+from buddy_mlir import runtime as rt
 
 from .operation import *
 from .type import *
@@ -460,7 +460,7 @@ class GraphImporter:
         _func_name (str): Name of the generated MLIR function.
         _inputs (List[TensorMeta]): Input tensor(s) of the FX graph.
         _num_input_visited (int): Number of input nodes that have been visited.
-        _module (mlir.ir.Module): The generated MLIR module.
+        _module (buddy_mlir.ir.Module): The generated MLIR module.
         _ops_registry (dict): Registry for the candidate operations.
     """
 
@@ -511,7 +511,7 @@ def _str_to_mlir_dtype(self, dtype: str) -> ir.Type:
             dtype (str): The tensor type.
 
         Returns:
-            mlir.ir.Type: The corresponding MLIR data type.
+            buddy_mlir.ir.Type: The corresponding MLIR data type.
 
         Raises:
             NotImplementedError: If the given dtype is not supported.
@@ -568,7 +568,7 @@ def import_graph(self) -> ir.Module:
         Imports buddy graph and generates an MLIR module in high-level dialects.
 
         Returns:
-            mlir.ir.Module: An MLIR module in high-level dialects.
+            buddy_mlir.ir.Module: An MLIR module in high-level dialects.
         """
         assert self._do_param_pack == False
         with ir.InsertionPoint(self._module.body):
@@ -643,7 +643,7 @@ def import_main_graph(self) -> ir.Module:
         module in high-level dialects with memref.
 
         Returns:
-            mlir.ir.Module: An MLIR module in high-level dialects.
+            buddy_mlir.ir.Module: An MLIR module in high-level dialects.
         """
         with ir.InsertionPoint(self._module.body):
             arguments = []
@@ -702,7 +702,7 @@ def _import_placeholder(
         Parameters:
         - node (PlaceholderOp): The PlaceholderOp node representing the
         placeholder.
-        - args_list (List[mlir.ir.BlockArgument]): List of input memrefs.
+        - args_list (List[buddy_mlir.ir.BlockArgument]): List of input memrefs.
 
         Returns:
         None

frontend/Python/graph/graph_driver.py

@@ -20,7 +20,7 @@
 #
 # ===---------------------------------------------------------------------------
 
-from mlir import ir
+from buddy_mlir import ir
 from collections import deque, defaultdict
 
 from .graph import Graph, GraphImporter, TensorMeta

frontend/Python/ops/func.py

@@ -20,8 +20,8 @@
 
 from typing import Tuple
 import functools
-from mlir.dialects import func, memref
-from mlir import ir
+from buddy_mlir.dialects import func, memref
+from buddy_mlir import ir
 from ..graph import FuncOp, CallOp, CallExternalOp, PlaceholderOp
 from .utils import *
 

frontend/Python/ops/linalg.py

@@ -20,8 +20,8 @@
 
 from typing import Dict, Tuple, List
 
-import mlir.ir as ir
-from mlir.dialects import (
+import buddy_mlir.ir as ir
+from buddy_mlir.dialects import (
     tosa,
     linalg,
     arith,

frontend/Python/ops/math.py

@@ -18,7 +18,7 @@
 #
 # ===---------------------------------------------------------------------------
 
-from mlir.dialects import math
+from buddy_mlir.dialects import math
 
 
 def erf_op(node, symbol_table):
@@ -209,7 +209,7 @@ def trunc_op(node, symbol_table):
 
 def abs_op(node, symbol_table):
     """abs(x) using math.AbsFOp for float, math.AbsIOp for int"""
-    import mlir.ir as ir
+    import buddy_mlir.ir as ir
 
     input_tensor = symbol_table.get((str(node.args[0]), 0))
     element_type = ir.RankedTensorType(input_tensor.type).element_type

frontend/Python/ops/tosa.py

@@ -23,9 +23,9 @@
 import numpy
 import sys
 
-import mlir.ir as ir
-from mlir.ir import IndexType, F32Type
-from mlir.dialects import (
+import buddy_mlir.ir as ir
+from buddy_mlir.ir import IndexType, F32Type
+from buddy_mlir.dialects import (
     tensor,
     tosa,
     arith,

frontend/Python/ops/utils.py

@@ -19,7 +19,7 @@
 # ===---------------------------------------------------------------------------
 
 from typing import Dict
-import mlir.ir as ir
+import buddy_mlir.ir as ir
 
 from ..graph import TensorDType
 

midend/python/CMakeLists.txt

@@ -105,7 +105,7 @@ declare_mlir_python_extension(BuddyMLIRPythonSources.Extension
 
 add_mlir_python_common_capi_library(BuddyMLIRPythonCAPI
   INSTALL_COMPONENT BuddyMLIRPythonModules
-  INSTALL_DESTINATION python_packages/buddy_mlir/_mlir_libs
+  INSTALL_DESTINATION buddy_mlir/_mlir_libs
   OUTPUT_DIRECTORY "${BUDDY_BUILD_DIR}/python_packages/buddy_mlir/_mlir_libs"
   RELATIVE_INSTALL_ROOT "../../.."
   DECLARED_SOURCES
@@ -121,7 +121,7 @@ link_directories(${LLVM_BINARY_DIR}/tools/mlir/python_packages/mlir_core/mlir/_m
 
 add_mlir_python_modules(BuddyMLIRPythonModules
   ROOT_PREFIX "${BUDDY_BUILD_DIR}/python_packages/buddy_mlir"
-  INSTALL_PREFIX "python_packages/buddy_mlir"
+  INSTALL_PREFIX "buddy_mlir"
   DECLARED_SOURCES
     BuddyMLIRPythonSources
     MLIRPythonSources

pyproject.toml

@@ -1,2 +1,27 @@
 [tool.black]
 line-length = 80
+
+[build-system]
+requires = [
+  "scikit-build-core>=0.7",
+  "build>=1.2",
+  "cmake>=3.22",
+  "ninja",
+  "setuptools>=61",
+  "wheel",
+  # required at build time for midend/python
+  "numpy<2",
+  "pybind11==2.11.1",
+  "nanobind>=2.4",
+]
+build-backend = "scikit_build_core.build"
+
+[project]
+name = "buddy"
+version = "0.0.1"
+description = "An MLIR-based compiler framework bridges DSLs (domain-specific languages) to DSAs (domain-specific architectures)."
+requires-python = ">=3.10"
+
+[tool.scikit-build]
+cmake.minimum-version = "3.22"
+wheel.expand-macos-universal-tags = false

tests/CMakeLists.txt

@@ -59,6 +59,7 @@ if(BUDDY_ENABLE_E2E_TESTS)
   buddy-translate
   mlir-runner
   buddy-lenet-run-test-cpu
+  python-package
   )
 
   add_lit_testsuite(check-e2e "Running E2E tests for Models directory..."