support Darwin Aarch64/ARM (#427)

joelberkeley · Nov 29, 2024 · 33e3cf7 · 33e3cf7
1 parent b68d523
commit 33e3cf7
Show file tree

Hide file tree

Showing 15 changed files with 267 additions and 75 deletions.
diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml
@@ -14,7 +14,7 @@ jobs:
         run: |
           shopt -s extglob nullglob globstar
           shellcheck **/*.sh
-  pjrt-linux:
+  pjrt-linux-x86_64:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
@@ -23,7 +23,7 @@ jobs:
     - 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.so"
+          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
@@ -34,14 +34,35 @@ jobs:
             sh -c "spidr/backend/build.sh"
         fi
 
-        mv libc_xla-linux.so libc_xla.so
+        mv libc_xla-linux-x86_64.so libc_xla.so
     - name: Upload binary
       uses: actions/upload-artifact@v4
       with:
-        name: libc_xla.so
+        name: libc_xla-linux-x86_64
         path: libc_xla.so
         if-no-files-found: error
-  pjrt-plugin-xla-cpu-linux:
+  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
@@ -52,7 +73,7 @@ jobs:
         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.so"
+          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 \
@@ -61,14 +82,37 @@ jobs:
             sh -c "pjrt-plugins/xla-cpu/build.sh"
         fi
 
-        mv pjrt_plugin_xla_cpu-linux.so pjrt_plugin_xla_cpu.so
+        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.so
+        name: pjrt_plugin_xla_cpu-linux-x86_64
         path: pjrt_plugin_xla_cpu.so
         if-no-files-found: error
-  pjrt-plugin-xla-cuda-linux:
+  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-xla-cuda-linux-x86_64:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
@@ -79,7 +123,7 @@ jobs:
           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.so"
+            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
@@ -89,59 +133,85 @@ jobs:
               sh -c "pjrt-plugins/xla-cuda/build.sh"
           fi
 
-          mv pjrt_plugin_xla_cuda-linux.so pjrt_plugin_xla_cuda.so
+          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.so
+          name: pjrt_plugin_xla_cuda-linux-x86_64
           path: pjrt_plugin_xla_cuda.so
           if-no-files-found: error
-  build-tests-xla-cpu:
+  build-tests-xla-cpu-linux-x86_64:
     runs-on: ubuntu-latest
     container: ghcr.io/stefan-hoeck/idris2-pack
     steps:
     - uses: actions/checkout@v4
-    - name: Build tests
-      working-directory: test/xla-cpu
+    - 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.tar.gz
+        name: tests-xla-cpu-linux-x86_64
         path: test/xla-cpu/tests-xla-cpu.tar.gz
         if-no-files-found: error
-  build-tests-xla-cuda:
+  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-xla-cuda-linux-x86_64:
     runs-on: ubuntu-latest
     container: ghcr.io/stefan-hoeck/idris2-pack
     steps:
     - uses: actions/checkout@v4
-    - name: Build tests
-      working-directory: test/xla-cuda
+    - 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.tar.gz -C build/exec .
+        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.tar.gz
-        path: test/xla-cuda/tests-xla-cuda.tar.gz
+        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:
+  test-xla-cpu-linux-x86_64:
     needs:
-      - pjrt-linux
-      - pjrt-plugin-xla-cpu-linux
-      - build-tests-xla-cpu
+      - 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: |
@@ -150,17 +220,37 @@ jobs:
       run: |
         tar xfz tests-xla-cpu.tar.gz . && rm tests-xla-cpu.tar.gz
         ./test
-  test-xla-cuda-linux:
+  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-xla-cuda-linux-x86_64:
     needs:
-      - pjrt-linux
-      - pjrt-plugin-xla-cuda-linux
-      - build-tests-xla-cuda
+      - 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: |

diff --git a/.github/workflows/release-c-xla.yml b/.github/workflows/release-c-xla.yml
@@ -23,7 +23,7 @@ jobs:
         tag_name: c-xla-v${{ env.C_XLA_VERSION }}
         release_name: C XLA release ${{ env.C_XLA_VERSION }}
         body: "C XLA release ${{ env.C_XLA_VERSION }}"
-  c-xla:
+  c-xla-linux-x86_64:
     needs: create-release
     runs-on: ubuntu-latest
     container: tensorflow/build:latest-python3.9
@@ -37,6 +37,22 @@ jobs:
         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       with:
         upload_url: ${{ needs.create-release.outputs.upload_url }}
-        asset_path: libc_xla-linux.so
-        asset_name: libc_xla-linux.so
+        asset_path: libc_xla-linux-x86_64.so
+        asset_name: libc_xla-linux-x86_64.so
         asset_content_type: application/x-sharedlib
+  c-xla-darwin-aarch64:
+    needs: create-release
+    runs-on: macos-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Build C XLA lib
+        run: ./spidr/backend/build.sh
+      - name: Publish C XLA lib
+        uses: actions/upload-release-asset@v1
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        with:
+          upload_url: ${{ needs.create-release.outputs.upload_url }}
+          asset_path: libc_xla-darwin-aarch64.dylib
+          asset_name: libc_xla-darwin-aarch64.dylib
+          asset_content_type: application/x-sharedlib
diff --git a/.github/workflows/release-xla.yml b/.github/workflows/release-xla.yml
@@ -26,12 +26,12 @@ jobs:
         tag_name: xla-${{ env.XLA_REV }}
         release_name: XLA revision ${{ env.XLA_REV }}
         body: "XLA revision ${{ env.XLA_REV }}"
-  pjrt-plugin-xla-cpu-linux:
+  pjrt-plugin-xla-cpu-linux-x86_64:
     needs: create-release
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
-    - name: Build XLA CPU plugin for Linux
+    - name: Build XLA CPU plugin for Linux x86_64
       run: |
         rm -rf /opt/hostedtoolcache
         docker run \
@@ -44,15 +44,31 @@ jobs:
         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       with:
         upload_url: ${{ needs.create-release.outputs.upload_url }}
-        asset_path: pjrt_plugin_xla_cpu-linux.so
-        asset_name: pjrt_plugin_xla_cpu-linux.so
+        asset_path: pjrt_plugin_xla_cpu-linux-x86_64.so
+        asset_name: pjrt_plugin_xla_cpu-linux-x86_64.so
         asset_content_type: application/x-sharedlib
-  pjrt-plugin-xla-cuda-linux:
+  pjrt-plugin-xla-cpu-darwin-aarch64:
+    needs: create-release
+    runs-on: macos-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Build XLA CPU plugin for Darwin AArch64
+      run: ./pjrt-plugins/xla-cpu/build.sh
+    - name: Publish XLA CPU plugin
+      uses: actions/upload-release-asset@v1
+      env:
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      with:
+        upload_url: ${{ needs.create-release.outputs.upload_url }}
+        asset_path: pjrt_plugin_xla_cpu-darwin-aarch64.dylib
+        asset_name: pjrt_plugin_xla_cpu-darwin-aarch64.dylib
+        asset_content_type: application/x-sharedlib
+  pjrt-plugin-xla-cuda-linux-x86_64:
     needs: create-release
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
-    - name: Build XLA CUDA plugin for Linux
+    - name: Build XLA CUDA plugin for Linux x86_64
       run: |
         rm -rf /opt/hostedtoolcache
         docker run \
@@ -65,6 +81,6 @@ jobs:
         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       with:
         upload_url: ${{ needs.create-release.outputs.upload_url }}
-        asset_path: pjrt_plugin_xla_cuda-linux.so
-        asset_name: pjrt_plugin_xla_cuda-linux.so
+        asset_path: pjrt_plugin_xla_cuda-linux-x86_64.so
+        asset_name: pjrt_plugin_xla_cuda-linux-x86_64.so
         asset_content_type: application/x-sharedlib
diff --git a/.gitignore b/.gitignore
@@ -9,4 +9,5 @@ bazel
 bazel-*
 bazelisk
 *.so
+*.dylib
 *.tar.gz
diff --git a/README.md b/README.md
@@ -8,14 +8,12 @@ See the [online reference](https://joelberkeley.github.io/spidr/) for API docume
 
 ### Install
 
-Installation requires curl and [pack](https://github.com/stefan-hoeck/idris2-pack). To install spidr, install a PJRT plugin. A plugin executes a spidr program, and determines what hardware your program will run on. You can install the CPU plugin with
+Installation requires curl and [pack](https://github.com/stefan-hoeck/idris2-pack). To install spidr, install a PJRT plugin. A plugin executes a spidr program, and determines what hardware your program will run on. You can install the CPU plugin, on Linux or Mac with Apple silicon, with
 ```
 pack install pjrt-plugin-xla-cpu
 ```
 or read the [plugin documentation](pjrt-plugins/README.md) for the CUDA-enabled GPU plugin and custom plugin builds.
 
-We have tested spidr on Ubuntu Linux 22.04. It may work on other Linux distributions. If spidr doesn't work on your platform, [get in touch](#contact).
-
 ### Motivation
 
 We made spidr to try out modern programming language capabilities in machine learning systems. To this end, we chose [Idris](https://github.com/idris-lang/Idris2) for the API; Idris is a general-purpose purely functional programming language with a particularly expressive type system. We also wanted to build something performant enough for working machine learning practitioners. Implementing efficient low-level linear algebra is not one of the project goals, so we opted to use [OpenXLA](https://openxla.org/)'s PJRT, an abstract interface for machine learning compilers for hardware accelerators.

diff --git a/XLA_VERSION b/XLA_VERSION
@@ -1 +1 @@
-43517d94ad963a96a2a308b7b33d77ecd7de4b4a
+46bd34518db14456697224b9684f47bf4cdcf6ba
diff --git a/pjrt-plugins/xla-cpu/README.md b/pjrt-plugins/xla-cpu/README.md
@@ -4,7 +4,7 @@ This is the PJRT plugin for CPU. It uses the XLA compiler.
 
 ## Install
 
-On Linux, run
+On Linux, or MacOS with AArch64 (ARM64, Apple Silicon), run
 ```
 pack install pjrt-plugin-xla-cpu
 ```
-Original file line number
+Diff line change
@@ Expand Up / @@ -9,4 +9,5 @@ bazel @@
     bazel-*
     bazelisk
     *.so
+    *.dylib
     *.tar.gz
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		43517d94ad963a96a2a308b7b33d77ecd7de4b4a
		46bd34518db14456697224b9684f47bf4cdcf6ba