test: stablehlo + apple metal #3162
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| on: | |
| pull_request: | |
| branches: master | |
| jobs: | |
| shellcheck: | |
| runs-on: ubuntu-latest | |
| steps: | |
| - uses: actions/checkout@v4 | |
| - name: Install shellcheck | |
| run: sudo apt-get install shellcheck | |
| - name: Run shellcheck | |
| shell: bash | |
| run: | | |
| shopt -s extglob nullglob globstar | |
| shellcheck **/*.sh | |
| pjrt-linux-x86_64: | |
| runs-on: ubuntu-latest | |
| steps: | |
| - uses: actions/checkout@v4 | |
| with: | |
| fetch-depth: 2 | |
| - name: Build PJRT + XLA binary | |
| run: | | |
| if [ ! "$(git diff --exit-code HEAD^ spidr/backend/VERSION)" ]; then | |
| curl -LO --fail-with-body "https://github.com/joelberkeley/spidr/releases/download/c-xla-v$(cat spidr/backend/VERSION)/libc_xla-linux-x86_64.so" | |
| else | |
| # free up space not used if running in Docker, see | |
| # https://github.com/orgs/community/discussions/25678#discussioncomment-5242449 | |
| rm -rf /opt/hostedtoolcache | |
| docker run \ | |
| -v $(pwd):/spidr -w /spidr \ | |
| tensorflow/build:latest-python3.9 \ | |
| sh -c "spidr/backend/build.sh" | |
| fi | |
| mv libc_xla-linux-x86_64.so libc_xla.so | |
| - name: Upload binary | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: libc_xla-linux-x86_64 | |
| path: libc_xla.so | |
| if-no-files-found: error | |
| pjrt-darwin-aarch64: | |
| runs-on: macos-latest | |
| steps: | |
| - uses: actions/checkout@v4 | |
| with: | |
| fetch-depth: 2 | |
| - name: Build PJRT + XLA binary | |
| run: | | |
| if [ ! "$(git diff --exit-code HEAD^ spidr/backend/VERSION)" ]; then | |
| curl -LO --fail-with-body "https://github.com/joelberkeley/spidr/releases/download/c-xla-v$(cat spidr/backend/VERSION)/libc_xla-darwin-aarch64.dylib" | |
| else | |
| ./spidr/backend/build.sh | |
| fi | |
| mv libc_xla-darwin-aarch64.dylib libc_xla.dylib | |
| - name: Upload binary | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: libc_xla-darwin-aarch64 | |
| path: libc_xla.dylib | |
| if-no-files-found: error | |
| pjrt-plugin-xla-cpu-linux-x86_64: | |
| runs-on: ubuntu-latest | |
| steps: | |
| - uses: actions/checkout@v4 | |
| with: | |
| fetch-depth: 2 | |
| - name: Build or fetch XLA CPU PJRT plugin | |
| run: | | |
| if [ ! "$(git diff --exit-code HEAD^ XLA_VERSION)" ]; then | |
| . ./dev.sh | |
| rev=$(cat XLA_VERSION) | |
| curl -LO --fail-with-body "https://github.com/joelberkeley/spidr/releases/download/xla-$(short_revision $rev)/pjrt_plugin_xla_cpu-linux-x86_64.so" | |
| else | |
| rm -rf /opt/hostedtoolcache | |
| docker run \ | |
| -v $(pwd):/spidr -w /spidr \ | |
| tensorflow/build:latest-python3.9 \ | |
| sh -c "pjrt-plugins/xla-cpu/build.sh" | |
| fi | |
| mv pjrt_plugin_xla_cpu-linux-x86_64.so pjrt_plugin_xla_cpu.so | |
| - name: Upload binary | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: pjrt_plugin_xla_cpu-linux-x86_64 | |
| path: pjrt_plugin_xla_cpu.so | |
| if-no-files-found: error | |
| pjrt-plugin-xla-cpu-darwin-aarch64: | |
| runs-on: macos-latest | |
| steps: | |
| - uses: actions/checkout@v4 | |
| with: | |
| fetch-depth: 2 | |
| - name: Build or fetch XLA CPU PJRT plugin | |
| run: | | |
| if [ ! "$(git diff --exit-code HEAD^ XLA_VERSION)" ]; then | |
| . ./dev.sh | |
| rev=$(cat XLA_VERSION) | |
| curl -LO --fail-with-body "https://github.com/joelberkeley/spidr/releases/download/xla-$(short_revision $rev)/pjrt_plugin_xla_cpu-darwin-aarch64.dylib" | |
| else | |
| ./pjrt-plugins/xla-cpu/build.sh | |
| fi | |
| mv pjrt_plugin_xla_cpu-darwin-aarch64.dylib pjrt_plugin_xla_cpu.dylib | |
| - name: Upload binary | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: pjrt_plugin_xla_cpu-darwin-aarch64 | |
| path: pjrt_plugin_xla_cpu.dylib | |
| if-no-files-found: error | |
| pjrt-plugin-apple-metal: | |
| runs-on: macos-latest | |
| steps: | |
| - uses: actions/checkout@v4 | |
| with: | |
| fetch-depth: 2 | |
| - name: Build or fetch Apple Metal PJRT plugin | |
| run: | | |
| prefix=jax_metal-0.1.0-py3-none-macosx_11_0_arm64 | |
| curl -fsL "https://files.pythonhosted.org/packages/80/af/ed482a421a868726e7ca3f51ac19b0c9a8e37f33f54413312c37e9056acc/jax_metal-0.1.0-py3-none-macosx_11_0_arm64.whl" \ | |
| -o "$prefix.zip" | |
| unzip "$prefix.zip" | |
| mv "jax_plugins/metal_plugin/pjrt_plugin_metal_14.dylib" pjrt_plugin_apple_metal.dylib | |
| - name: Upload binary | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: pjrt_plugin_apple_metal | |
| path: pjrt_plugin_apple_metal.dylib | |
| if-no-files-found: error | |
| pjrt-plugin-xla-cuda-linux-x86_64: | |
| runs-on: ubuntu-latest | |
| steps: | |
| - uses: actions/checkout@v4 | |
| with: | |
| fetch-depth: 2 | |
| - name: Build or fetch XLA CUDA PJRT plugin | |
| run: | | |
| if [ ! "$(git diff --exit-code HEAD^ XLA_VERSION)" ]; then | |
| . ./dev.sh | |
| rev=$(cat XLA_VERSION) | |
| curl -LO --fail-with-body "https://github.com/joelberkeley/spidr/releases/download/xla-$(short_revision $rev)/pjrt_plugin_xla_cuda-linux-x86_64.so" | |
| else | |
| rm -rf /opt/hostedtoolcache | |
| # note this implies specific versions of CUDA and cuDNN | |
| docker run \ | |
| -v $(pwd):/spidr -w /spidr \ | |
| tensorflow/build:latest-python3.9 \ | |
| sh -c "pjrt-plugins/xla-cuda/build.sh" | |
| fi | |
| mv pjrt_plugin_xla_cuda-linux-x86_64.so pjrt_plugin_xla_cuda.so | |
| - name: Upload binary | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: pjrt_plugin_xla_cuda-linux-x86_64 | |
| path: pjrt_plugin_xla_cuda.so | |
| if-no-files-found: error | |
| build-tests-xla-cpu-linux-x86_64: | |
| runs-on: ubuntu-latest | |
| container: ghcr.io/stefan-hoeck/idris2-pack | |
| steps: | |
| - uses: actions/checkout@v4 | |
| - name: Install build dependencies | |
| run: | | |
| apt-get update && apt-get install -y curl | |
| pack switch HEAD | |
| - name: Build tests | |
| working-directory: test/xla-cpu | |
| run: | | |
| SPIDR_INSTALL_SUPPORT_LIBS=false pack --no-prompt build xla-cpu.ipkg | |
| tar cfz tests-xla-cpu.tar.gz -C build/exec . | |
| - name: Upload tests | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: tests-xla-cpu-linux-x86_64 | |
| path: test/xla-cpu/tests-xla-cpu.tar.gz | |
| if-no-files-found: error | |
| build-tests-xla-cpu-darwin-aarch64: | |
| runs-on: macos-latest | |
| steps: | |
| - uses: actions/checkout@v4 | |
| - name: Install build dependencies | |
| run: | | |
| brew install chezscheme | |
| git clone https://github.com/stefan-hoeck/idris2-pack.git | |
| (cd idris2-pack && make micropack SCHEME=chez) | |
| ~/.pack/bin/pack switch HEAD | |
| - name: Build tests | |
| working-directory: test/xla-cpu | |
| run: | | |
| SPIDR_INSTALL_SUPPORT_LIBS=false ~/.pack/bin/pack --no-prompt build xla-cpu.ipkg | |
| tar cfz tests-xla-cpu.tar.gz -C build/exec . | |
| - name: Upload tests | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: tests-xla-cpu-darwin-aarch64 | |
| path: test/xla-cpu/tests-xla-cpu.tar.gz | |
| if-no-files-found: error | |
| build-tests-apple-metal: | |
| runs-on: macos-latest | |
| steps: | |
| - uses: actions/checkout@v4 | |
| - name: Install build dependencies | |
| run: | | |
| brew install chezscheme | |
| git clone https://github.com/stefan-hoeck/idris2-pack.git | |
| (cd idris2-pack && make micropack SCHEME=chez) | |
| ~/.pack/bin/pack switch HEAD | |
| - name: Build tests | |
| working-directory: test/apple-metal | |
| run: | | |
| SPIDR_INSTALL_SUPPORT_LIBS=false ~/.pack/bin/pack --no-prompt build apple-metal.ipkg | |
| tar cfz tests-apple-metal.tar.gz -C build/exec . | |
| - name: Upload tests | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: tests-apple-metal | |
| path: test/apple-metal/tests-apple-metal.tar.gz | |
| if-no-files-found: error | |
| build-tests-xla-cuda-linux-x86_64: | |
| runs-on: ubuntu-latest | |
| container: ghcr.io/stefan-hoeck/idris2-pack | |
| steps: | |
| - uses: actions/checkout@v4 | |
| - name: Install build dependencies | |
| run: | | |
| apt-get update && apt-get install -y curl | |
| pack switch HEAD | |
| - name: Build tests | |
| working-directory: test/xla-cuda | |
| run: | | |
| SPIDR_INSTALL_SUPPORT_LIBS=false pack --no-prompt build xla-cuda.ipkg | |
| tar cfz tests-xla-cuda-linux-x86_64.tar.gz -C build/exec . | |
| - name: Upload tests | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: tests-xla-cuda-linux-x86_64 | |
| path: test/xla-cuda/tests-xla-cuda-linux-x86_64.tar.gz | |
| if-no-files-found: error | |
| test-xla-cpu-linux-x86_64: | |
| needs: | |
| - pjrt-linux-x86_64 | |
| - pjrt-plugin-xla-cpu-linux-x86_64 | |
| - build-tests-xla-cpu-linux-x86_64 | |
| runs-on: ubuntu-latest | |
| steps: | |
| - name: Download artifacts | |
| uses: actions/download-artifact@v4 | |
| with: | |
| pattern: "{libc_xla-linux-x86_64,*cpu-linux-x86_64}" | |
| merge-multiple: true | |
| - name: Install runtime dependencies | |
| run: | | |
| sudo apt-get install chezscheme | |
| - name: Run tests | |
| run: | | |
| tar xfz tests-xla-cpu.tar.gz . && rm tests-xla-cpu.tar.gz | |
| ./test | |
| test-xla-cpu-darwin-aarch64: | |
| needs: | |
| - pjrt-darwin-aarch64 | |
| - pjrt-plugin-xla-cpu-darwin-aarch64 | |
| - build-tests-xla-cpu-darwin-aarch64 | |
| runs-on: macos-latest | |
| steps: | |
| - name: Download artifacts | |
| uses: actions/download-artifact@v4 | |
| with: | |
| pattern: "*darwin-aarch64" | |
| merge-multiple: true | |
| - name: Install runtime dependencies | |
| run: | | |
| brew install chezscheme | |
| - name: Run tests | |
| run: | | |
| tar xfz tests-xla-cpu.tar.gz && rm tests-xla-cpu.tar.gz | |
| ./test | |
| test-apple-metal: | |
| needs: | |
| - pjrt-darwin-aarch64 | |
| - pjrt-plugin-apple-metal | |
| - build-tests-apple-metal | |
| runs-on: macos-latest | |
| steps: | |
| - name: Download artifacts | |
| uses: actions/download-artifact@v4 | |
| with: | |
| pattern: "{libc_xla-darwin-aarch64,pjrt_plugin_apple_metal,tests-apple-metal}" | |
| merge-multiple: true | |
| - name: Install runtime dependencies | |
| run: | | |
| brew install chezscheme | |
| - name: Run tests | |
| run: | | |
| tar xfz tests-apple-metal.tar.gz && rm tests-apple-metal.tar.gz | |
| ./test | |
| test-xla-cuda-linux-x86_64: | |
| needs: | |
| - pjrt-linux-x86_64 | |
| - pjrt-plugin-xla-cuda-linux-x86_64 | |
| - build-tests-xla-cuda-linux-x86_64 | |
| runs-on: ubuntu-latest # needs a CUDA runner | |
| container: nvcr.io/nvidia/tensorrt:23.11-py3 | |
| steps: | |
| - name: Download artifacts | |
| uses: actions/download-artifact@v4 | |
| with: | |
| pattern: "{libc_xla-linux-x86_64,*cuda-linux-x86_64}" | |
| merge-multiple: true | |
| - name: Install runtime dependencies | |
| run: | | |
| apt-get update && apt-get install chezscheme | |
| - name: Run tests | |
| run: | | |
| tar xfz tests-xla-cuda.tar.gz . && rm tests-xla-cuda.tar.gz | |
| exit 0 # we can't run tests without a GPU | |
| readme: | |
| runs-on: ubuntu-latest | |
| container: ghcr.io/stefan-hoeck/idris2-pack | |
| steps: | |
| - uses: actions/checkout@v4 | |
| - name: Type-check README | |
| run: | | |
| apt-get update && apt-get install -y curl | |
| pack switch HEAD | |
| SPIDR_INSTALL_SUPPORT_LIBS=false pack --no-prompt typecheck readme.ipkg | |
| tutorials: | |
| runs-on: ubuntu-latest | |
| container: ghcr.io/stefan-hoeck/idris2-pack | |
| steps: | |
| - uses: actions/checkout@v4 | |
| - name: Type-check tutorials | |
| run: | | |
| apt-get update && apt-get install -y curl | |
| pack switch HEAD | |
| export SPIDR_INSTALL_SUPPORT_LIBS=false | |
| res=0; for f in tutorials/*.ipkg; do pack --no-prompt typecheck $f || res=$?; done; $(exit $res) |