Add `torch-mlir-no-jit-importer` build case for mac os wheels (#1902)

* add flags to setup.py for out-of-tree build

* - fix build_ext bug
- add wheels script cases for mac wheels

.github/workflows/buildRelease.yml

@@ -63,7 +63,7 @@ jobs:
       run: mkdir dist
     - name: Copy releases to publish to dist directory
       if: github.event.inputs.release_id != ''
-      run: cp build_tools/python_deploy/wheelhouse/torch_mlir-*.whl dist/
+      run: cp build_tools/python_deploy/wheelhouse/torch_mlir*.whl dist/
 
     # Wheels must be published from a linux environment.
     #
@@ -120,7 +120,7 @@ jobs:
       run: mkdir dist
     - name: Copy releases to publish to dist directory
       if: github.event.inputs.release_id != ''
-      run: cp build_tools/python_deploy/wheelhouse/torch_mlir-*.whl dist/
+      run: cp build_tools/python_deploy/wheelhouse/torch_mlir*.whl dist/
 
     # Wheels must be published from a linux environment.
     #
@@ -179,7 +179,7 @@ jobs:
       continue-on-error: true
     - name: Copy releases to publish to dist directory
       if: github.event.inputs.release_id != ''
-      run: cp ./wheelhouse/torch_mlir-*.whl dist/
+      run: cp ./wheelhouse/torch_mlir*.whl dist/
 
     # Wheels must be published from a linux environment.
     #

build_tools/python_deploy/build_linux_packages.sh

@@ -48,7 +48,7 @@ TM_PYTHON_VERSIONS="${TM_PYTHON_VERSIONS:-cp38-cp38 cp310-cp310 cp311-cp311}"
 # Location to store Release wheels
 TM_OUTPUT_DIR="${TM_OUTPUT_DIR:-${this_dir}/wheelhouse}"
 # What "packages to build"
-TM_PACKAGES="${TM_PACKAGES:-torch-mlir}"
+TM_PACKAGES="${TM_PACKAGES:-torch-mlir torch-mlir-no-jit}"
 # Use pre-built Pytorch
 TM_USE_PYTORCH_BINARY="${TM_USE_PYTORCH_BINARY:-ON}"
 # Skip running tests if you want quick iteration
@@ -84,6 +84,11 @@ function run_on_host() {
       export USERID=0
       export GROUPID=0
       ;;
+    torch-mlir-no-jit)
+      TM_CURRENT_DOCKER_IMAGE=${TM_RELEASE_DOCKER_IMAGE}
+      export USERID=0
+      export GROUPID=0
+      ;;
     out-of-tree)
       TM_CURRENT_DOCKER_IMAGE=${TM_CI_DOCKER_IMAGE}
       # CI uses only Python3.10
@@ -159,6 +164,12 @@ function run_in_docker() {
 
           clean_build torch_mlir "$python_version"
           ;;
+        torch-mlir-no-jit)
+          clean_wheels torch_mlir_no_jit "$python_version"
+          build_torch_mlir_no_jit
+          run_audit_wheel torch_mlir_no_jit "$python_version"
+          clean_build torch_mlir_no_jit "$python_version"
+          ;;
         out-of-tree)
           setup_venv "$python_version"
           build_out_of_tree "$TM_USE_PYTORCH_BINARY" "$python_version"
@@ -373,6 +384,24 @@ function run_audit_wheel() {
   rm "$generic_wheel"
 }
 
+function build_torch_mlir_no_jit() {
+  python -m pip install --no-cache-dir -r /main_checkout/torch-mlir/build-requirements.txt
+  CMAKE_GENERATOR=Ninja \
+  TORCH_MLIR_PYTHON_PACKAGE_VERSION=${TORCH_MLIR_PYTHON_PACKAGE_VERSION} \
+  TORCH_MLIR_ENABLE_JIT_IR_IMPORTER=0 \
+  TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS=1 \
+  python -m pip wheel -v -w /wheelhouse /main_checkout/torch-mlir
+}
+
+function run_audit_wheel_no_jit() {
+  local wheel_basename="$1"
+  local python_version="$2"
+  generic_wheel="/wheelhouse/${wheel_basename}-${TORCH_MLIR_PYTHON_PACKAGE_VERSION}-${python_version}-linux_x86_64.whl"
+  echo ":::: Auditwheel $generic_wheel"
+  auditwheel repair -w /wheelhouse "$generic_wheel"
+  rm "$generic_wheel"
+}
+
 function clean_wheels() {
   local wheel_basename="$1"
   local python_version="$2"

build_tools/python_deploy/build_macos_packages.sh

@@ -22,7 +22,7 @@ this_dir="$(cd "$(dirname "$0")" && pwd)"
 repo_root="$(cd "$this_dir"/../../ && pwd)"
 python_versions="${TORCH_MLIR_PYTHON_VERSIONS:-3.9 3.10 3.11}"
 output_dir="${output_dir:-${this_dir}/wheelhouse}"
-packages="${packages:-torch-mlir}"
+packages="${packages:-torch-mlir torch-mlir-no-jit-importer}"
 
 PKG_VER_FILE="${repo_root}"/torch_mlir_package_version ; [ -f "$PKG_VER_FILE" ] && . "$PKG_VER_FILE"
 export TORCH_MLIR_PYTHON_PACKAGE_VERSION="${TORCH_MLIR_PYTHON_PACKAGE_VERSION:-0.0.1}"
@@ -61,6 +61,11 @@ function run() {
           build_torch_mlir torch_mlir "$python_version"
           run_audit_wheel torch_mlir "$python_version"
           ;;
+        torch-mlir-no-jit-importer)
+          clean_wheels torch_mlir_no_jit "$python_version"
+          build_torch_mlir_no_jit torch_mlir_no_jit "$python_version"
+          run_audit_wheel_no_jit torch_mlir_no_jit "$python_version"
+          ;;
         *)
           echo "Unrecognized package '$package'"
           exit 1
@@ -88,6 +93,25 @@ function build_torch_mlir() {
   rm -rf "$output_dir"/build_venv
 }
 
+function build_torch_mlir_no_jit() {
+  local wheel_basename="$1"
+  local python_version="$2"
+  rm -rf "$output_dir"/build_venv
+  python"${python_version}" -m venv "$output_dir"/build_venv
+  source "$output_dir"/build_venv/bin/activate
+  python"${python_version}" -m pip install -U pip delocate
+  python"${python_version}" -m pip install -r "$repo_root"/build-requirements.txt
+  CMAKE_GENERATOR=Ninja \
+  TORCH_MLIR_PYTHON_PACKAGE_VERSION=${TORCH_MLIR_PYTHON_PACKAGE_VERSION} \
+  MACOSX_DEPLOYMENT_TARGET=$MACOSX_DEPLOYMENT_TARGET \
+  CMAKE_OSX_ARCHITECTURES=$CMAKE_OSX_ARCHITECTURES \
+  TORCH_MLIR_ENABLE_JIT_IR_IMPORTER=0 \
+  TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS=1 \
+  python"${python_version}" -m pip wheel -v -w "$output_dir" "$repo_root"
+  deactivate
+  rm -rf "$output_dir"/build_venv
+}
+
 function clean_wheels() {
   local wheel_basename="$1"
   local python_version="$2"
@@ -117,4 +141,24 @@ function run_audit_wheel() {
   fi
 }
 
+function run_audit_wheel_no_jit() {
+  set +x
+  local wheel_basename="$1"
+  local python_version="$2"
+  generic_wheel=$(ls "$output_dir"/"${wheel_basename}"-* | grep "${python_version//./}")
+  echo "Looking for $generic_wheel"
+  if [ -f "$generic_wheel" ]; then
+    echo "$generic_wheel found. Delocating it.."
+    rm -rf "$output_dir"/test_venv
+    python"${python_version}" -m venv "$output_dir"/test_venv
+    source "$output_dir"/test_venv/bin/activate
+    python"${python_version}" -m pip install -U pip delocate
+    python"${python_version}" -m pip install -r "$repo_root"/build-requirements.txt
+    python"${python_version}" -m pip install "$generic_wheel"
+    DYLD_LIBRARY_PATH="$output_dir"/test_venv/lib/python"${python_version}"/site-packages/torch/lib delocate-wheel -v "$generic_wheel"
+    deactivate
+    rm -rf "$output_dir"/test_venv
+  fi
+}
+
 run

setup.py

@@ -146,7 +146,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
 
 
 setup(
-    name="torch-mlir",
+    name="torch-mlir" if not TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS else "torch-mlir-no-jit",
     version=f"{PACKAGE_VERSION}",
     author="Sean Silva",
     author_email="silvasean@google.com",
@@ -161,7 +161,7 @@ setup(
     },
     ext_modules=[
         CMakeExtension("torch_mlir._mlir_libs._jit_ir_importer"),
-    ] if not TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS else [CMakeExtension("_")],
+    ] if not TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS else [CMakeExtension("torch_mlir._mlir_libs._torchMlir")],
     install_requires=["numpy", ] + (
         [f"torch=={torch.__version__}".split("+", 1)[0], ] if not TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS else []),
     zip_safe=False,