_build.yaml

name: Build PJRT GPU libraries

on:
  workflow_call:
    secrets:
      BUILDBUDDY_API_KEY:
        required: true
    inputs:
      xla_commit:
        required: false
        type: string
      rocm_xla_commit:
        required: false
        type: string
      bazel_extra_args:
        required: false
        type: string

env:
  XLA_COMMIT: ${{ inputs.xla_commit || '9e9a0fb13b164c9875349246d1d8616a267c9bf5' }} # main from 2026-03-12
  ROCM_XLA_COMMIT: ${{ inputs.rocm_xla_commit || 'c40171e49b63e8a12498fae2b64f90c2406802ec' }} # rocm-jaxlib-v0.8.0 from 2025-12-30
  TF_ROCM_AMDGPU_TARGETS: "gfx942,gfx950,gfx1100,gfx1101,gfx1150,gfx1151,gfx1200,gfx1201"
  BUILDBUDDY_API_KEY: ${{ secrets.BUILDBUDDY_API_KEY }}

jobs:
  pjrt-artifacts:
    runs-on: ${{ matrix.pjrt.runs_on }}
    strategy:
      fail-fast: false
      matrix:
        pjrt:
          - target: cuda
            config: "--enable_workspace --spawn_strategy=local --config=cuda_nvcc" 
            artifact: bazel-bin/xla/pjrt/c/libpjrt_c_api_gpu_plugin.so
            renamed_artifact: libpjrt_cuda.so
            runs_on: ["runs-on", "runner=32cpu-linux-x64", "image=ubuntu22-amd64"]
            platform: linux-amd64
            bazel_target: //xla/pjrt/c:pjrt_c_api_gpu_plugin
          - target: rocm
            config: "--spawn_strategy=local --disk_cache=/home/runner/.cache/bazel-disk --config=rocm_ci"
            artifact: bazel-bin/xla/pjrt/c/libpjrt_c_api_gpu_plugin.so
            renamed_artifact: libpjrt_rocm.so
            runs_on: ["runs-on", "runner=32cpu-linux-x64", "image=ubuntu22-amd64"]
            platform: linux-amd64
            bazel_target: //xla/pjrt/c:pjrt_c_api_gpu_plugin
          - target: cpu
            config: "--config=hermetic_macos_arm64 --config=remote --config=bzlmod --remote_header=x-buildbuddy-api-key=$BUILDBUDDY_API_KEY"
            artifact: bazel-bin/xla/pjrt/c/libpjrt_c_api_cpu_plugin.dylib
            runs_on: ["runs-on", "runner=32cpu-linux-x64", "image=ubuntu22-amd64"]
            renamed_artifact: libpjrt_cpu.dylib
            platform: darwin-arm64
            bazel_target: //xla/pjrt/c:pjrt_c_api_cpu_plugin
          - target: cpu
            config: "--config=hermetic_macos_x86 --config=remote --config=bzlmod --remote_header=x-buildbuddy-api-key=$BUILDBUDDY_API_KEY --//xla/tsl/framework/contraction:disable_onednn_contraction_kernel=True"
            artifact: bazel-bin/xla/pjrt/c/libpjrt_c_api_cpu_plugin.dylib
            renamed_artifact: libpjrt_cpu.dylib
            runs_on: ["runs-on", "runner=32cpu-linux-x64", "image=ubuntu22-amd64"]
            platform: darwin-amd64
            bazel_target: //xla/pjrt/c:pjrt_c_api_cpu_plugin
          - target: cpu
            config: "--config=hermetic_linux_x86 --config=remote --config=bzlmod --remote_header=x-buildbuddy-api-key=$BUILDBUDDY_API_KEY"
            artifact: bazel-bin/xla/pjrt/c/libpjrt_c_api_cpu_plugin.so
            renamed_artifact: libpjrt_cpu.so
            runs_on: ["runs-on", "runner=32cpu-linux-x64", "image=ubuntu22-amd64"]
            platform: linux-amd64
            bazel_target: //xla/pjrt/c:pjrt_c_api_cpu_plugin
          - target: cpu
            config: "--config=hermetic_linux_arm64 --config=remote --config=bzlmod --remote_header=x-buildbuddy-api-key=$BUILDBUDDY_API_KEY"
            artifact: bazel-bin/xla/pjrt/c/libpjrt_c_api_cpu_plugin.so
            renamed_artifact: libpjrt_cpu.so
            runs_on: ["runs-on", "runner=32cpu-linux-x64", "image=ubuntu22-amd64"]
            platform: linux-arm64
            bazel_target: //xla/pjrt/c:pjrt_c_api_cpu_plugin
    steps:
      - uses: runs-on/action@v1

      - name: "Checking out repository"
        uses: actions/checkout@v4
        with:
          path: "pjrt-artifacts"

      - name: "Checking out openxla repository"
        uses: actions/checkout@v4
        with:
          ref: ${{ matrix.pjrt.target == 'rocm' && env.ROCM_XLA_COMMIT || env.XLA_COMMIT }}
          repository: ${{ matrix.pjrt.target == 'rocm' && 'ROCm/xla' || 'openxla/xla' }}
          path: "xla"

      - name: Apply patches to openxla
        id: patches
        working-directory: ./xla
        env:
          PATCH_DIRECTORY: ${{ matrix.pjrt.target == 'rocm' && 'rocm' || 'upstream' }}
        run: |
          xla_commit=$(git rev-parse HEAD)
          echo "XLA_COMMIT_ID=$xla_commit" >> $GITHUB_OUTPUT
          echo ::notice::Applying patches to openxla $xla_commit
          if compgen -G "../pjrt-artifacts/openxla/patches/$PATCH_DIRECTORY/*.patch" > /dev/null; then
            for patch in ../pjrt-artifacts/openxla/patches/$PATCH_DIRECTORY/*.patch; do
                echo "Applying patch $patch"
                git apply "$patch"
            done
          fi
          # apply patches from upstream-cpu if target = cpu
          if [ "${{ matrix.pjrt.target }}" = "cpu" ]; then
            if compgen -G "../pjrt-artifacts/openxla/patches/upstream-cpu/*.patch" > /dev/null; then
              for patch in ../pjrt-artifacts/openxla/patches/upstream-cpu/*.patch; do
                  echo "Applying patch $patch"
                  git apply "$patch"
              done
            fi
          fi

      - uses: runs-on/snapshot@v1
        if: matrix.pjrt.target != 'cpu'
        with:
          path: /home/runner/.cache/bazel-disk
          volume_size: 350

      - name: TMP fix for the mount user
        if: matrix.pjrt.target != 'cpu'
        run: sudo chown -R runner:runner /home/runner/.cache/

      - name: Download ROCm toolchain (not fully hermetic)
        if: matrix.pjrt.target == 'rocm'
        run: |
          export DEBIAN_FRONTEND=noninteractive
          sudo bash -c "$(wget -O - https://apt.llvm.org/llvm.sh)" -- 18
          sudo apt update -y
          sudo apt install llvm-18 lld-18 clang-18 libnuma-dev -y
          wget https://repo.radeon.com/amdgpu-install/7.1.1/ubuntu/jammy/amdgpu-install_7.1.1.70101-1_all.deb
          sudo apt install ./amdgpu-install_7.1.1.70101-1_all.deb -y
          sudo amdgpu-install --usecase=rocm,rocmdev,hiplibsdk -y --no-dkms

      - uses: bazel-contrib/setup-bazel@0.15.0
        with:
          bazelisk-version: 1.26.0
          bazelisk-cache: false
          bazelrc: |
            common --verbose_failures
            common --color=yes
            common --show_timestamps

      - name: "Setup bazelrc for openxla"
        env:
          BAZERLC_DIRECTORY: ${{ matrix.pjrt.target == 'rocm' && 'rocm' || 'upstream' }}
        run: |
          # if .bazelrc exists copy it
          if [ -f "pjrt-artifacts/openxla/bazelrc/${BAZERLC_DIRECTORY}/.bazelrc" ]; then
            cp pjrt-artifacts/openxla/bazelrc/${BAZERLC_DIRECTORY}/.bazelrc xla/
          fi
          cp pjrt-artifacts/openxla/bazelrc/${BAZERLC_DIRECTORY}/*.bazelrc xla/

      - name: "Build ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }}"
        working-directory: ./xla
        env:
          USE_BAZEL_VERSION: ${{ matrix.pjrt.target == 'cpu' && '8.5.1' || '7.7.0' }}
          BAZEL_EXTRA_ARGS: ${{ inputs.bazel_extra_args || '' }}
        run: |
          bazel build ${{ matrix.pjrt.config }} ${{ env.BAZEL_EXTRA_ARGS }} ${{ matrix.pjrt.bazel_target }}

      - name: Set file at the right path
        working-directory: ./xla
        run: |
          cp ${{ matrix.pjrt.artifact }} ../${{ matrix.pjrt.renamed_artifact }}

      - name: Create compressed ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} release file
        uses: a7ul/tar-action@v1.2.0
        with:
          command: c
          cwd: ./
          files: |
            ./${{ matrix.pjrt.renamed_artifact }}
          outPath: pjrt-${{ matrix.pjrt.target }}_${{ matrix.pjrt.platform }}.tar.gz

      - name: Upload ${{ matrix.pjrt.target }} ${{ matrix.pjrt.platform }} artifacts
        uses: actions/upload-artifact@v4
        with:
          name: pjrt-${{ matrix.pjrt.target }}_${{ matrix.pjrt.platform }}.tar.gz
          path: pjrt-${{ matrix.pjrt.target }}_${{ matrix.pjrt.platform }}.tar.gz