fix buildwheels (#27)

.github/workflows/build.yml

@@ -9,9 +9,6 @@ on:
   workflow_dispatch:
     branches:
       - main
-  pull_request:
-    branches:
-      - main
 
 jobs:
 
@@ -22,6 +19,11 @@ jobs:
     strategy:
       matrix:
         os: [ ubuntu-latest, macos-latest ]
+        arch: [ AArch64, X86 ]
+        py_version: [ "3.11" ]
+        exclude:
+          - os: macos-latest
+            arch: X86
 
     steps:
       - name: Checkout
@@ -30,7 +32,7 @@ jobs:
       # Used to host cibuildwheel
       - uses: actions/setup-python@v3
         with:
-          python-version: "3.10"
+          python-version: ${{ matrix.py_version }}
 
       - name: Install Ninja
         uses: llvm/actions/install-ninja@6a57890d0e3f9f35dfc72e7e48bc5e1e527cdd6c # Jan 17
@@ -42,22 +44,92 @@ jobs:
           # LLVM needs serious cache size
           max-size: 6G
 
-      - name: Install terminfo
-        if: contains(matrix.os, 'ubuntu')
+      - name: Install cross-compilation toolchain if necessary
+        if: contains(matrix.os, 'ubuntu') && contains(matrix.arch, 'AArch64')
         run: |
           sudo apt-get update
-          sudo apt -y install terminfo
+          sudo apt-get install -y binutils-aarch64-linux-gnu \
+            g++-aarch64-linux-gnu gcc-aarch64-linux-gnu
 
-      - name: Build wheels
+      - name: Get host llvmtblgen/mlirtblgen if necessary
+        if: contains(matrix.os, 'ubuntu') && contains(matrix.arch, 'AArch64')
+        run: |
+          TORCH_MLIR_COMMIT=$(git ls-tree HEAD externals/torch-mlir --object-only --abbrev=8)
+          wget "https://github.com/nod-ai/PI/releases/download/torch-mlir-${TORCH_MLIR_COMMIT}/torch-mlir-${TORCH_MLIR_COMMIT}-${{ matrix.os }}-X86.tar.xz" -O host-torch-mlir-install.tar.xz
+          mkdir host-torch-mlir-install && tar -xvf host-torch-mlir-install.tar.xz -C host-torch-mlir-install
+          echo "TORCH_MLIR_HOST_MAIN_INSTALL_DIR=${PWD}/host-torch-mlir-install/torch_mlir_install" | tee -a $GITHUB_ENV
+
+      - name: Choose compiler and set env vars
         run: |
-          if [ x"${{ matrix.os }}" == x"ubuntu-latest" ]; then
-            CMAKE_GENERATOR=Ninja pip wheel . -w wheelhouse -v
+          
+          PY_VERSION=$(echo ${{ matrix.py_version }} | sed 's/\.//g')
+          
+          if [ x"${{ matrix.os }}" == x"macos-latest" ]; then
+            echo "CXX_COMPILER=clang++" | tee -a $GITHUB_ENV
+            echo "C_COMPILER=clang" | tee -a $GITHUB_ENV
+            echo "LLVM_DEFAULT_TARGET_TRIPLE=arm64-apple-darwin21.6.0" | tee -a $GITHUB_ENV
+            echo "LLVM_HOST_TRIPLE=arm64-apple-darwin21.6.0" | tee -a $GITHUB_ENV
+            echo "ARCH=AArch64" | tee -a $GITHUB_ENV
+            echo "PY_BUILD_EXT_SUFFIX=.cpython-${PY_VERSION}-darwin.so" | tee -a $GITHUB_ENV
+            echo "PLAT_NAME=macosx_12_0_arm64" | tee -a $GITHUB_ENV
+          elif [ x"${{ matrix.os }}" == x"ubuntu-latest" ] && [ x"${{ matrix.arch }}" == x"AArch64" ]; then
+            echo "CXX_COMPILER=aarch64-linux-gnu-g++" | tee -a $GITHUB_ENV
+            echo "C_COMPILER=aarch64-linux-gnu-gcc" | tee -a $GITHUB_ENV
+            echo "LLVM_DEFAULT_TARGET_TRIPLE=aarch64-linux-gnu" | tee -a $GITHUB_ENV
+            echo "LLVM_HOST_TRIPLE=aarch64-linux-gnu" | tee -a $GITHUB_ENV
+            echo "ARCH=AArch64" | tee -a $GITHUB_ENV
+            echo "PY_BUILD_EXT_SUFFIX=.cpython-${PY_VERSION}-aarch64-linux-gnu.so" | tee -a $GITHUB_ENV
+            echo "PLAT_NAME=linux_aarch64" | tee -a $GITHUB_ENV
           else
-            pip install -r requirements.txt
-            cibuildwheel --platform macos
+            echo "CXX_COMPILER=g++" | tee -a $GITHUB_ENV
+            echo "C_COMPILER=gcc" | tee -a $GITHUB_ENV
+            echo "LLVM_DEFAULT_TARGET_TRIPLE=x86_64-unknown-linux-gnu" | tee -a $GITHUB_ENV
+            echo "LLVM_HOST_TRIPLE=x86_64-unknown-linux-gnu" | tee -a $GITHUB_ENV
+            echo "ARCH=X86" | tee -a $GITHUB_ENV
+            echo "PY_BUILD_EXT_SUFFIX=.cpython-${PY_VERSION}-x86_64-linux-gnu.so" | tee -a $GITHUB_ENV
+            echo "PLAT_NAME=linux_x86_64" | tee -a $GITHUB_ENV
+          fi
+
+      - name: Build wheels
+        run: |
+          CMAKE_ARGS="\
+            -DCMAKE_BUILD_TYPE=Release \
+            -DCMAKE_CXX_COMPILER=$CXX_COMPILER \
+            -DCMAKE_C_COMPILER=$C_COMPILER \
+            -DLLVM_CCACHE_BUILD=ON \
+            -DLLVM_DEFAULT_TARGET_TRIPLE=$LLVM_DEFAULT_TARGET_TRIPLE \
+            -DLLVM_ENABLE_ZLIB=OFF \
+            -DLLVM_ENABLE_ZSTD=OFF \
+            -DLLVM_HOST_TRIPLE=$LLVM_HOST_TRIPLE \
+            -DLLVM_TARGETS_TO_BUILD=$ARCH \
+            -DLLVM_TARGET_ARCH=$ARCH \
+            -DLLVM_USE_HOST_TOOLS=ON \
+            -DPY_BUILD_EXT_SUFFIX=$PY_BUILD_EXT_SUFFIX"
+          
+          if [ x"${{ matrix.os }}" == x"macos-latest" ]; then
+            CMAKE_ARGS="${CMAKE_ARGS} \
+              -DCMAKE_OSX_ARCHITECTURES=arm64 \
+              -DMACOSX_DEPLOYMENT_TARGET=12.0"
+          elif [ x"${{ matrix.os }}" == x"ubuntu-latest" ] && [ x"${{ matrix.arch }}" == x"AArch64" ]; then
+            CMAKE_ARGS="${CMAKE_ARGS} \
+              -DLLVM_TABLEGEN=$TORCH_MLIR_HOST_MAIN_INSTALL_DIR/bin/llvm-tblgen \
+              -DMLIR_LINALG_ODS_YAML_GEN=$TORCH_MLIR_HOST_MAIN_INSTALL_DIR/bin/mlir-linalg-ods-yaml-gen \
+              -DMLIR_LINALG_ODS_YAML_GEN_EXE=$TORCH_MLIR_HOST_MAIN_INSTALL_DIR/bin/mlir-linalg-ods-yaml-gen \
+              -DMLIR_PDLL_TABLEGEN=$TORCH_MLIR_HOST_MAIN_INSTALL_DIR/bin/mlir-pdll \
+              -DMLIR_TABLEGEN=$TORCH_MLIR_HOST_MAIN_INSTALL_DIR/bin/mlir-tblgen"
           fi
+          
+          echo $CMAKE_ARGS
+          
+          export CMAKE_GENERATOR=Ninja
+          export CMAKE_ARGS=$CMAKE_ARGS
+          export TORCH_MLIR_DISTRO_ARCHITECTURE=$ARCH
+          pip install -r requirements.txt
+          python3 setup.py bdist_wheel --plat-name=$PLAT_NAME --dist-dir wheelhouse
+          
+          PIVERSION="$(python setup.py --version)"
           mkdir -p ${{ github.sha }}
-          mv wheelhouse/pi*.whl ${{ github.sha }}/
+          mv wheelhouse/pi-$PIVERSION-*.whl ${{ github.sha }}/
 
       - name: Upload an artifact
         uses: actions/upload-artifact@v3

.github/workflows/build_torch_mlir.yml

@@ -10,10 +10,10 @@ on:
     paths:
       - externals/**
       - .github/workflows/build_torch_mlir.yml
-  pull_request:
-    paths:
-      - externals/**
-      - .github/workflows/build_torch_mlir.yml
+#  pull_request:
+#    paths:
+#      - externals/**
+#      - .github/workflows/build_torch_mlir.yml
 
 jobs:
 

CMakeLists.txt

@@ -1,43 +1,48 @@
 cmake_minimum_required(VERSION 3.13.4)
 
-if (POLICY CMP0068)
-    cmake_policy(SET CMP0068 NEW)
-    set(CMAKE_BUILD_WITH_INSTALL_NAME_DIR ON)
-endif ()
+if(POLICY CMP0068)
+  cmake_policy(SET CMP0068 NEW)
+  set(CMAKE_BUILD_WITH_INSTALL_NAME_DIR ON)
+endif()
 
-if (POLICY CMP0075)
-    cmake_policy(SET CMP0075 NEW)
-endif ()
+if(POLICY CMP0075)
+  cmake_policy(SET CMP0075 NEW)
+endif()
 
-if (POLICY CMP0077)
-    cmake_policy(SET CMP0077 NEW)
-endif ()
+if(POLICY CMP0077)
+  cmake_policy(SET CMP0077 NEW)
+endif()
 
-if (POLICY CMP0116)
-    cmake_policy(SET CMP0116 NEW)
-endif ()
+if(POLICY CMP0116)
+  cmake_policy(SET CMP0116 NEW)
+endif()
 
-if (POLICY CMP0135)
-    cmake_policy(SET CMP0116 OLD)
-endif ()
+if(POLICY CMP0135)
+  cmake_policy(SET CMP0116 OLD)
+endif()
 
 project(PI LANGUAGES CXX C)
 
 set(CMAKE_BUILD_WITH_INSTALL_NAME_DIR ON)
 
 set(CMAKE_CXX_STANDARD
-        17
-        CACHE STRING "C++ standard to conform to")
+    17
+    CACHE STRING "C++ standard to conform to")
 
-set(TORCH_MLIR_INSTALL_DIR "" CACHE STRING "Path to TorchMLIR install dir")
+set(TORCH_MLIR_INSTALL_DIR
+    ""
+    CACHE STRING "Path to TorchMLIR install dir")
 
-if ((TORCH_MLIR_INSTALL_DIR STREQUAL "") OR (NOT EXISTS TORCH_MLIR_INSTALL_DIR))
-    set(TorchMLIR_DIR ".")
-    find_package(TorchMLIR)
-endif ()
+set(PY_BUILD_EXT_SUFFIX ".so" CACHE INTERNAL "Cross python lib extension")
+
+if((${TORCH_MLIR_INSTALL_DIR} STREQUAL "") OR (NOT EXISTS
+                                               ${TORCH_MLIR_INSTALL_DIR}))
+  set(TorchMLIR_DIR ".")
+  find_package(TorchMLIR)
+endif()
 
 message(STATUS "TORCH_MLIR_INSTALL_DIR: ${TORCH_MLIR_INSTALL_DIR}")
-list(APPEND CMAKE_PREFIX_PATH "${TORCH_MLIR_INSTALL_DIR}")
+list(APPEND CMAKE_PREFIX_PATH ${TORCH_MLIR_INSTALL_DIR})
 find_package(MLIR REQUIRED CONFIG)
 
 message(STATUS "Using MLIRConfig.cmake in: ${MLIR_DIR}")

TorchMLIRConfig.cmake

@@ -9,11 +9,14 @@ endif()
 
 if(CMAKE_SYSTEM_PROCESSOR MATCHES "(AMD64|x86_64)")
   set(ARCH X86)
-elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)")
+elseif((CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)")
+       OR (CMAKE_OSX_ARCHITECTURES MATCHES "arm64"))
   set(ARCH AArch64)
 endif()
 
-if (CMAKE_OSX_ARCHITECTURES MATCHES "arm64")
+if($ENV{TORCH_MLIR_DISTRO_ARCHITECTURE} MATCHES "^(AMD64|x86_64|X86)")
+  set(ARCH X86)
+elseif($ENV{TORCH_MLIR_DISTRO_ARCHITECTURE} MATCHES "^(aarch64|arm64|AArch64)")
   set(ARCH AArch64)
 endif()
 
@@ -43,10 +46,10 @@ message(STATUS "Using torch-mlir commit ${TORCH_MLIR_COMMIT}")
 
 # Try to download torch-mlir distro
 if(NOT EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/torch-mlir-install.tar.xz)
-  message(STATUS "Downloading torch-mlir distro.")
   set(TORCH_MLIR_INSTALL_URL
       "https://github.com/nod-ai/PI/releases/download/torch-mlir-${TORCH_MLIR_COMMIT}/torch-mlir-${TORCH_MLIR_COMMIT}-${OS}-latest-${ARCH}.tar.xz"
   )
+  message(STATUS "Downloading torch-mlir distro from ${TORCH_MLIR_INSTALL_URL}")
   file(
     DOWNLOAD ${TORCH_MLIR_INSTALL_URL}
     ${CMAKE_CURRENT_SOURCE_DIR}/torch-mlir-install.tar.xz
@@ -59,6 +62,7 @@ if(NOT EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/torch-mlir-install.tar.xz)
     file(REMOVE ${CMAKE_CURRENT_SOURCE_DIR}/torch-mlir-install.tar.xz)
   endif()
 else()
+  message(STATUS "Already downloaded torch-mlir distro.")
   set(TORCH_MLIR_DOWNLOAD_STATUS_CODE 0)
   set(TORCH_MLIR_DOWNLOAD_STATUS_STRING "No error")
 endif()
@@ -77,7 +81,10 @@ if(TORCH_MLIR_DOWNLOAD_STATUS_CODE EQUAL 0)
   set(TORCH_MLIR_INSTALL_DIR
       "${CMAKE_CURRENT_SOURCE_DIR}/torch_mlir_install/torch_mlir_install")
 else()
-  message(STATUS "Failed to download torch-mlir because ${TORCH_MLIR_DOWNLOAD_STATUS_STRING}.")
+  message(
+    STATUS
+      "Failed to download torch-mlir because ${TORCH_MLIR_DOWNLOAD_STATUS_STRING}."
+  )
   message(STATUS "Will unshallow submodule and build and install.")
 
   execute_process(
@@ -123,8 +130,7 @@ else()
       -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}
       -DCMAKE_INSTALL_PREFIX=${TORCH_MLIR_INSTALL_DIR} -DLLVM_CCACHE_BUILD=ON
       -DLLVM_ENABLE_ASSERTIONS=ON -DLLVM_ENABLE_PROJECTS=mlir
-      -DLLVM_ENABLE_TERMINFO=OFF
-      -DLLVM_ENABLE_ZSTD=OFF
+      -DLLVM_ENABLE_TERMINFO=OFF -DLLVM_ENABLE_ZSTD=OFF
       -DLLVM_EXTERNAL_PROJECTS=torch-mlir\;torch-mlir-dialects
       -DLLVM_EXTERNAL_TORCH_MLIR_DIALECTS_SOURCE_DIR=${TORCH_MLIR_MAIN_SRC_DIR}/externals/llvm-external-projects/torch-mlir-dialects
       -DLLVM_EXTERNAL_TORCH_MLIR_SOURCE_DIR=${TORCH_MLIR_MAIN_SRC_DIR}

cpp_ext/CMakeLists.txt

@@ -1,7 +1,14 @@
 include(MLIRDetectPythonEnv)
 include(AddMLIRPython)
-mlir_configure_python_dev_packages()
 mlir_detect_pybind11_install()
+mlir_configure_python_dev_packages()
+
+# for cross-compiling
+message(STATUS "PY_BUILD_EXT_SUFFIX ${PY_BUILD_EXT_SUFFIX}")
+set(PYTHON_MODULE_EXTENSION
+    ${PY_BUILD_EXT_SUFFIX}
+    CACHE INTERNAL ".so" FORCE)
+message(STATUS "PYTHON_MODULE_EXTENSION ${PYTHON_MODULE_EXTENSION}")
 
 set(PI_MLIR_PYTHON_PACKAGES_DIR "${CMAKE_CURRENT_BINARY_DIR}/python_packages")
 

pyproject.toml

@@ -19,15 +19,15 @@ log_cli_date_format = "%Y-%m-%d %H:%M:%S"
 before-build = "pip install -r requirements.txt -v"
 # HOLY FUCK
 # https://github.com/pypa/pip/issues/5229#issuecomment-387301397
-environment = { PIP_NO_BUILD_ISOLATION = "false" }
+#environment = { PIP_NO_BUILD_ISOLATION = "true" }
 build-verbosity = "3"
 
-#[tool.cibuildwheel.linux]
-#archs = ["x86_64"]
-#build = "cp311-manylinux_x86_64"
+[tool.cibuildwheel.linux]
+build = "cp311-manylinux_x86_64"
+environment = { CIBW_ENVIRONMENT_PASS_LINUX = "TORCH_MLIR_DISTRO_ARCHITECTURE TORCH_MLIR_INSTALL_DIR" }
 
 [tool.cibuildwheel.macos]
 build = "cp311-*"
 archs = ["arm64"]
-environment = { PIP_NO_BUILD_ISOLATION = "false", ARCHFLAGS = "-arch arm64" , MACOSX_DEPLOYMENT_TARGET= 12.6 }
+environment = { CMAKE_OSX_ARCHITECTURES = "arm64", MACOSX_DEPLOYMENT_TARGET = 12.6 }
 repair-wheel-command = "delocate-wheel --ignore-missing-dependencies --require-archs {delocate_archs} -w {dest_dir} -v {wheel}"