build_tools/benchmarks/common/linux_device_utils_test.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

import unittest

from unittest import mock

from common.benchmark_definition import DeviceInfo, PlatformType
from common.linux_device_utils import canonicalize_gpu_name, get_linux_cpu_arch, get_linux_cpu_features

LSCPU_OUTPUT = ("Architecture:                    x86_64\n"
                "Vendor ID:                       AuthenticAMD\n"
                "Flags:                           fpu vme de pse tsc\n")


class LinuxDeviceUtilsTest(unittest.TestCase):

  def test_get_linux_cpu_arch(self):
    self.assertEqual(get_linux_cpu_arch(LSCPU_OUTPUT), "x86_64")

  def test_get_linux_cpu_features(self):
    self.assertEqual(get_linux_cpu_features(LSCPU_OUTPUT),
                     ["fpu", "vme", "de", "pse", "tsc"])

  def test_canonicalize_gpu_name(self):
    self.assertEqual(canonicalize_gpu_name("Tesla  V100-SXM2-16GB"),
                     "Tesla-V100-SXM2-16GB")


if __name__ == "__main__":
  unittest.main()