# Copyright 2025 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 Sharktank Tests
on:
  workflow_call:
    inputs:
      artifact_run_id:
        type: string
        default: ""
  workflow_dispatch:
    inputs:
      artifact_run_id:
        type: string
        default: ""

jobs:
  tests:
    name: "sharktank_tests :: ${{ matrix.name }}"
    runs-on: ${{ matrix.runs-on }}
    strategy:
      fail-fast: false
      matrix:
        include:
          # CPU
          - name: cpu_task
            target: target_cpu
            rocm-chip: none
            models-config-file: models_cpu_llvm_task.json
            backend: cpu
            iree_test_files: /home/nod/iree_tests_cache
            sku: cpu
            runs-on:
              - self-hosted # must come first
              - persistent-cache
              - Linux
              - X64

          # AMD GPU
          - name: amdgpu_rocm_mi308_gfx942
            rocm-chip: gfx942
            backend: rocm
            iree_test_files: /home/esaimana/iree_tests_cache
            sku: mi308
            target: target_hip
            runs-on: nodai-amdgpu-mi308-x86-64

    env:
      VENV_DIR: ${{ github.workspace }}/venv
      PACKAGE_DOWNLOAD_DIR: ${{ github.workspace }}/.packages
      TEST_OUTPUT_ARTIFACTS: ${{ github.workspace }}/model_output_artifacts
      IREE_TEST_FILES: ${{ matrix.iree_test_files }}
    steps:
      - name: Checking out IREE repository
        uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3
        with:
          submodules: false

      - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v5.4.0
        with:
          # Must match the subset of versions built in pkgci_build_packages.
          python-version: "3.11"

      - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
        if: ${{ inputs.artifact_run_id == '' }}
        with:
          name: linux_x86_64_release_packages
          path: ${{ env.PACKAGE_DOWNLOAD_DIR }}

      - name: Checkout test suites repository
        uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3
        with:
          repository: iree-org/iree-test-suites
          ref: 612f93663b32ea86ea442e87be549f1056a5de1c
          path: iree-test-suites
          lfs: false

      - name: Setup Python virtual environment
        run: |
          ./build_tools/pkgci/setup_venv.py ${VENV_DIR} \
            --artifact-path=${PACKAGE_DOWNLOAD_DIR} \
            --fetch-gh-workflow=${{ inputs.artifact_run_id }}
          source ${VENV_DIR}/bin/activate
          pip install -e ${{ github.workspace }}/iree-test-suites/sharktank_models
        env:
          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}

      - name: Run sharktank model tests
        run: |
          source ${VENV_DIR}/bin/activate
          pytest iree-test-suites/sharktank_models \
              -rA \
              --log-cli-level=info \
              --override-ini=xfail_strict=false \
              -m ${{ matrix.target }} \
              --timeout=120 \
              --durations=0 \
              --ignore=iree-test-suites/sharktank_models/benchmarks \
              --ignore=iree-test-suites/sharktank_models/quality_tests
        env:
          HIP_TARGET: ${{ matrix.rocm-chip }}
          ASSET_PATH: ${{ github.workspace }}/iree-test-suites/sharktank_models

      - name: Run quality tests
        run: |
          source ${VENV_DIR}/bin/activate
          pytest \
            iree-test-suites/sharktank_models/quality_tests \
            -rpFe \
            --log-cli-level=info \
            --durations=0 \
            --timeout=1200 \
            --capture=no \
            --test-file-directory=${{ github.workspace }}/tests/external/iree-test-suites/sharktank_models/quality_tests \
            --external-file-directory=${{ github.workspace }}/tests/external/iree-test-suites/test_suite_files
        env:
          ROCM_CHIP: ${{ matrix.rocm-chip }}
          SKU: ${{ matrix.sku }}
          IREE_TEST_FILES: ${{ env.IREE_TEST_FILES }}
          BACKEND: ${{ matrix.backend }}


      # Since the benchmark tests require compiled model files, this step must run after the quality tests
      - name: Run benchmark tests
        run: |
          source ${VENV_DIR}/bin/activate
          pytest \
            iree-test-suites/sharktank_models/benchmarks \
            --log-cli-level=info \
            --retries=7 \
            --timeout=600 \
            --test-file-directory=${{ github.workspace }}/tests/external/iree-test-suites/sharktank_models/benchmarks \
            --external-file-directory=${{ github.workspace }}/tests/external/iree-test-suites/test_suite_files

          echo "$(<${JOB_SUMMARY_PATH}/job_summary.md )" >> $GITHUB_STEP_SUMMARY
        env:
          ROCM_CHIP: ${{ matrix.rocm-chip }}
          SKU: ${{ matrix.sku }}
          IREE_TEST_FILES: ${{ env.IREE_TEST_FILES }}
          BACKEND: ${{ matrix.backend }}
          JOB_SUMMARY_PATH: ${{ github.workspace }}/iree-test-suites


  # These temporary tests are designed for targets who do not have cached
  # artifacts and unable to run quality and benchmark tests
  # TODO (geomin12): add cached artifact support for w7900
  model-test:
    name: "sharktank_model_tests :: ${{ matrix.name }}"
    runs-on: ${{ matrix.runs-on }}
    strategy:
      fail-fast: false
      matrix:
        include:
          - name: rocm_hip_w7900
            rocm-chip: gfx1100
            target: target_hip
            runs-on: [Linux, X64, iree-w7900x2, persistent-cache]
    env:
      VENV_DIR: ${{ github.workspace }}/venv
      PACKAGE_DOWNLOAD_DIR: ${{ github.workspace }}/.packages
      TEST_OUTPUT_ARTIFACTS: ${{ github.workspace }}/model_output_artifacts
    steps:
      - name: Checking out IREE repository
        uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3
        with:
          submodules: false

      - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v5.4.0
        with:
          # Must match the subset of versions built in pkgci_build_packages.
          python-version: "3.11"

      - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
        if: ${{ inputs.artifact_run_id == '' }}
        with:
          name: linux_x86_64_release_packages
          path: ${{ env.PACKAGE_DOWNLOAD_DIR }}

      - name: Checkout test suites repository
        uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3
        with:
          repository: iree-org/iree-test-suites
          ref: 612f93663b32ea86ea442e87be549f1056a5de1c
          path: iree-test-suites
          lfs: false

      - name: Setup Python virtual environment
        run: |
          ./build_tools/pkgci/setup_venv.py ${VENV_DIR} \
            --artifact-path=${PACKAGE_DOWNLOAD_DIR} \
            --fetch-gh-workflow=${{ inputs.artifact_run_id }}
          source ${VENV_DIR}/bin/activate
          pip install -e ${{ github.workspace }}/iree-test-suites/sharktank_models
        env:
          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}

      - name: Run sharktank model tests
        run: |
          source ${VENV_DIR}/bin/activate
          pytest iree-test-suites/sharktank_models \
              -rA \
              --log-cli-level=info \
              --override-ini=xfail_strict=false \
              -m ${{ matrix.target }} \
              --timeout=120 \
              --durations=0 \
              --ignore=iree-test-suites/sharktank_models/benchmarks \
              --ignore=iree-test-suites/sharktank_models/quality_tests
        env:
          HIP_TARGET: ${{ matrix.rocm-chip }}
          ASSET_PATH: ${{ github.workspace }}/iree-test-suites/sharktank_models