# Copyright 2024 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

name: PkgCI Test ONNX
on:
  workflow_call:
    inputs:
      artifact_run_id:
        type: string
        default: ""
  workflow_dispatch:
    inputs:
      artifact_run_id:
        type: string
        default: ""

jobs:
  test_onnx_ops:
    name: "test_onnx_ops :: ${{ matrix.name }}"
    runs-on: ${{ matrix.runs-on }}
    strategy:
      fail-fast: false
      matrix:
        include:
          # CPU
          - name: cpu_llvm_sync_O0
            config-file: onnx_ops_cpu_llvm_sync_O0.json
            numprocesses: auto
            runs-on: ubuntu-24.04
          - name: cpu_llvm_sync_O2
            config-file: onnx_ops_cpu_llvm_sync_O2.json
            numprocesses: auto
            runs-on: ubuntu-24.04

          # AMD GPU
          - name: amdgpu_hip_rdna3_O3
            numprocesses: 1
            config-file: onnx_ops_gpu_hip_rdna3_O3.json
            runs-on: [Linux, X64, gfx1100]
          - name: amdgpu_hip_rdna4_O3
            numprocesses: 1
            config-file: onnx_ops_gpu_hip_rdna4_O3.json
            runs-on: [Linux, X64, gfx1201]
          - name: amdgpu_vulkan_rdna3_O0
            numprocesses: 1
            config-file: onnx_ops_gpu_vulkan_rdna3_O0.json
            runs-on: [Linux, X64, rdna3]

          # NVIDIA GPU
          # TODO(#18238): migrate to new runner cluster
          # - name: nvidiagpu_cuda_O0
          #   config-file: onnx_ops_gpu_cuda_O0.json
          #   numprocesses: 4
          #   runs-on:
          #     - self-hosted # must come first
          #     - environment=prod
          #     - gpu # TODO(scotttodd): qualify further with vendor/model
          #     - os-family=Linux
          # - name: nvidiagpu_vulkan_O0
          #   config-file: onnx_ops_gpu_vulkan_O0.json
          #   numprocesses: 4
          #   runs-on:
          #     - self-hosted # must come first
          #     - environment=prod
          #     - gpu # TODO(scotttodd): qualify further with vendor/model
          #     - os-family=Linux
    env:
      PACKAGE_DOWNLOAD_DIR: ${{ github.workspace }}/.packages
      CONFIG_FILE_PATH: tests/external/iree-test-suites/onnx_ops/${{ matrix.config-file }}
      NUMPROCESSES: ${{ matrix.numprocesses }}
      LOG_FILE_PATH: /tmp/test_onnx_ops_${{ matrix.name }}_logs.json
      VENV_DIR: ${{ github.workspace }}/venv
      GH_TOKEN: ${{ github.token }}
    steps:
      - name: Checking out IREE repository
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          submodules: false
      - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
        with:
          # Must match the subset of versions built in pkgci_build_packages.
          python-version: "3.11"
      - uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2
        if: ${{ inputs.artifact_run_id == '' }}
        with:
          name: linux_x86_64_release_packages
          path: ${{ env.PACKAGE_DOWNLOAD_DIR }}
      - name: Setup venv
        run: |
          ./build_tools/pkgci/setup_venv.py ${VENV_DIR} \
            --artifact-path=${PACKAGE_DOWNLOAD_DIR} \
            --fetch-gh-workflow=${{ inputs.artifact_run_id }}
        env:
          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}

      - name: Checkout test suites repository
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          repository: iree-org/iree-test-suites
          ref: 612f93663b32ea86ea442e87be549f1056a5de1c
          path: iree-test-suites
      - name: Install ONNX ops test suite requirements
        run: |
          source ${VENV_DIR}/bin/activate
          python -m pip install -r iree-test-suites/onnx_ops/requirements.txt
      - name: Run ONNX ops test suite
        env:
          # Enable hip logging on errors.
          AMD_LOG_LEVEL: 1
        run: |
          source ${VENV_DIR}/bin/activate
          pytest iree-test-suites/onnx_ops/ \
              -rpfE \
              --numprocesses ${NUMPROCESSES} \
              --timeout=30 \
              --durations=20 \
              --config-files=${CONFIG_FILE_PATH} \
              --report-log=${LOG_FILE_PATH}
      - name: "Updating config file with latest XFAIL lists"
        if: failure()
        run: |
          source ${VENV_DIR}/bin/activate
          python iree-test-suites/onnx_ops/update_config_xfails.py \
            --log-file=${LOG_FILE_PATH} \
            --config-file=${CONFIG_FILE_PATH}
          cat ${CONFIG_FILE_PATH}
      - name: "Uploading new config file"
        if: failure()
        uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2
        with:
          name: ${{ matrix.config-file }}
          path: ${{ env.CONFIG_FILE_PATH }}

  test_onnx_models:
    name: "test_onnx_models :: ${{ matrix.name }}"
    runs-on: ${{ matrix.runs-on }}
    strategy:
      fail-fast: false
      matrix:
        include:
          # CPU
          - name: cpu_llvm_task
            config-file: onnx_models_cpu_llvm_task.json
            runs-on:
              - self-hosted # must come first
              - persistent-cache
              - Linux
              - X64

          # AMD GPU
          - name: amdgpu_hip_rdna3
            config-file: onnx_models_gpu_hip_rdna3.json
            runs-on: [Linux, X64, gfx1100, persistent-cache]
          - name: amdgpu_hip_rdna4
            config-file: onnx_models_gpu_hip_rdna4.json
            runs-on: [Linux, X64, gfx1201, persistent-cache]
          # TODO(#24464): There are vulkan driver issues on shark55.
          #- name: amdgpu_vulkan_rdna3
          #  config-file: onnx_models_gpu_vulkan.json
          #  runs-on: [Linux, X64, rdna3, persistent-cache]

          # NVIDIA GPU
          # TODO(#18238): migrate to new runner cluster
    env:
      VENV_DIR: ${{ github.workspace }}/venv
      CONFIG_FILE_PATH: tests/external/iree-test-suites/onnx_models/${{ matrix.config-file }}
    steps:
      - name: Checking out IREE repository
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          submodules: false
      - uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
        with:
          # Must match the subset of versions built in pkgci_build_packages.
          python-version: "3.11"
      - uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2
        with:
          name: linux_x86_64_release_packages
          path: ${{ env.PACKAGE_DOWNLOAD_DIR }}
      - name: Setup venv
        run: |
          ./build_tools/pkgci/setup_venv.py ${VENV_DIR} \
            --artifact-path=${PACKAGE_DOWNLOAD_DIR} \
            --fetch-gh-workflow=${{ inputs.artifact_run_id }}
        env:
          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}

      - name: Checkout test suites repository
        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          repository: iree-org/iree-test-suites
          ref: 612f93663b32ea86ea442e87be549f1056a5de1c
          path: iree-test-suites
      - name: Install ONNX models test suite requirements
        run: |
          source ${VENV_DIR}/bin/activate
          python -m pip install -r iree-test-suites/onnx_models/requirements.txt
      - name: Run ONNX models test suite
        env:
          # Enable hip logging on errors.
          AMD_LOG_LEVEL: 1
        run: |
          source ${VENV_DIR}/bin/activate
          pytest iree-test-suites/onnx_models/ \
              -rA \
              --log-cli-level=info \
              --timeout=120 \
              --durations=0 \
              --test-config-file=${CONFIG_FILE_PATH}