Python wheel packaging for ptoas, in manylinux docker image

CMakeLists.txt

@@ -56,7 +56,7 @@ add_subdirectory(tools)
 option(PTO_ENABLE_PYTHON_BINDING "Enable Python bindings" ON)
 
 if(PTO_ENABLE_PYTHON_BINDING)
-  find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
+  find_package(Python3 COMPONENTS Interpreter Development.Module REQUIRED)
   find_package(pybind11 CONFIG REQUIRED)
 
   # 设置 Python 扩展的输出目录，方便调试

README.md

@@ -97,7 +97,7 @@ cmake -G Ninja -S llvm -B $LLVM_BUILD_DIR \
     -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
     -DPython3_EXECUTABLE=$(which python3) \
     -DCMAKE_BUILD_TYPE=Release \
-    -DLLVM_TARGETS_TO_BUILD="Host"
+    -DLLVM_TARGETS_TO_BUILD="host"
 
 # 4. 编译 LLVM (这一步耗时较长)
 ninja -C $LLVM_BUILD_DIR
@@ -124,6 +124,8 @@ cmake -G Ninja \
     -B build \
     -DLLVM_DIR=$LLVM_BUILD_DIR/lib/cmake/llvm \
     -DMLIR_DIR=$LLVM_BUILD_DIR/lib/cmake/mlir \
+    -DPython3_EXECUTABLE=$(which python3) \
+    -DPython3_FIND_STRATEGY=LOCATION \
     -Dpybind11_DIR="${PYBIND11_CMAKE_DIR}" \
     -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
     -DMLIR_PYTHON_PACKAGE_DIR=$LLVM_BUILD_DIR/tools/mlir/python_packages/mlir_core \
@@ -174,7 +176,7 @@ export PYTHONPATH=$MLIR_PYTHON_ROOT:$PTO_PYTHON_ROOT:$PYTHONPATH
 export LD_LIBRARY_PATH=$LLVM_BUILD_DIR/lib:$PTO_INSTALL_DIR/lib:$LD_LIBRARY_PATH
 
 # 3. PATH: 将 ptoas 添加到命令行路径
-export PATH=$PTO_INSTALL_DIR/bin:$PATH
+export PATH=$PTO_SOURCE_DIR/build/tools/ptoas:$PATH
 
 ```
 
@@ -213,11 +215,11 @@ with Context() as ctx, Location.unknown():
 
 ```bash
 # 运行python binding 测试
-python3 ./test/samples/MatMul/tmatmulk.py > ./test/samples/MatMul/tmatmulk.pto
+cd ./test/samples/MatMul/
+python3 ./tmatmulk.py > ./tmatmulk.pto
 
 # 运行ptoas 测试
-./build/tools/ptoas/ptoas ./test/samples/Matmul/tmatmulk.pto -o ./test/samples/Matmul/tmatmulk.cpp
-
+./build/tools/ptoas/ptoas ./tmatmulk.pto -o ./tmatmulk.cpp
 ```
 
 ---

docker/Dockerfile

@@ -0,0 +1,95 @@
+FROM quay.io/pypa/manylinux_2_34_x86_64
+
+# NOTE: change $PY_VER for different Python versions (3.8 - 3.14 available)
+ARG PY_VER=cp312-cp312
+ARG LLVM_TAG=llvmorg-19.1.7
+
+## -- usually no need to change below --
+
+ENV PY_PATH="/opt/python/${PY_VER}"
+ENV PATH="${PY_PATH}/bin:${PATH}"
+
+# dependency
+RUN dnf install -y ninja-build cmake git ccache gcc-c++ lld && dnf clean all
+RUN pip install --no-cache-dir numpy pybind11 nanobind
+
+# setting build directories
+ENV WORKSPACE_DIR=/llvm-workspace
+ENV LLVM_SOURCE_DIR=$WORKSPACE_DIR/llvm-project
+ENV LLVM_BUILD_DIR=$LLVM_SOURCE_DIR/build-shared
+
+ENV PTO_SOURCE_DIR=$WORKSPACE_DIR/PTOAS
+ENV PTO_INSTALL_DIR=$PTO_SOURCE_DIR/install
+
+# build LLVM
+WORKDIR $WORKSPACE_DIR
+RUN git clone --depth 1 --branch ${LLVM_TAG} https://github.com/llvm/llvm-project.git
+
+WORKDIR $LLVM_SOURCE_DIR
+RUN cmake -G Ninja -S llvm -B $LLVM_BUILD_DIR \
+    -DLLVM_ENABLE_PROJECTS="mlir;clang" \
+    -DBUILD_SHARED_LIBS=ON \
+    -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
+    -DPython3_EXECUTABLE=${PY_PATH}/bin/python \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DLLVM_TARGETS_TO_BUILD="host"
+
+RUN ninja -C $LLVM_BUILD_DIR
+
+# build ptoas
+WORKDIR $WORKSPACE_DIR
+RUN git clone -b wheel https://github.com/learning-chip/PTOAS.git
+# TODO: tag git commit of PTOAS repo
+
+WORKDIR $PTO_SOURCE_DIR
+RUN cmake -G Ninja \
+    -S . \
+    -B build \
+    -DLLVM_DIR=$LLVM_BUILD_DIR/lib/cmake/llvm \
+    -DMLIR_DIR=$LLVM_BUILD_DIR/lib/cmake/mlir \
+    -DPython3_ROOT_DIR=${PY_PATH} \
+    -DPython3_EXECUTABLE=${PY_PATH}/bin/python \
+    -DPython3_FIND_STRATEGY=LOCATION \
+    -Dpybind11_DIR=$(python -m pybind11 --cmakedir) \
+    -DMLIR_PYTHON_PACKAGE_DIR=${LLVM_BUILD_DIR}/tools/mlir/python_packages/mlir_core \
+    -DCMAKE_INSTALL_PREFIX=${PTO_INSTALL_DIR}
+
+RUN ninja -C build && ninja -C build install
+
+# test ptoas cli
+ENV PATH=$PTO_SOURCE_DIR/build/tools/ptoas:$PATH
+RUN which ptoas
+
+# create python wheel
+ENV PY_PACKAGE_DIR=$LLVM_BUILD_DIR/tools/mlir/python_packages/mlir_core/
+
+# copy pto.py, _pto_ops_gen.py 
+RUN cp $PTO_INSTALL_DIR/mlir/dialects/*.py $PY_PACKAGE_DIR/mlir/dialects/
+
+COPY ./setup.py $PY_PACKAGE_DIR/
+
+WORKDIR $PY_PACKAGE_DIR
+RUN pip install --no-cache-dir setuptools wheel auditwheel
+RUN python setup.py bdist_wheel
+
+# fix missing so
+RUN export LD_LIBRARY_PATH=$LLVM_BUILD_DIR/lib:$PTO_INSTALL_DIR/lib:$LD_LIBRARY_PATH \
+    && auditwheel repair dist/ptoas*.whl
+
+RUN pip install wheelhouse/ptoas*.whl
+
+# test import without adjusting $PYTHONPATH and $LD_LIBRARY_PATH, in some irrelevant dir
+WORKDIR /test
+RUN python -c "import mlir.ir"
+RUN python -c "from mlir.dialects import pto"
+
+WORKDIR $PTO_SOURCE_DIR/test/samples/MatMul/
+RUN python ./tmatmulk.py > ./tmatmulk.pto
+RUN ptoas ./tmatmulk.pto -o ./tmatmulk.cpp
+# TODO: test compilation of the cpp file using `bisheng`
+# TODO: also test `ptoas --enable-insert-sync`
+
+# show the built wheel by default
+WORKDIR $PY_PACKAGE_DIR/dist
+
+CMD ["/bin/bash"]

docker/README.md

@@ -0,0 +1,14 @@
+Build:
+
+```bash
+docker build . -t ptoas:py3.12
+
+# optional, to change python version
+docker build . -t ptoas:py3.11 --build-arg PY_VER=cp311-cp311
+```
+
+Use:
+
+```bash
+docker run --rm -it ptoas:py3.12 /bin/bash
+```

docker/setup.py

@@ -0,0 +1,21 @@
+from setuptools import setup, find_namespace_packages
+
+setup(
+    name="ptoas",
+    version="0.1.0",
+    description="PTO Assembler & Optimizer",
+    # NOTE: find_namespace_packages detects folders even without __init__.py
+    packages=find_namespace_packages(),
+    # NOTE: The * at the end captures .so.22, .so.22.1, etc.
+    package_data={
+            "mlir": [
+                "**/*.so*",
+                "**/*.pyd",
+                "**/*.py",
+                "_mlir_libs/*.so*", 
+            ],
+        },
+    include_package_data=True,
+    zip_safe=False,
+    python_requires=">=3.9",
+)

lib/Bindings/Python/CMakeLists.txt

@@ -1,6 +1,6 @@
 # lib/Bindings/Python/CMakeLists.txt
 
-find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
+find_package(Python3 COMPONENTS Interpreter Development.Module REQUIRED)
 find_package(pybind11 CONFIG REQUIRED)
 
 include(TableGen)