build_tools/benchmarks/run_benchmarks_on_android.py

#!/usr/bin/env python3
# Copyright 2021 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
"""Runs all matched benchmark suites on an Android device.

This script probes the Android phone via `adb` and uses the device information
to filter and run suitable benchmarks and optionally captures Tracy traces on
the Android phone.

It expects that `adb` is installed, and there is iree tools cross-compiled
towards Android. If to capture traces, another set of tracing-enabled iree
tools and the Tracy `capture` tool should be cross-compiled towards Android.

Example usages:

  # Without trace generation
  python3 run_benchmarks.py \
    --normal_benchmark_tool_dir=/path/to/normal/android/target/tools/dir \
    /path/to/host/build/dir

  # With trace generation
  python3 run_benchmarks.py \
    --normal_benchmark_tool_dir=/path/to/normal/android/target/tools/dir \
    --traced_benchmark_tool_dir=/path/to/tracy/android/target/tools/dir \
    --trace_capture_tool=/path/to/host/build/tracy/capture \
    /path/to/host/build/dir
"""

import sys
import pathlib

# Add build_tools python dir to the search path.
sys.path.insert(0, str(pathlib.Path(__file__).parent.with_name("python")))

import atexit
import json
import shutil
import subprocess
import tarfile
from typing import Any, List, Optional, Sequence, Tuple

from common.benchmark_config import BenchmarkConfig
from common.benchmark_driver import BenchmarkDriver
from common.benchmark_definition import (
    DriverInfo,
    execute_cmd,
    execute_cmd_and_get_stdout,
    execute_cmd_and_get_output,
    get_git_commit_hash,
    get_iree_benchmark_module_arguments,
    wait_for_iree_benchmark_module_start,
    parse_iree_benchmark_metrics,
)
from common.benchmark_suite import MODEL_FLAGFILE_NAME, BenchmarkCase, BenchmarkSuite
from common.android_device_utils import (
    get_android_device_model,
    get_android_device_info,
    get_android_gpu_name,
)
import common.common_arguments
from e2e_test_artifacts import iree_artifacts
from e2e_test_framework import serialization
from e2e_test_framework.definitions import common_definitions, iree_definitions
from e2e_test_framework.device_specs import device_parameters

# Root directory to perform benchmarks in on the Android device.
ANDROID_TMPDIR = pathlib.PurePosixPath("/data/local/tmp/iree-benchmarks")

NORMAL_TOOL_REL_DIR = pathlib.PurePosixPath("normal-tools")
TRACED_TOOL_REL_DIR = pathlib.PurePosixPath("traced-tools")


def adb_push_to_tmp_dir(
    content: pathlib.Path,
    relative_dir: pathlib.PurePosixPath = pathlib.PurePosixPath(),
    verbose: bool = False,
) -> pathlib.PurePosixPath:
    """Pushes content onto the Android device.

    Args:
      content: the full path to the source file.
      relative_dir: the directory to push to; relative to ANDROID_TMPDIR.

    Returns:
      The full path to the content on the Android device.
    """
    filename = content.name
    android_path = ANDROID_TMPDIR / relative_dir / filename
    # When the output is a TTY, keep the default progress info output.
    # In other cases, redirect progress info to null to avoid bloating log files.
    stdout_redirect = None if sys.stdout.isatty() else subprocess.DEVNULL
    execute_cmd(
        ["adb", "push", content.resolve(), android_path],
        verbose=verbose,
        stdout=stdout_redirect,
    )
    return android_path


def adb_execute_and_get_output(
    cmd_args: Sequence[str],
    relative_dir: pathlib.PurePosixPath = pathlib.PurePosixPath(),
    verbose: bool = False,
) -> Tuple[str, str]:
    """Executes command with adb shell.

    Switches to `relative_dir` relative to the android tmp directory before
    executing. Waits for completion and returns the command stdout.

    Args:
      cmd_args: a list containing the command to execute and its parameters
      relative_dir: the directory to execute the command in; relative to
        ANDROID_TMPDIR.

    Returns:
      Strings for stdout and stderr.
    """
    cmd = ["adb", "shell", "cd", ANDROID_TMPDIR / relative_dir, "&&"]
    cmd.extend(cmd_args)
    return execute_cmd_and_get_output(cmd, verbose=verbose)


def adb_execute(
    cmd_args: Sequence[str],
    relative_dir: pathlib.PurePosixPath = pathlib.PurePosixPath(),
    verbose: bool = False,
) -> subprocess.CompletedProcess:
    """Executes command with adb shell.

    Switches to `relative_dir` relative to the android tmp directory before
    executing. Waits for completion. Output is streamed to the terminal.

    Args:
      cmd_args: a list containing the command to execute and its parameters
      relative_dir: the directory to execute the command in; relative to
        ANDROID_TMPDIR.

    Returns:
      The completed process.
    """
    cmd = ["adb", "shell", "cd", ANDROID_TMPDIR / relative_dir, "&&"]
    cmd.extend(cmd_args)
    return execute_cmd(cmd, verbose=verbose)


def is_magisk_su():
    """Returns true if the Android device has a Magisk SU binary."""
    stdout, _ = adb_execute_and_get_output(["su", "--help"])
    return "MagiskSU" in stdout


def adb_execute_as_root(cmd_args: Sequence[Any]) -> subprocess.CompletedProcess:
    """Executes the given command as root."""
    cmd = ["su", "-c" if is_magisk_su() else "root"]
    cmd.extend(cmd_args)
    return adb_execute(cmd)


def adb_start_cmd(
    cmd_args: Sequence[str],
    relative_dir: pathlib.PurePosixPath = pathlib.PurePosixPath(),
    verbose: bool = False,
) -> subprocess.Popen:
    """Executes command with adb shell in a directory and returns the handle
    without waiting for completion.

    Args:
      cmd_args: a list containing the command to execute and its parameters
      relative_dir: the directory to execute the command in; relative to
        ANDROID_TMPDIR.

    Returns:
      A Popen object for the started command.
    """
    cmd = ["adb", "shell", "cd", ANDROID_TMPDIR / relative_dir, "&&"]
    cmd.extend(cmd_args)

    if verbose:
        print(f"cmd: {cmd}")
    return subprocess.Popen(cmd, stdout=subprocess.PIPE, text=True)


def get_vmfb_full_path_for_benchmark_case(
    benchmark_case_dir: pathlib.Path,
) -> pathlib.Path:
    flagfile = benchmark_case_dir / MODEL_FLAGFILE_NAME
    for line in flagfile.read_text().splitlines():
        flag_name, flag_value = line.strip().split("=")
        if flag_name == "--module":
            # Realpath canonicalization matters. The caller may rely on that to track
            # which files it already pushed.
            return (benchmark_case_dir / flag_value).resolve()
    raise ValueError(f"{flagfile} does not contain a --module flag")


class AndroidBenchmarkDriver(BenchmarkDriver):
    """Android benchmark driver."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.already_pushed_files = {}

    def run_benchmark_case(
        self,
        benchmark_case: BenchmarkCase,
        benchmark_results_filename: Optional[pathlib.Path],
        capture_filename: Optional[pathlib.Path],
    ) -> None:
        benchmark_case_dir = benchmark_case.benchmark_case_dir
        android_case_dir = pathlib.PurePosixPath(
            benchmark_case_dir.relative_to(self.config.root_benchmark_dir)
        )

        run_config = benchmark_case.run_config
        self.__check_and_push_file(
            benchmark_case_dir / iree_artifacts.MODULE_FILENAME, android_case_dir
        )
        taskset = self.__deduce_taskset_from_run_config(run_config)
        run_args = run_config.materialize_run_flags()
        run_args.append(f"--module={iree_artifacts.MODULE_FILENAME}")

        if benchmark_results_filename is not None:
            self.__run_benchmark(
                android_case_dir=android_case_dir,
                tool_name=benchmark_case.benchmark_tool_name,
                driver_info=benchmark_case.driver_info,
                run_args=run_args,
                results_filename=benchmark_results_filename,
                taskset=taskset,
            )

        if capture_filename is not None:
            self.__run_capture(
                android_case_dir=android_case_dir,
                tool_name=benchmark_case.benchmark_tool_name,
                run_args=run_args,
                capture_filename=capture_filename,
                taskset=taskset,
            )

    def __run_benchmark(
        self,
        android_case_dir: pathlib.PurePosixPath,
        tool_name: str,
        driver_info: DriverInfo,
        run_args: Sequence[str],
        results_filename: pathlib.Path,
        taskset: str,
    ):
        if self.config.normal_benchmark_tool_dir is None:
            raise ValueError("normal_benchmark_tool_dir can't be None.")

        host_tool_path = self.config.normal_benchmark_tool_dir / tool_name
        android_tool = self.__check_and_push_file(host_tool_path, NORMAL_TOOL_REL_DIR)
        cmd = ["taskset", taskset, android_tool]
        cmd += run_args
        if tool_name == "iree-benchmark-module":
            cmd += get_iree_benchmark_module_arguments(
                results_filename=f"'{results_filename.name}'",
                driver_info=driver_info,
                benchmark_min_time=self.config.benchmark_min_time,
            )

        benchmark_stdout, benchmark_stderr = adb_execute_and_get_output(
            cmd, android_case_dir, verbose=self.verbose
        )
        benchmark_metrics = parse_iree_benchmark_metrics(
            benchmark_stdout, benchmark_stderr
        )
        if self.verbose:
            print(benchmark_metrics)
        results_filename.write_text(json.dumps(benchmark_metrics.to_json_object()))

    def __run_capture(
        self,
        android_case_dir: pathlib.PurePosixPath,
        tool_name: str,
        capture_filename: pathlib.Path,
        run_args: Sequence[str],
        taskset: str,
    ):
        capture_config = self.config.trace_capture_config
        if capture_config is None:
            raise ValueError("capture_config can't be None.")

        host_tool_path = capture_config.traced_benchmark_tool_dir / tool_name
        android_tool = self.__check_and_push_file(host_tool_path, TRACED_TOOL_REL_DIR)
        run_cmd = [
            "TRACY_NO_EXIT=1",
            f"IREE_PRESERVE_DYLIB_TEMP_FILES={ANDROID_TMPDIR}",
            "taskset",
            taskset,
            android_tool,
        ]
        run_cmd += run_args

        # Just launch the traced benchmark tool with TRACY_NO_EXIT=1 without
        # waiting for the adb command to complete as that won't happen.
        process = adb_start_cmd(run_cmd, android_case_dir, verbose=self.verbose)

        wait_for_iree_benchmark_module_start(process, self.verbose)

        # Now it's okay to collect the trace via the capture tool. This will
        # send the signal to let the previously waiting benchmark tool to
        # complete.
        capture_cmd = [capture_config.trace_capture_tool, "-f", "-o", capture_filename]
        # If verbose, just let the subprocess print its output. The subprocess
        # may need to detect if the output is a TTY to decide whether to log
        # verbose progress info and use ANSI colors, so it's better to use
        # stdout redirection than to capture the output in a string.
        stdout_redirect = None if self.verbose else subprocess.DEVNULL
        execute_cmd(capture_cmd, verbose=self.verbose, stdout=stdout_redirect)

    # TODO(#13187): These logics are inherited from the legacy benchmark suites,
    # which only work for a few specific phones. We should define the topology
    # in their device specs.
    def __deduce_taskset_from_run_config(
        self, run_config: iree_definitions.E2EModelRunConfig
    ) -> str:
        """Deduces the CPU mask according to device and execution config."""

        device_spec = run_config.target_device_spec
        # For GPU benchmarks, use the most performant core.
        if device_spec.architecture.type == common_definitions.ArchitectureType.GPU:
            return "80"

        device_params = device_spec.device_parameters
        single_thread = "1-thread" in run_config.module_execution_config.tags
        if device_parameters.ARM_BIG_CORES in device_params:
            return "80" if single_thread else "f0"
        elif device_parameters.ARM_LITTLE_CORES in device_params:
            return "08" if single_thread else "0f"

        raise ValueError(f"Unsupported config to deduce taskset: '{run_config}'.")

    def __check_and_push_file(
        self, host_path: pathlib.Path, relative_dir: pathlib.PurePosixPath
    ):
        """Checks if the file has been pushed and pushes it if not."""
        android_path = self.already_pushed_files.get(host_path)
        if android_path is not None:
            return android_path

        android_path = adb_push_to_tmp_dir(
            host_path, relative_dir=relative_dir, verbose=self.verbose
        )
        self.already_pushed_files[host_path] = android_path
        return android_path


def set_cpu_frequency_scaling_governor(governor: str):
    git_root = execute_cmd_and_get_stdout(["git", "rev-parse", "--show-toplevel"])
    cpu_script = (
        pathlib.Path(git_root)
        / "build_tools"
        / "benchmarks"
        / "set_android_scaling_governor.sh"
    )
    android_path = adb_push_to_tmp_dir(cpu_script)
    adb_execute_as_root([android_path, governor])


def set_gpu_frequency_scaling_policy(policy: str):
    git_root = execute_cmd_and_get_stdout(["git", "rev-parse", "--show-toplevel"])
    device_model = get_android_device_model()
    gpu_name = get_android_gpu_name()
    benchmarks_tool_dir = pathlib.Path(git_root) / "build_tools" / "benchmarks"
    if device_model == "Pixel-6" or device_model == "Pixel-6-Pro":
        gpu_script = benchmarks_tool_dir / "set_pixel6_gpu_scaling_policy.sh"
    elif gpu_name.lower().startswith("adreno"):
        gpu_script = benchmarks_tool_dir / "set_adreno_gpu_scaling_policy.sh"
    else:
        raise RuntimeError(
            f"Unsupported device '{device_model}' for setting GPU scaling policy"
        )
    android_path = adb_push_to_tmp_dir(gpu_script)
    adb_execute_as_root([android_path, policy])


def main(args):
    device_info = get_android_device_info(args.verbose)
    if args.verbose:
        print(device_info)

    commit = get_git_commit_hash("HEAD")
    benchmark_config = BenchmarkConfig.build_from_args(args, commit)
    benchmark_groups = json.loads(args.execution_benchmark_config.read_text())
    benchmark_group = benchmark_groups.get(args.target_device_name)
    if benchmark_group is None:
        raise ValueError("Target device not found in the benchmark config.")
    run_configs = serialization.unpack_and_deserialize(
        data=benchmark_group["run_configs"],
        root_type=List[iree_definitions.E2EModelRunConfig],
    )
    benchmark_suite = BenchmarkSuite.load_from_run_configs(
        run_configs=run_configs, root_benchmark_dir=benchmark_config.root_benchmark_dir
    )

    benchmark_driver = AndroidBenchmarkDriver(
        device_info=device_info,
        benchmark_config=benchmark_config,
        benchmark_suite=benchmark_suite,
        benchmark_grace_time=1.0,
        verbose=args.verbose,
    )

    if args.pin_cpu_freq:
        set_cpu_frequency_scaling_governor("performance")
        atexit.register(set_cpu_frequency_scaling_governor, "schedutil")
    if args.pin_gpu_freq:
        set_gpu_frequency_scaling_policy("performance")
        atexit.register(set_gpu_frequency_scaling_policy, "default")

    # Clear the benchmark directory on the Android device first just in case
    # there are leftovers from manual or failed runs.
    execute_cmd_and_get_stdout(
        ["adb", "shell", "rm", "-rf", ANDROID_TMPDIR], verbose=args.verbose
    )

    if not args.no_clean:
        # Clear the benchmark directory on the Android device.
        atexit.register(
            execute_cmd_and_get_stdout,
            ["adb", "shell", "rm", "-rf", ANDROID_TMPDIR],
            verbose=args.verbose,
        )
        # Also clear temporary directory on the host device.
        atexit.register(shutil.rmtree, args.tmp_dir)

    # Tracy client and server communicate over port 8086 by default. If we want
    # to capture traces along the way, forward port via adb.
    trace_capture_config = benchmark_config.trace_capture_config
    if trace_capture_config:
        execute_cmd_and_get_stdout(
            ["adb", "forward", "tcp:8086", "tcp:8086"], verbose=args.verbose
        )
        atexit.register(
            execute_cmd_and_get_stdout,
            ["adb", "forward", "--remove", "tcp:8086"],
            verbose=args.verbose,
        )

    benchmark_driver.run()

    benchmark_results = benchmark_driver.get_benchmark_results()
    if args.output is not None:
        with open(args.output, "w") as f:
            f.write(benchmark_results.to_json_str())

    if args.verbose:
        print(benchmark_results.commit)
        print(benchmark_results.benchmarks)

    if trace_capture_config:
        # Put all captures in a tarball and remove the original files.
        with tarfile.open(trace_capture_config.capture_tarball, "w:gz") as tar:
            for capture_filename in benchmark_driver.get_capture_filenames():
                tar.add(capture_filename)

    benchmark_errors = benchmark_driver.get_benchmark_errors()
    if benchmark_errors:
        print("Benchmarking completed with errors", file=sys.stderr)
        raise RuntimeError(benchmark_errors)


if __name__ == "__main__":
    main(common.common_arguments.Parser().parse_args())