[WIP] ARM OSX Migrator

[regro-cf-autotick-bot]

This feedstock is being rebuilt as part of the ARM OSX migration.

Feel free to merge the PR if CI is all green, but please don't close it
without reaching out the the ARM OSX team first at @conda-forge/help-osx-arm64.

If this PR was opened in error or needs to be updated please add the bot-rerun label to this PR. The bot will close this PR and schedule another one. If you do not have permissions to add this label, you can use the phrase @conda-forge-admin, please rerun bot in a PR comment to have the conda-forge-admin add it for you.

_{This PR was created by the regro-cf-autotick-bot.
The regro-cf-autotick-bot is a service to automatically track the dependency graph, migrate packages, and propose package version updates for conda-forge. If you would like a local version of this bot, you might consider using rever. Rever is a tool for automating software releases and forms the backbone of the bot's conda-forge PRing capability. Rever is both conda (conda install -c conda-forge rever) and pip (pip install re-ver) installable.
Finally, feel free to drop us a line if there are any issues!
This PR was generated by https://github.com/regro/autotick-bot/actions/runs/638294645, please use this URL for debugging}

[conda-forge-linter]

Hi! This is the friendly automated conda-forge-linting service.

I just wanted to let you know that I linted all conda-recipes in your PR (recipe) and found it was in an excellent condition.

[xhochy]

@ericmjl Do you know whether the wheels have CUDA support?

[conda-forge-linter]

Hi! This is the friendly automated conda-forge-linting service.

I was trying to look for recipes to lint for you, but it appears we have a merge conflict.
Please try to merge or rebase with the base branch to resolve this conflict.

Please ping the 'conda-forge/core' team (using the @ notation in a comment) if you believe this is a bug.

[conda-forge-linter]

Hi! This is the friendly automated conda-forge-linting service.

I just wanted to let you know that I linted all conda-recipes in your PR (recipe) and found it was in an excellent condition.

[xhochy]

We need bazel from conda-forge, this depends currently on conda-forge/openjdk-feedstock#78

[ericmjl]

@ericmjl Do you know whether the wheels have CUDA support?

Hey @xhochy! The pip-distributed wheels are CPU-only. According to the JAX developer docs, the pre-compiled GPU wheels are provided on Google's blob storage.

[xhochy]

😌 That makes it easier than to keep the package functionality in line with what the wheels contained. We can add CUDA support as an option to the package here later as we did with other packages on conda-forge, too. Just the default variant should be CUDA-free.

[xhochy]

I've got bazel building over at conda-forge/bazel-feedstock#60 but compilation on the M1 fails with

[1,620 / 6,348] Compiling src/cpu/x64/gemm/f32/jit_avx512_core_f32_copy_at_kern_autogen.cpp; 145s local ... (8 actions running)
ERROR: /Users/uwe/mambaforge/conda-bld/jaxlib_1616437349574/_build_env/share/bazel/e8631f0932ec352f4790a4cefa08bf36/external/mkl_dnn_v1/BUILD.bazel:112:11: Compiling src/cpu/x64/jit_utils/jitprofiling/jitprofiling.c failed: (Exit 1): wrapped_clang failed: error executing command
  (cd /Users/uwe/mambaforge/conda-bld/jaxlib_1616437349574/_build_env/share/bazel/e8631f0932ec352f4790a4cefa08bf36/execroot/__main__ && \
  exec env - \
    APPLE_SDK_PLATFORM=MacOSX \
    APPLE_SDK_VERSION_OVERRIDE=11.1 \
    PATH=/Users/uwe/mambaforge/conda-bld/jaxlib_1616437349574/_build_env/bin:/Users/uwe/mambaforge/conda-bld/jaxlib_1616437349574/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_place/bin:/Users/uwe/mambaforge/conda-bld/jaxlib_1616437349574/_build_env:/Users/uwe/mambaforge/conda-bld/jaxlib_1616437349574/_h_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_place:/opt/google-cloud-sdk/bin:/Users/uwe/mambaforge/condabin:/opt/homebrew/bin:/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin:/Applications/Keybase.app/Contents/SharedSupport/bin:/Library/Apple/usr/bin:/Users/uwe/.local/bin \
    TF_CUDA_COMPUTE_CAPABILITIES=3.5,5.2,6.0,6.1,7.0 \
    TF_ROCM_AMDGPU_TARGETS=gfx803,gfx900,gfx906,gfx1010 \
    XCODE_VERSION_OVERRIDE=12.4.0.12D4e \
  external/local_config_cc/wrapped_clang '-D_FORTIFY_SOURCE=1' -fstack-protector -fcolor-diagnostics -Wall -Wthread-safety -Wself-assign -fno-omit-frame-pointer -g0 -O2 '-D_FORTIFY_SOURCE=1' -DNDEBUG '-DNS_BLOCK_ASSERTIONS=1' -fdebug-compilation-dir . -iquote external/mkl_dnn_v1 -iquote bazel-out/darwin-opt/bin/external/mkl_dnn_v1 -isystem external/mkl_dnn_v1/include -isystem bazel-out/darwin-opt/bin/external/mkl_dnn_v1/include -isystem external/mkl_dnn_v1/src -isystem bazel-out/darwin-opt/bin/external/mkl_dnn_v1/src -isystem external/mkl_dnn_v1/src/common -isystem bazel-out/darwin-opt/bin/external/mkl_dnn_v1/src/common -isystem external/mkl_dnn_v1/src/cpu -isystem bazel-out/darwin-opt/bin/external/mkl_dnn_v1/src/cpu -isystem external/mkl_dnn_v1/src/cpu/gemm -isystem bazel-out/darwin-opt/bin/external/mkl_dnn_v1/src/cpu/gemm -isystem external/mkl_dnn_v1/src/cpu/gemm/f32 -isystem bazel-out/darwin-opt/bin/external/mkl_dnn_v1/src/cpu/gemm/f32 -isystem external/mkl_dnn_v1/src/cpu/gemm/s8x8s32 -isystem bazel-out/darwin-opt/bin/external/mkl_dnn_v1/src/cpu/gemm/s8x8s32 -isystem external/mkl_dnn_v1/src/cpu/matmul -isystem bazel-out/darwin-opt/bin/external/mkl_dnn_v1/src/cpu/matmul -isystem external/mkl_dnn_v1/src/cpu/rnn -isystem bazel-out/darwin-opt/bin/external/mkl_dnn_v1/src/cpu/rnn -isystem external/mkl_dnn_v1/src/cpu/x64 -isystem bazel-out/darwin-opt/bin/external/mkl_dnn_v1/src/cpu/x64 -isystem external/mkl_dnn_v1/src/cpu/x64/jit_utils -isystem bazel-out/darwin-opt/bin/external/mkl_dnn_v1/src/cpu/x64/jit_utils -isystem external/mkl_dnn_v1/src/cpu/x64/jit_utils/jitprofiling -isystem bazel-out/darwin-opt/bin/external/mkl_dnn_v1/src/cpu/x64/jit_utils/jitprofiling -isystem external/mkl_dnn_v1/src/cpu/x64/xbyak -isystem bazel-out/darwin-opt/bin/external/mkl_dnn_v1/src/cpu/x64/xbyak -MD -MF bazel-out/darwin-opt/bin/external/mkl_dnn_v1/_objs/dnnl_single_threaded/jitprofiling.d '-frandom-seed=bazel-out/darwin-opt/bin/external/mkl_dnn_v1/_objs/dnnl_single_threaded/jitprofiling.o' -isysroot __BAZEL_XCODE_SDKROOT__ -F__BAZEL_XCODE_SDKROOT__/System/Library/Frameworks -F__BAZEL_XCODE_DEVELOPER_DIR__/Platforms/MacOSX.platform/Developer/Library/Frameworks '-mmacosx-version-min=10.9' -Wno-sign-compare -fexceptions -DDNNL_ENABLE_MAX_CPU_ISA -no-canonical-prefixes -Wno-builtin-macro-redefined '-D__DATE__="redacted"' '-D__TIMESTAMP__="redacted"' '-D__TIME__="redacted"' -c external/mkl_dnn_v1/src/cpu/x64/jit_utils/jitprofiling/jitprofiling.c -o bazel-out/darwin-opt/bin/external/mkl_dnn_v1/_objs/dnnl_single_threaded/jitprofiling.o)
Execution platform: @local_execution_config_platform//:platform
In file included from external/mkl_dnn_v1/src/cpu/x64/jit_utils/jitprofiling/jitprofiling.c:38:
external/mkl_dnn_v1/src/cpu/x64/jit_utils/jitprofiling/ittnotify_config.h:349:12: error: implicit declaration of function '__TBB_machine_fetchadd4' is invalid in C99 [-Werror,-Wimplicit-function-declaration]
    return __TBB_machine_fetchadd4(ptr, 1) + 1L;
           ^
1 error generated.

[ericmjl]

@xhochy I'm going to pull in another person who's successfully built on macOS ARM before: @jotsif.

@jotsif, I realize this might be a bit abrupt pulling you in like this. That said, I'm wondering if you might be kind enough to help us debug what's going wrong on the build? Or perhaps could you comment here whether you got things working with an officially-released bazel/TensorFlow, or if you had to rely on a specific commit instead? (The latter is what I'm guessing happened in this issue.) Getting things working here will help us get jaxlib working for the entire community 😄.

[jotsif]

I got that working by adding --copt=-Wno-implicit-function-declaration to the bazel build configs.

Regarding getting jaxlib to work in general I have some difficulties with my build though, in runtime. Looking into that and will update the issue on Jax with details.

[xhochy]

Thanks! With the information on the linked issue, I was able to build this fully. Waiting for conda-forge/bazel-feedstock#60 to turn green, then I'll push things here. My compilation is still only directly on the M1, we need cross-compilation here but we'll get there soon.

[xhochy]

Merged master and the build now fails with the same error as the latest oneDNN build also fails on osx-arm64:

n file included from external/mkl_dnn_v1/src/cpu/cpu_batch_normalization_list.cpp:32:
In file included from external/mkl_dnn_v1/src/cpu/aarch64/jit_uni_batch_normalization.hpp:28:
In file included from external/mkl_dnn_v1/src/cpu/aarch64/cpu_isa_traits.hpp:33:
external/mkl_dnn_v1/src/cpu/aarch64/xbyak_aarch64/xbyak_aarch64_util.h:20:10: fatal error: 'sys/prctl.h' file not found
#include <sys/prctl.h>
         ^~~~~~~~~~~~~
1 error generated.

This is fixed on oneDNN's master, so I would rather wait / push for a oneDNN release.

[xhochy]

Meanwhile I'll work on adding a cross-compile toolchain for Bazel here.

[ericmjl]

Thanks @xhochy! I've added a "WIP" on the PR to make it clear that there's stuff still waiting.

[xhochy]

Cross-referencing conda-forge/onednn-feedstock#3 here. Once that has been solved, we should be able to proceed further here.

[regro-cf-autotick-bot]

Due to the bot-rerun label I'm closing this PR. I will make another one as appropriate. This was generated by https://github.com/regro/autotick-bot/actions/runs/1035132257