build macOS arm64 wheels (#2616)

Cross-build macOS arm64 wheels.

- [x] Manually test is required.

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

.github/workflows/build_wheel.yml

@@ -22,6 +22,11 @@ jobs:
             python: 311
             platform_id: macosx_x86_64
             dp_variant: cpu
+          # macos-arm64
+          - os: macos-latest
+            python: 311
+            platform_id: macosx_arm64
+            dp_variant: cpu
           # win-64
           - os: windows-2019
             python: 311

backend/find_tensorflow.py

@@ -45,7 +45,16 @@ def find_tensorflow() -> Tuple[Optional[str], List[str]]:
     requires = []
 
     tf_spec = None
-    if os.environ.get("TENSORFLOW_ROOT") is not None:
+    if os.environ.get("CIBUILDWHEEL", "0") == "1" and os.environ.get(
+        "CIBW_BUILD", ""
+    ).endswith("macosx_arm64"):
+        # cibuildwheel cross build
+        site_packages = Path(os.environ.get("RUNNER_TEMP")) / "tensorflow"
+        tf_spec = FileFinder(str(site_packages)).find_spec("tensorflow")
+
+    if (tf_spec is None or not tf_spec) and os.environ.get(
+        "TENSORFLOW_ROOT"
+    ) is not None:
         site_packages = Path(os.environ.get("TENSORFLOW_ROOT")).parent.absolute()
         tf_spec = FileFinder(str(site_packages)).find_spec("tensorflow")
 
@@ -100,21 +109,20 @@ def get_tf_requirement(tf_version: str = "") -> dict:
     if tf_version == "":
         return {
             "cpu": [
-                "tensorflow-cpu; platform_machine!='aarch64'",
-                "tensorflow; platform_machine=='aarch64'",
+                "tensorflow-cpu; platform_machine!='aarch64' and (platform_machine!='arm64' or platform_system != 'Darwin')",
+                "tensorflow; platform_machine=='aarch64' or (platform_machine=='arm64' and platform_system == 'Darwin')",
             ],
             "gpu": [
-                "tensorflow; platform_machine!='aarch64'",
-                "tensorflow; platform_machine=='aarch64'",
+                "tensorflow",
+                "tensorflow-metal; platform_machine=='arm64' and platform_system == 'Darwin'",
             ],
         }
     elif tf_version in SpecifierSet("<1.15") or tf_version in SpecifierSet(
         ">=2.0,<2.1"
     ):
         return {
             "cpu": [
-                f"tensorflow=={tf_version}; platform_machine!='aarch64'",
-                f"tensorflow=={tf_version}; platform_machine=='aarch64'",
+                f"tensorflow=={tf_version}",
             ],
             "gpu": [
                 f"tensorflow-gpu=={tf_version}; platform_machine!='aarch64'",
@@ -124,12 +132,12 @@ def get_tf_requirement(tf_version: str = "") -> dict:
     else:
         return {
             "cpu": [
-                f"tensorflow-cpu=={tf_version}; platform_machine!='aarch64'",
-                f"tensorflow=={tf_version}; platform_machine=='aarch64'",
+                f"tensorflow-cpu=={tf_version}; platform_machine!='aarch64' and (platform_machine!='arm64' or platform_system != 'Darwin')",
+                f"tensorflow=={tf_version}; platform_machine=='aarch64'  or (platform_machine=='arm64' and platform_system == 'Darwin')",
             ],
             "gpu": [
-                f"tensorflow=={tf_version}; platform_machine!='aarch64'",
-                f"tensorflow=={tf_version}; platform_machine=='aarch64'",
+                f"tensorflow=={tf_version}",
+                "tensorflow-metal; platform_machine=='arm64' and platform_system == 'Darwin'",
             ],
         }
 

doc/install/easy-install.md

@@ -101,10 +101,10 @@ pip install deepmd-kit[cpu]
 ```bash
 pip install deepmd-kit[gpu,cu11,lmp,ipi]
 ```
-MPICH is required for parallel running.
+MPICH is required for parallel running. (The macOS arm64 package doesn't support MPI yet.)
 
 It is suggested to install the package into an isolated environment.
-The supported platform includes Linux x86-64 and aarch64 with GNU C Library 2.28 or above, macOS x86-64, and Windows x86-64.
+The supported platform includes Linux x86-64 and aarch64 with GNU C Library 2.28 or above, macOS x86-64 and arm64, and Windows x86-64.
 A specific version of TensorFlow which is compatible with DeePMD-kit will be also installed.
 
 :::{Warning}

pyproject.toml

@@ -76,8 +76,13 @@ manylinux-aarch64-image = "quay.io/pypa/manylinux_2_28_aarch64:2022-11-19-1b19e8
 
 [tool.cibuildwheel.macos]
 environment = { PIP_PREFER_BINARY="1", DP_LAMMPS_VERSION="stable_23Jun2022_update4", DP_ENABLE_IPI="1" }
-before-all = ["brew install mpich"]
-repair-wheel-command = "delocate-wheel --require-archs {delocate_archs} -w {dest_dir} -v {wheel} --ignore-missing-dependencies"
+before-all = [
+    """if [[ "$CIBW_BUILD" != *macosx_arm64* ]]; then brew install mpich; fi""",
+]
+before-build = [
+    """if [[ "$CIBW_BUILD" == *macosx_arm64* ]]; then python -m pip install "tensorflow-macos>=2.13.0rc0" --platform macosx_12_0_arm64 --no-deps --target=$RUNNER_TEMP/tensorflow; fi""",
+]
+repair-wheel-command = """if [[ "$CIBW_BUILD" == *macosx_arm64* ]]; then rm -rf $RUNNER_TEMP/tensorflow; fi && delocate-wheel --require-archs {delocate_archs} -w {dest_dir} -v {wheel} --ignore-missing-dependencies"""
 
 [tool.cibuildwheel.linux]
 repair-wheel-command = "auditwheel repair --exclude libtensorflow_framework.so.2 --exclude libtensorflow_framework.so.1 --exclude libtensorflow_framework.so --exclude _pywrap_tensorflow_internal.so --exclude libtensorflow_cc.so.2 -w {dest_dir} {wheel}"

setup.py

@@ -69,6 +69,7 @@
     find_libpython_requires = []
 else:
     find_libpython_requires = ["find_libpython"]
+cmake_args.append(f"-DTENSORFLOW_VERSION={tf_version}")
 
 
 class bdist_wheel_abi3(bdist_wheel):

source/cmake/Findtensorflow.cmake

@@ -280,20 +280,22 @@ else(BUILD_CPP_IF)
 endif(BUILD_CPP_IF)
 
 # detect TensorFlow version
-try_run(
-  TENSORFLOW_VERSION_RUN_RESULT_VAR TENSORFLOW_VERSION_COMPILE_RESULT_VAR
-  ${CMAKE_CURRENT_BINARY_DIR}/tf_version
-  "${CMAKE_CURRENT_LIST_DIR}/tf_version.cpp"
-  CMAKE_FLAGS "-DINCLUDE_DIRECTORIES:STRING=${TensorFlow_INCLUDE_DIRS}"
-  RUN_OUTPUT_VARIABLE TENSORFLOW_VERSION
-  COMPILE_OUTPUT_VARIABLE TENSORFLOW_VERSION_COMPILE_OUTPUT_VAR)
-if(NOT ${TENSORFLOW_VERSION_COMPILE_RESULT_VAR})
-  message(
-    FATAL_ERROR "Failed to compile: \n ${TENSORFLOW_VERSION_COMPILE_OUTPUT_VAR}"
-  )
-endif()
-if(NOT ${TENSORFLOW_VERSION_RUN_RESULT_VAR} EQUAL "0")
-  message(FATAL_ERROR "Failed to run, return code: ${TENSORFLOW_VERSION}")
+if(NOT DEFINED TENSORFLOW_VERSION)
+  try_run(
+    TENSORFLOW_VERSION_RUN_RESULT_VAR TENSORFLOW_VERSION_COMPILE_RESULT_VAR
+    ${CMAKE_CURRENT_BINARY_DIR}/tf_version
+    "${CMAKE_CURRENT_LIST_DIR}/tf_version.cpp"
+    CMAKE_FLAGS "-DINCLUDE_DIRECTORIES:STRING=${TensorFlow_INCLUDE_DIRS}"
+    RUN_OUTPUT_VARIABLE TENSORFLOW_VERSION
+    COMPILE_OUTPUT_VARIABLE TENSORFLOW_VERSION_COMPILE_OUTPUT_VAR)
+  if(NOT ${TENSORFLOW_VERSION_COMPILE_RESULT_VAR})
+    message(
+      FATAL_ERROR
+        "Failed to compile: \n ${TENSORFLOW_VERSION_COMPILE_OUTPUT_VAR}")
+  endif()
+  if(NOT ${TENSORFLOW_VERSION_RUN_RESULT_VAR} EQUAL "0")
+    message(FATAL_ERROR "Failed to run, return code: ${TENSORFLOW_VERSION}")
+  endif()
 endif()
 
 # print message