build_tools/benchmarks/common/linux_device_utils.py

#!/usr/bin/env python3
# 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
"""Utils for accessing Linux device information."""

import re
from typing import Optional, Sequence

from .benchmark_definition import (execute_cmd_and_get_stdout, DeviceInfo,
                                   PlatformType)


def _get_lscpu_field(lscpu_output: str, field_name: str) -> str:
  (value,) = re.findall(f"^{field_name}:\s*(.+)", lscpu_output, re.MULTILINE)
  return value


def get_linux_cpu_arch(lscpu_output: str) -> str:
  """Returns CPU Architecture, e.g., 'x86_64'."""
  return _get_lscpu_field(lscpu_output, "Architecture")


def get_linux_cpu_features(lscpu_output: str) -> Sequence[str]:
  """Returns CPU feature lists, e.g., ['mmx', 'fxsr', 'sse', 'sse2']."""
  return _get_lscpu_field(lscpu_output, "Flags").split(" ")


def canonicalize_gpu_name(gpu_name: str) -> str:
  # Replace all consecutive non-word characters with a single hyphen.
  return re.sub(r"\W+", "-", gpu_name)


def get_linux_device_info(device_model: str = "Unknown",
                          cpu_uarch: Optional[str] = None,
                          gpu_id: str = "0",
                          verbose: bool = False) -> DeviceInfo:
  """Returns device info for the Linux device.

    Args:
    - device_model: the device model name, e.g., 'ThinkStation P520'
    - cpu_uarch: the CPU microarchitecture, e.g., 'CascadeLake'
    - gpu_id: the target GPU ID, e.g., '0' or 'GPU-<UUID>'
  """
  lscpu_output = execute_cmd_and_get_stdout(["lscpu"], verbose)

  try:
    gpu_name = execute_cmd_and_get_stdout([
        "nvidia-smi", "--query-gpu=name", "--format=csv,noheader",
        f"--id={gpu_id}"
    ], verbose)
  except FileNotFoundError:
    # Set GPU name to Unknown if the tool "nvidia-smi" doesn't exist.
    gpu_name = "Unknown"

  return DeviceInfo(
      PlatformType.LINUX,
      # Includes CPU model as it is the key factor of the device performance.
      model=device_model,
      # Currently we only have x86, so CPU ABI = CPU arch.
      cpu_abi=get_linux_cpu_arch(lscpu_output),
      cpu_uarch=cpu_uarch,
      cpu_features=get_linux_cpu_features(lscpu_output),
      gpu_name=canonicalize_gpu_name(gpu_name))