# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

name: CI ROCm
permissions:
  contents: read
on:
  pull_request:
  workflow_dispatch:
  workflow_call:
    inputs:
      rocm_xla_branch:
        description: 'Branch name to fetch ROCm XLA scripts from'
        required: false
        default: 'rocm-dev-infra'
        type: string
      runner_label:
        description: 'Runner label for GPU jobs'
        required: false
        default: 'amd-do-linux-xla-gpu-gfx950-1'
        type: string

concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
  cancel-in-progress: ${{ github.ref != 'main' }}

jobs:
  rocm-config:
    runs-on: ubuntu-latest
    outputs:
      docker-image: ${{ steps.out.outputs.docker-image }}
      docker-image-jax: ${{ steps.out.outputs.docker-image-jax }}
    steps:
      - id: out
        shell: bash
        run: |
          # hermetic llvm based on  ghcr.io/rocm/jax-ubu22.rocm7.2.4:latest
          echo "docker-image=ghcr.io/rocm/jax-ubu22.rocm7.2.4@sha256:248db525342fb7438ccbc6f5f1396ec82c783f7e0ae5667cb698edcddc8fe2fd" >> "$GITHUB_OUTPUT"
          echo "docker-image-jax=ghcr.io/rocm/jax-ubu22.rocm7.2.4@sha256:248db525342fb7438ccbc6f5f1396ec82c783f7e0ae5667cb698edcddc8fe2fd" >> "$GITHUB_OUTPUT"

  jax:
    needs: rocm-config
    if: ${{ github.event_name != 'pull_request' || github.base_ref == 'main' }}
    name: JAX Linux x86 AMD Instinct GPU
    runs-on: ${{ inputs.runner_label || 'amd-do-linux-xla-gpu-gfx950-1' }}
    timeout-minutes: 80
    # Remove this once the JAX CI is stable.
    continue-on-error: true
    env:
      DOCKER_IMAGE: ${{ needs.rocm-config.outputs.docker-image-jax }}
    container:
      image: ${{ needs.rocm-config.outputs.docker-image-jax }} # zizmor: ignore[unpinned-images] image is pinned in the rocm-config job.
      volumes: &rocm_container_volumes
        - /data/ci-cert.crt:/data/ci-cert.crt
        - /data/ci-cert.key:/data/ci-cert.key
        - /data:/data
      options: &rocm_container_options >-
        --device=/dev/dri
        --device=/dev/kfd
        --ipc=host
        --shm-size=64G
        --cap-add=SYS_PTRACE
        --security-opt=seccomp=unconfined
        --tmpfs /root/.cache/bazel:rw,exec,size=80g
        --group-add video
        --env-file /etc/podinfo/gha-gpu-isolation-settings
    defaults:
      run:
        shell: bash
    steps:
      - name: Checking out openxla/xla
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          repository: openxla/xla
          persist-credentials: false

      - name: Checkout JAX
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          repository: jax-ml/jax
          path: jax
          persist-credentials: false

      - name: CPU Info
        run: lscpu

      - name: Run JAX Unit Tests
        timeout-minutes: 60
        working-directory: jax
        run: |
          ./ci/run_bazel_test_rocm_rbe.sh \
            --override_repository=xla="${GITHUB_WORKSPACE}" \
            --override_module=xla="${GITHUB_WORKSPACE}" \
            --config=single_gpu \
            --//jax:build_jaxlib=wheel \
            --//jax:build_jax=true \
            --local_test_jobs=1 \
            --action_env=JAX_ENABLE_X64=0 \
            --repo_env=HERMETIC_PYTHON_VERSION=3.14 \
            --repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" \
            --jobs=20 \
            --bes_keywords=jax \
            --bes_keywords=upstream \
            --bes_keywords=gpu \
            --bes_keywords=pull_request
  xla:
    needs: rocm-config
    name: XLA Linux x86 AMD Instinct GPU
    runs-on: ${{ inputs.runner_label || 'amd-do-linux-xla-gpu-gfx950-1' }}
    timeout-minutes: 220
    env:
      DOCKER_IMAGE: ${{ needs.rocm-config.outputs.docker-image }}
      EXECUTE_CI_BUILD_URL: https://raw.githubusercontent.com/ROCm/xla/refs/heads/${{ inputs.rocm_xla_branch || 'rocm-dev-infra' }}/build_tools/rocm/execute_ci_build_upstream.sh
    container:
      image: ${{ needs.rocm-config.outputs.docker-image }} # zizmor: ignore[unpinned-images] image is pinned in the rocm-config job.
      volumes: *rocm_container_volumes
      options: *rocm_container_options
    defaults:
      run:
        shell: bash
    steps:
      - name: Checking out openxla/xla
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          repository: openxla/xla
          persist-credentials: false

      - name: Fetch ROCm XLA Upstream Script from ROCm/xla
        run: |
          wget -O build_tools/rocm/execute_ci_build_upstream.sh "${EXECUTE_CI_BUILD_URL}"
          chmod +x build_tools/rocm/execute_ci_build_upstream.sh

      - name: CPU Info
        run: lscpu

      - name: Test XLA [single_gpu]
        timeout-minutes: 120
        run: |
          build_tools/rocm/execute_ci_build_upstream.sh \
            --config=rocm_ci \
            --config=rocm_rbe \
            --config=ci_single_gpu \
            --local_test_jobs=1 \
            --repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" \
            --repo_env=ROCM_PATH="/opt/rocm" \
            --local_test_jobs=1 \
            --internal_spawn_scheduler \
            --strategy=TestRunner=dynamic \
            --bes_keywords=xla \
            --bes_keywords=upstream \
            --bes_keywords=gpu \
            --bes_keywords=pull_request \
      - name: Test XLA [rocm_cpu]
        timeout-minutes: 80
        run: |
          build_tools/rocm/execute_ci_build_upstream.sh \
            --config=rocm_ci \
            --config=rocm_rbe \
            --config=ci_rocm_cpu \
            --local_test_jobs=200 \
            --repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" \
            --repo_env=ROCM_PATH="/opt/rocm" \
            --bes_keywords=xla \
            --bes_keywords=upstream \
            --bes_keywords=cpu \
            --bes_keywords=pull_request
