Page MenuHomePhabricator

[Bazel] Add support for targeting macOS arm64
ClosedPublic

Authored by dfm on Sep 15 2021, 10:20 AM.

Details

Summary

In attempting to build JAX on Apple Silicon, we discovered an issue with
the bazel configuration in llvm-project-overlay. This patch fixes the logic,
at least when building JAX. More context is included on the following GitHub
issue: https://github.com/google/jax/issues/5501

Diff Detail

Event Timeline

dfm created this revision.Sep 15 2021, 10:20 AM
dfm published this revision for review.Sep 15 2021, 10:26 AM
Herald added a project: Restricted Project. · View Herald TranscriptSep 15 2021, 10:26 AM
GMNGeoffrey added inline comments.Sep 15 2021, 2:18 PM
utils/bazel/llvm-project-overlay/llvm/BUILD.bazel
21 ↗(On Diff #372744)

We already use @bazel_tools//src/conditions:foo in config.bzl. Looking at the history of src/conditions/BUILD, there's some bug in Bazel that means that cross-compiling from darwinn x86 to darwinn arm is broken and this is worked around in src/conditions by adding a --darwin_arm64 flag (https://github.com/bazelbuild/bazel/issues/12655 and https://github.com/bazelbuild/bazel/commit/52457b18bd), but only after Bazel 4.0 (the version we're currently on). Is that cross-compiling something you actually need? If not, I'd prefer we use the src/conditions constraints for now (at least until such a time as we decide to define all our own config settings)

utils/bazel/llvm-project-overlay/llvm/config.bzl
79–81

+Chandler. Mostly FYI since I know you cared about being able to compile on your mac :-)

GMNGeoffrey added a subscriber: chandlerc.
dfm added inline comments.Sep 15 2021, 5:09 PM
utils/bazel/llvm-project-overlay/llvm/BUILD.bazel
21 ↗(On Diff #372744)

Thanks for your response! Yeah, that's exactly what we're looking to do (I came to this when trying to build binaries for a Python library using x86 hardware), although the existing implementation wouldn't work even when building on an arm machine. I also tried your suggested change (using @bazel_tools//src/conditions:darwin_arm64) and I was still getting the wrong compiler flags when building locally, but I'm far from an expert in any of this though so there's a good shout that I was doing something wrong. I understand the hesitation to copy this over and I'd be happy to look into alternative ideas if you have suggestions. Thanks again!

GMNGeoffrey added inline comments.Sep 15 2021, 5:16 PM
utils/bazel/llvm-project-overlay/llvm/BUILD.bazel
21 ↗(On Diff #372744)

Ah in that case, I think we should probably upgrade to fix the Bazel bug :-) (in addition to adding the right selects in config.bzl). I would've done it already except that support for rbe_autoconfig was dropped for versions >4.0.0. I posted on bazel-discuss about how that squares with it being an LTS release.

So to clarify, with these edits you can cross-compile from darwin x86_64 to arm? That makes sense to me given the linked issue in Bazel that was worked around. How blocking is this for you? Like can you patch this in locally and then continue on your way for a bit while we fiddle around with using a newer version of Bazel? Can you just build with Bazel >4.2 (ignore the .bazelversion file) after adding the patch to use the src/conditions configs? I think the default Bazel install these days has a shell wrapper that automatically uses the version in the .bazelrc file, but if you specifically use a bazel-4.2 binary then it should work.

dfm added inline comments.Sep 15 2021, 5:24 PM
utils/bazel/llvm-project-overlay/llvm/BUILD.bazel
21 ↗(On Diff #372744)

That sounds like a good plan. It's definitely not a major blocker, so definitely no pressure. I'll play around with bazel versions and confirm my claims about cross-compiling tomorrow, just to be sure. Thanks!

dfm added inline comments.Sep 21 2021, 11:54 AM
utils/bazel/llvm-project-overlay/llvm/BUILD.bazel
21 ↗(On Diff #372744)

Just a quick update on this. I've done some more playing around and I can confirm that cross-compiling from darwin x86_64 to arm works with bazel version 4.2.1 (and probably others, but that's what I tested) using only the proposed change to config.bzl, without introducing these config settings. Is there a reason to not update the conditions to darwin_arm64 and darwin_x86_64 instead of plain darwin?

(Side note: I couldn't get any of this to work with the --cpu=darwin_arm64 and needed to use the --platforms instead; I remain confused by this, but maybe this is the expected behavior.)

utils/bazel/llvm-project-overlay/llvm/config.bzl
79–81

This should be

+     "@bazel_tools//src/conditions:darwin_arm64": native_arch_defines("AArch64", "arm64-apple-darwin"),
+     "@bazel_tools//src/conditions:darwin_x86_64": native_arch_defines("X86", "x86_64-apple-darwin"),

(rather than macos) but, otherwise, it looks like this would be a sufficient change to support both platforms.

GMNGeoffrey added inline comments.Sep 21 2021, 12:05 PM
utils/bazel/llvm-project-overlay/llvm/BUILD.bazel
21 ↗(On Diff #372744)

Is there a reason to not update the conditions to darwin_arm64 and darwin_x86_64 instead of plain darwin?

Nope, that's what I was suggesting :-)

(Side note: I couldn't get any of this to work with the --cpu=darwin_arm64 and needed to use the --platforms instead; I remain confused by this, but maybe this is the expected behavior.)

Yeah no idea, unfortunately

utils/bazel/llvm-project-overlay/llvm/config.bzl
79–81

SGTM.

GMNGeoffrey retitled this revision from Add support for targeting macOS arm64 with bazel to [Bazel] Add support for targeting macOS arm64.Sep 21 2021, 12:06 PM
dfm updated this revision to Diff 374003.Sep 21 2021, 12:17 PM

Updated in response to review. I'm new to this system so apologies if this isn't the right workflow!

Updated in response to review. I'm new to this system so apologies if this isn't the right workflow!

You got it right :-)

LGTM. I assume you don't have commit access and need me to land this for you? What author would you like me to use for the commit? From your GitHub history, my guess is "Dan F-M <foreman.mackey@gmail.com>"

GMNGeoffrey accepted this revision.Sep 21 2021, 12:21 PM
This revision is now accepted and ready to land.Sep 21 2021, 12:21 PM
dfm added a comment.Sep 21 2021, 12:22 PM

That's the one. Thanks!

This revision was automatically updated to reflect the committed changes.

FYI, TF reverted this with a patch on their side as a temporary fix because somehow this broke their mac build: https://github.com/tensorflow/tensorflow/commit/62093ca262. They're investigating the root cause further.

dfm added a comment.Sep 30 2021, 9:55 AM

Interesting! Do you know if there is a public discussion about this anywhere? I would like to know how they end up solving it, because I believe that their current build also must not work on arm.

@bkramer @jurahul could you move the discussion of fixing this somewhere public? Basically the issue is that the Bazel selects aren't working as expected and it's getting the linux triple instead. Probably some combination of the older Bazel version and the specific flags being passed on their mac build