# Copyright 2022 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
#
# Workflow for all benchmark-related jobs. It is designed to be called from
# the main workflow ci.yml. The concurrency of this workflow is controlled by
# the caller's job.

name: Benchmarks

on:
  workflow_call:
    inputs:
      runner-group:
        required: true
        type: string
      runner-env:
        required: true
        type: string
      build-dir:
        required: true
        type: string
      build-dir-archive:
        required: true
        type: string
      build-dir-gcs-artifact:
        required: true
        type: string
      tf-binaries-dir:
        required: true
        type: string
      tf-binaries-archive:
        required: true
        type: string
      tf-binaries-gcs-artifact:
        required: true
        type: string
    outputs:
      benchmarks-dir:
        description: "Benchmarks directory"
        value: ${{ jobs.build_suites.outputs.benchmarks-dir }}
      benchmarks-archive:
        description: "Benchmarks archive tarball"
        value: ${{ jobs.build_suites.outputs.benchmarks-archive }}
      benchmarks-gcs-artifact:
        description: "Benchmarks GCS artifact"
        value: ${{ jobs.build_suites.outputs.benchmarks-gcs-artifact }}

env:
  # This duplicates the variable from ci.yml. The variable needs to be in env
  # instead of the outputs of setup because it contains the run attempt and we
  # want that to be the current attempt, not whatever attempt the setup step
  # last ran in. It therefore can't be passed in via inputs because the env
  # context isn't available there.
  GCS_DIR: gs://iree-github-actions-${{ github.event_name == 'pull_request' && 'presubmit' || 'postsubmit' }}-artifacts/${{ github.run_id }}/${{ github.run_attempt }}

jobs:
  build_suites:
    runs-on:
      - self-hosted # must come first
      - runner-group=${{ inputs.runner-group }}
      - environment=${{ inputs.runner-env }}
      - cpu
      - os-family=Linux
    env:
      BUILD_DIR: ${{ inputs.build-dir }}
      BUILD_DIR_ARCHIVE: ${{ inputs.build-dir-archive }}
      BUILD_DIR_GCS_ARTIFACT: ${{ inputs.build-dir-gcs-artifact }}
      TF_BINARIES_DIR: ${{ inputs.tf-binaries-dir }}
      TF_BINARIES_ARCHIVE: ${{ inputs.tf-binaries-archive }}
      TF_BINARIES_GCS_ARTIFACT: ${{ inputs.tf-binaries-gcs-artifact }}
    outputs:
      benchmarks-dir: ${{ steps.build.outputs.benchmarks-dir }}
      benchmarks-archive: ${{ steps.archive.outputs.benchmarks-archive }}
      benchmarks-gcs-artifact: ${{ steps.upload.outputs.benchmarks-gcs-artifact }}
    steps:
      - name: "Checking out repository"
        uses: actions/checkout@7884fcad6b5d53d10323aee724dc68d8b9096a2e # v2
      - name: "Checking out runtime submodules"
        run: ./build_tools/scripts/git/update_runtime_submodules.sh
      - name: "Downloading build dir archive"
        run: gcloud alpha storage cp "${BUILD_DIR_GCS_ARTIFACT}" "${BUILD_DIR_ARCHIVE}"
      - name: "Extracting install from build dir archive"
        run: tar -xf "${BUILD_DIR_ARCHIVE}" "${BUILD_DIR}/install"
      - name: "Downloading TF binaries archive"
        run: gcloud alpha storage cp "${TF_BINARIES_GCS_ARTIFACT}" "${TF_BINARIES_ARCHIVE}"
      - name: "Extracting TF binaries archive"
        run: tar -xf "${TF_BINARIES_ARCHIVE}"
      - name: "Building benchmarks"
        id: build
        env:
          BENCHMARKS_OUTPUT_DIR: build-benchmarks
        run: |
          build_tools/github_actions/docker_run.sh \
            --env "IREE_TF_BINARIES_DIR=${TF_BINARIES_DIR}" \
            --env "IREE_HOST_BINARY_ROOT=${BUILD_DIR}/install" \
            --env "BUILD_BENCHMARKS_DIR=${BENCHMARKS_OUTPUT_DIR}" \
            gcr.io/iree-oss/base@sha256:22c43975179265296e016d15eb6f65d18abd5f9d4d3a5fa5e478ca4862bb61c4 \
            build_tools/cmake/build_benchmarks.sh
          echo "::set-output name=benchmarks-dir::${BENCHMARKS_OUTPUT_DIR}"
      - name: "Creating archive of benchmarks"
        id: archive
        env:
          BENCHMARKS_ARCHIVE: benchmarks.tar
          BENCHMARKS_DIR: ${{ steps.build.outputs.benchmarks-dir }}
        run: |
          # TODO(#10988): Excludes the experimental `models` and iree` dirs for
          # now.
          tar --exclude '*.tflite' --exclude '*.mlir' --exclude '*-tf-model' \
            --exclude "${BENCHMARKS_DIR}/benchmark_suites/models" \
            --exclude "${BENCHMARKS_DIR}/benchmark_suites/iree" \
            -czf \
            "${BENCHMARKS_ARCHIVE}" "${BENCHMARKS_DIR}/benchmark_suites"
          echo "::set-output name=benchmarks-archive::${BENCHMARKS_ARCHIVE}"
      - name: "Uploading benchmarks archive"
        id: upload
        env:
          BENCHMARKS_ARCHIVE: ${{ steps.archive.outputs.benchmarks-archive }}
          BENCHMARKS_GCS_ARTIFACT: ${{ env.GCS_DIR }}/${{ steps.archive.outputs.benchmarks-archive }}
        run: |
          gcloud alpha storage cp "${BENCHMARKS_ARCHIVE}" "${BENCHMARKS_GCS_ARTIFACT}"
          echo "::set-output name=benchmarks-gcs-artifact::${BENCHMARKS_GCS_ARTIFACT}"

  build_tools:
    runs-on:
      - self-hosted  # must come first
      - runner-group=${{ inputs.runner-group }}
      - environment=${{ inputs.runner-env }}
      - cpu
      - os-family=Linux
    strategy:
      matrix:
        target:
          - platform: "riscv"
            architecture: "rv64"
            docker_image: "gcr.io/iree-oss/riscv@sha256:e11710158fd09a04ae7accc275a38c95cd59f83fd15a0630a083eb5623807451"
    outputs:
      benchmark-tools-gcs-artifacts: ${{ toJSON(steps.upload.outputs) }}
    env:
      PLATFORM: ${{ matrix.target.platform }}
      ARCHITECTURE: ${{ matrix.target.architecture }}
      DOCKER_IMAGE: ${{ matrix.target.docker_image }}
      BUILD_TOOLS_DIR: ${{ matrix.target.platform }}-${{ matrix.target.architecture }}-benchmark-tools-dir
      BUILD_DIR: ${{ inputs.build-dir }}
      BUILD_DIR_ARCHIVE: ${{ inputs.build-dir-archive }}
      BUILD_DIR_GCS_ARTIFACT: ${{ inputs.build-dir-gcs-artifact }}
    steps:
      - name: "Checking out repository"
        uses: actions/checkout@7884fcad6b5d53d10323aee724dc68d8b9096a2e # v2
      - name: "Checking out runtime submodules"
        run: ./build_tools/scripts/git/update_runtime_submodules.sh
      - name: "Downloading build dir archive"
        run: gcloud alpha storage cp "${BUILD_DIR_GCS_ARTIFACT}" "${BUILD_DIR_ARCHIVE}"
      - name: "Extracting host binaries"
        run: tar -xf "${BUILD_DIR_ARCHIVE}" "${BUILD_DIR}/install"
      - name: "Compiling the benchmark tools"
        id: build
        run: |
          ./build_tools/github_actions/docker_run.sh \
            --env "${PLATFORM^^}_ARCH=${ARCHITECTURE}" \
            --env "BUILD_${PLATFORM^^}_DIR=${BUILD_TOOLS_DIR}/build" \
            --env "BUILD_PRESET=benchmark" \
            --env "IREE_HOST_BINARY_ROOT=${BUILD_DIR}/install" \
            "${DOCKER_IMAGE}" \
            bash -euo pipefail -c "./build_tools/cmake/build_${PLATFORM}.sh"
      - name: "Compiling the benchmark tools with tracing"
        id: build-with-tracing
        run: |
          ./build_tools/github_actions/docker_run.sh \
            --env "${PLATFORM^^}_ARCH=${ARCHITECTURE}" \
            --env "BUILD_${PLATFORM^^}_DIR=${BUILD_TOOLS_DIR}/build-traced" \
            --env "BUILD_PRESET=benchmark-with-tracing" \
            --env "IREE_HOST_BINARY_ROOT=${BUILD_DIR}/install" \
            "${DOCKER_IMAGE}" \
            bash -euo pipefail -c "./build_tools/cmake/build_${PLATFORM}.sh"
      - name: "Creating the benchmark tools archive"
        id: archive
        env:
          BENCHMARK_TOOLS_ARCHIVE: ${{ matrix.target.platform }}-${{ matrix.target.architecture }}-benchmarks-tools.tar
        run: |
          tar -cf ${BENCHMARK_TOOLS_ARCHIVE} \
            ${BUILD_TOOLS_DIR}/*/tools/iree-benchmark-module
          echo "::set-output name=benchmark-tools-archive::${BENCHMARK_TOOLS_ARCHIVE}"
      - name: "Uploading the benchmark tools archive"
        id: upload
        env:
          BENCHMARK_TOOLS_ARCHIVE: ${{ steps.archive.outputs.benchmark-tools-archive }}
          BENCHMARK_TOOLS_GCS_ARTIFACT: ${{ env.GCS_DIR }}/${{ steps.archive.outputs.benchmark-tools-archive }}
        run: |
          gcloud alpha storage cp "${BENCHMARK_TOOLS_ARCHIVE}" "${BENCHMARK_TOOLS_GCS_ARTIFACT}"
          echo "::set-output name=${PLATFORM}-${ARCHITECTURE}-benchmark-tools-gcs-artifact::${BENCHMARK_TOOLS_GCS_ARTIFACT}"