Skip to content

Commit

Permalink
Support osx-arm64 native compile
Browse files Browse the repository at this point in the history
  • Loading branch information
xhochy committed Mar 23, 2021
1 parent 3db0c97 commit 1b30623
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 2 deletions.
3 changes: 1 addition & 2 deletions recipe/build.sh
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
set -euxo pipefail

# TODO(features):
# * --enable_mkl_dnn
# * --enable_cuda
# * --enable_rocm
# * --target_cpu_features (different settings depending on archspec)
Expand All @@ -13,5 +12,5 @@ set -euxo pipefail
# * MLIR
# * protobuf
# * mkl_dnn_v1 (nowadays called oneDNN)
${PYTHON} build/build.py --target_cpu_features default
${PYTHON} build/build.py --target_cpu_features default --enable_mkl_dnn
${PYTHON} -m pip install dist/jaxlib-*.whl
3 changes: 3 additions & 0 deletions recipe/meta.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ package:
source:
url: https://github.com/google/jax/archive/{{ name }}-v{{ version }}.tar.gz
sha256: e70d6da2ee6f065f2e5f6fb0bb1d541f362d68a4ca27398827f3a34b13f3896d
patches:
- osx-arm64.patch

build:
number: 0
Expand All @@ -17,6 +19,7 @@ requirements:
build:
- {{ compiler('c') }}
- {{ compiler('cxx') }}
- bazel
- python # [build_platform != target_platform]
- cross-python_{{ target_platform }} # [build_platform != target_platform]
- numpy # [build_platform != target_platform]
Expand Down
40 changes: 40 additions & 0 deletions recipe/osx-arm64.patch
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
--- build/build.py.orig 2021-03-23 13:57:54.000000000 +0100
+++ build/build.py 2021-03-23 13:59:39.000000000 +0100
@@ -271,6 +271,13 @@
# Workaround for gcc 10+ warnings related to upb.
# See https://github.com/tensorflow/tensorflow/issues/39467
build:linux --copt=-Wno-stringop-truncation
+
+# Apple silicon config
+build:macos_arm64 --config=posix
+build:macos_arm64 --config=macos
+build:macos_arm64 --apple_platform_type=macos
+build:macos_arm64 --cpu=darwin_arm64
+build --copt=-Wno-implicit-function-declaration
"""


--- build/build_wheel.py.orig 2021-03-23 13:58:02.000000000 +0100
+++ build/build_wheel.py 2021-03-23 14:00:17.000000000 +0100
@@ -169,12 +169,18 @@

def build_wheel(sources_path, output_path):
"""Builds a wheel in `output_path` using the source tree in `sources_path`."""
+ system = platform.system()
platform_name = {
"Linux": "manylinux2010",
- "Darwin": "macosx_10_9",
+ "Darwin": "macosx_11_0" if int(platform.mac_ver()[0].split(".")[0]) > 10 else "macosx_10_9",
"Windows": "win",
- }[platform.system()]
- cpu_name = "amd64" if platform.system() == "Windows" else "x86_64"
+ }[system]
+ if system == "Darwin":
+ cpu_name = platform.mac_ver()[2]
+ elif system == "Windows":
+ cpu_name = "amd64"
+ else:
+ cpu_name = "x86_64"
python_tag_arg = (f"--python-tag=cp{sys.version_info.major}"
f"{sys.version_info.minor}")
platform_tag_arg = f"--plat-name={platform_name}_{cpu_name}"

0 comments on commit 1b30623

Please sign in to comment.