migrate to OpenXLA (#393)

* use PJRT to compile and execute graph
* allow users to choose runtime device, via wrapped PJRT plugins
* allow users to create and use custom PJRT plugins
* remove source dependency on TensorFlow
* build all binaries in spidr, rather than using elixir-nx
* build (not test) spidr for CUDA in CI
* automatically install support libraries to idris path (not global for user)
* add tutorial on spidr internals

.github/workflows/checks-backend.yml → .github/workflows/c-xla-version.yml

@@ -2,14 +2,15 @@ on:
   pull_request:
     branches: master
     paths:
-      - 'backend/**'
+      - 'spidr/backend/**'
+      - 'XLA_VERSION'
 
 jobs:
   backend-version:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
       with:
         fetch-depth: 2
     - name: Check backend version is updated
-      run: if [[ -z $(git diff HEAD~1 backend/VERSION) ]]; then exit 1; fi
+      run: sh -c "! git diff --exit-code HEAD^ spidr/backend/VERSION"

.github/workflows/checks.yml

@@ -2,43 +2,190 @@ on:
   pull_request:
     branches: master
 
-env:
-  PACK_DIR: /root/.pack
-
 jobs:
-  test-unit:
+  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:
+    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 "https://github.com/joelberkeley/spidr/releases/download/c-xla-v$(cat spidr/backend/VERSION)/libc_xla-linux.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.so libc_xla.so
+    - name: Upload binary
+      uses: actions/upload-artifact@v4
+      with:
+        name: libc_xla.so
+        path: libc_xla.so
+        if-no-files-found: error
+  pjrt-plugin-xla-cpu-linux:
+    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 "https://github.com/joelberkeley/spidr/releases/download/xla-$(short_revision $rev)/pjrt_plugin_xla_cpu-linux.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.so pjrt_plugin_xla_cpu.so
+    - name: Upload binary
+      uses: actions/upload-artifact@v4
+      with:
+        name: pjrt_plugin_xla_cpu.so
+        path: pjrt_plugin_xla_cpu.so
+        if-no-files-found: error
+  pjrt-plugin-xla-cuda-linux:
+    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 "https://github.com/joelberkeley/spidr/releases/download/xla-$(short_revision $rev)/pjrt_plugin_xla_cuda-linux.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.so pjrt_plugin_xla_cuda.so
+      - name: Upload binary
+        uses: actions/upload-artifact@v4
+        with:
+          name: pjrt_plugin_xla_cuda.so
+          path: pjrt_plugin_xla_cuda.so
+          if-no-files-found: error
+  build-tests-xla-cpu:
     runs-on: ubuntu-latest
     container: ghcr.io/stefan-hoeck/idris2-pack
     steps:
-    - uses: actions/checkout@v3
-    - name: Build backend
-      working-directory: backend
+    - uses: actions/checkout@v4
+    - name: Build tests
+      working-directory: test/xla-cpu
       run: |
-        ./install_bazel.sh
-        ./build.sh cpu
+        apt-get update && apt-get install -y curl
+        pack switch HEAD
+        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
+        path: test/xla-cpu/tests-xla-cpu.tar.gz
+        if-no-files-found: error
+  build-tests-xla-cuda:
+    runs-on: ubuntu-latest
+    container: ghcr.io/stefan-hoeck/idris2-pack
+    steps:
+    - uses: actions/checkout@v4
     - name: Build tests
+      working-directory: test/xla-cuda
       run: |
+        apt-get update && apt-get install -y curl
         pack switch HEAD
-        pack --no-prompt build test.ipkg
+        pack --no-prompt build xla-cuda.ipkg
+        tar cfz tests-xla-cuda.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
+        if-no-files-found: error
+  test-xla-cpu-linux:
+    needs:
+      - pjrt-linux
+      - pjrt-plugin-xla-cpu-linux
+      - build-tests-xla-cpu
+    runs-on: ubuntu-latest
+    steps:
+    - name: Download artifacts
+      uses: actions/download-artifact@v4
+      with:
+        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-cuda-linux:
+    needs:
+      - pjrt-linux
+      - pjrt-plugin-xla-cuda-linux
+      - build-tests-xla-cuda
+    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:
+        merge-multiple: true
+    - name: Install runtime dependencies
+      run: |
+        apt-get update && apt-get install chezscheme
     - name: Run tests
-      run: pack run test.ipkg
-      env:
-        LD_LIBRARY_PATH: $LD_LIBRARY_PATH:backend/bazel-bin
+      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@v3
+    - uses: actions/checkout@v4
     - name: Type-check README
       run: |
+        apt-get update && apt-get install -y curl
         pack switch HEAD
         pack --no-prompt typecheck readme.ipkg
   tutorials:
     runs-on: ubuntu-latest
     container: ghcr.io/stefan-hoeck/idris2-pack
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
     - name: Type-check tutorials
       run: |
+        apt-get update && apt-get install -y curl
         pack switch HEAD
         res=0; for f in tutorials/*.ipkg; do pack --no-prompt typecheck $f || res=$?; done; $(exit $res)

.github/workflows/release-c-xla.yml

@@ -0,0 +1,42 @@
+on:
+  workflow_dispatch:
+  push:
+    branches: master
+    paths:
+      - 'spidr/backend/VERSION'
+
+jobs:
+  create-release:
+    runs-on: ubuntu-latest
+    outputs:
+      upload_url: ${{ steps.create-release.outputs.upload_url }}
+    steps:
+    - uses: actions/checkout@v4
+    - name: Read C XLA version
+      run: echo "C_XLA_VERSION=$(cat spidr/backend/VERSION)" >> $GITHUB_ENV
+    - name: Create C XLA release
+      id: create-release
+      uses: actions/create-release@v1
+      env:
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      with:
+        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:
+    needs: create-release
+    runs-on: ubuntu-latest
+    container: tensorflow/build:latest-python3.9
+    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-linux.so
+        asset_name: libc_xla-linux.so
+        asset_content_type: application/x-sharedlib

.github/workflows/release-docs.yml

@@ -12,11 +12,11 @@ jobs:
     runs-on: ubuntu-latest
     container: ghcr.io/stefan-hoeck/idris2-pack
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
     - run: |
         pack switch HEAD
         pack --no-prompt --with-docs install spidr
-        mv build/docs .
+        mv spidr/build/docs .
     - run: |
         git config --global --add safe.directory "$GITHUB_WORKSPACE"
         git add .