Build failure on Apple M1 at current HEAD

Building the jaxlib wheel from source is broken on my machine at current main HEAD:

➜ python build/build.py

     _   _  __  __
    | | / \ \ \/ /
 _  | |/ _ \ \  /
| |_| / ___ \/  \
 \___/_/   \/_/\_\


Bazel binary path: /opt/homebrew/bin/bazel
Bazel version: 5.0.0
Python binary path: /Users/nicholasjunge/Workspaces/python/jax/venv/bin/python
Python version: 3.9
NumPy version: 1.22.2
MKL-DNN enabled: yes
Target CPU: arm64
Target CPU features: release
CUDA enabled: no
TPU enabled: no
ROCm enabled: no

Building XLA and installing it in the jaxlib source tree...
/opt/homebrew/bin/bazel run --verbose_failures=true --config=mkl_open_source_only :build_wheel -- --output_path=/Users/nicholasjunge/Workspaces/python/jax/dist --cpu=arm64
INFO: Options provided by the client:
  Inherited 'common' options: --isatty=0 --terminal_columns=80
INFO: Reading rc options for 'run' from /Users/nicholasjunge/Workspaces/python/jax/.bazelrc:
  Inherited 'common' options: --experimental_repo_remote_exec
INFO: Reading rc options for 'run' from /Users/nicholasjunge/Workspaces/python/jax/.bazelrc:
  Inherited 'build' options: --apple_platform_type=macos --macos_minimum_os=10.9 --announce_rc --define open_source_build=true --spawn_strategy=standalone --enable_platform_specific_config --define=no_aws_support=true --define=no_gcp_support=true --define=no_hdfs_support=true --define=no_kafka_support=true --define=no_ignite_support=true --define=grpc_no_ares=true -c opt --config=short_logs --copt=-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.
INFO: Reading rc options for 'run' from /Users/nicholasjunge/Workspaces/python/jax/.jax_configure.bazelrc:
  Inherited 'build' options: --strategy=Genrule=standalone --repo_env PYTHON_BIN_PATH=/Users/nicholasjunge/Workspaces/python/jax/venv/bin/python --action_env=PYENV_ROOT --python_path=/Users/nicholasjunge/Workspaces/python/jax/venv/bin/python --distinct_host_configuration=false
INFO: Found applicable config definition build:short_logs in file /Users/nicholasjunge/Workspaces/python/jax/.bazelrc: --output_filter=DONT_MATCH_ANYTHING
INFO: Found applicable config definition build:mkl_open_source_only in file /Users/nicholasjunge/Workspaces/python/jax/.bazelrc: --define=tensorflow_mkldnn_contraction_kernel=1
INFO: Found applicable config definition build:macos in file /Users/nicholasjunge/Workspaces/python/jax/.bazelrc: --config=posix
INFO: Found applicable config definition build:posix in file /Users/nicholasjunge/Workspaces/python/jax/.bazelrc: --copt=-fvisibility=hidden --copt=-Wno-sign-compare --cxxopt=-std=c++14 --host_cxxopt=-std=c++14
Loading:
Loading: 0 packages loaded
Analyzing: target //build:build_wheel (0 packages loaded, 0 targets configured)
ERROR: /Users/nicholasjunge/Workspaces/python/jax/jaxlib/mlir/BUILD.bazel:58:15: no such target '@llvm-project//mlir/python:FuncPyFiles': target 'FuncPyFiles' not declared in package 'mlir/python' (did you mean 'QuantPyFiles'?) defined by /private/var/tmp/_bazel_nicholasjunge/270a4a78734ae0f3124fa7265b8a65ef/external/llvm-project/mlir/python/BUILD.bazel and referenced by '//jaxlib/mlir:_func_dialect_srcs'
ERROR: Analysis of target '//build:build_wheel' failed; build aborted:
INFO: Elapsed time: 0.154s
INFO: 0 processes.
FAILED: Build did NOT complete successfully (0 packages loaded, 0 targets configured)
ERROR: Build failed. Not running target
FAILED: Build did NOT complete successfully (0 packages loaded, 0 targets configured)
b''
Traceback (most recent call last):
  File "/Users/nicholasjunge/Workspaces/python/jax/build/build.py", line 527, in <module>
    main()
  File "/Users/nicholasjunge/Workspaces/python/jax/build/build.py", line 522, in main
    shell(command)
  File "/Users/nicholasjunge/Workspaces/python/jax/build/build.py", line 53, in shell
    output = subprocess.check_output(cmd)
  File "/opt/homebrew/Cellar/python@3.9/3.9.10/Frameworks/Python.framework/Versions/3.9/lib/python3.9/subprocess.py", line 424, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/opt/homebrew/Cellar/python@3.9/3.9.10/Frameworks/Python.framework/Versions/3.9/lib/python3.9/subprocess.py", line 528, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['/opt/homebrew/bin/bazel', 'run', '--verbose_failures=true', '--config=mkl_open_source_only', ':build_wheel', '--', '--output_path=/Users/nicholasjunge/Workspaces/python/jax/dist', '--cpu=arm64']' returned non-zero exit status 1.

Looks like a bazel problem, failing to discover an MLIR Python dependency.

jax on  main [⇡$] via jax
➜ git log
commit e8f1a020a2ebc920657b79ca94b2b6a4ea8dea03 (HEAD -> main, upstream/main)
Merge: dc2ca18d ad5144f5
Author: jax authors <[e-mail]>
Date:   Tue Mar 8 05:44:05 2022 -0800

    Merge pull request #9798 from hawkinsp:initdoc

    PiperOrigin-RevId: 433187449

Device is Apple M1 Pro @ macOS 12.2.1. Just let me know if you need more info, also happy to contribute a fix (if applicable).

[hawkinsp]

For GPU, StreamExecutor is currently the only path that works end to end.

Note that IREE can already run on Metal using a Vulkan-Metal compatibility layer.

[nicholasjng]

Thanks, I will test that out! Your fix in #9801 works, by the way, I just built the wheel successfully on my machine.