checks.yml

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-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-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-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)