TheRock/rocm-libraries/shared/rocroller/scripts/lib/rrperf/utils.py

# Copyright Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT

import datetime
import functools
import json
import os
import subprocess
from dataclasses import asdict, fields
from hashlib import sha1
from pathlib import Path

import rrperf


def empty():
    yield from ()


def sjoin(xs):
    return " ".join([str(x) for x in xs])


def load_suite(suite_name: str):
    """Load performance suite from rrsuites.py."""
    rrsuite = getattr(rrperf.rrsuites, suite_name, None)
    if rrsuite is None:
        rrsuite = getattr(rrperf.rrsuites_gfx1250, suite_name, None)
    return rrsuite()


def first_problem_from_suite(suite: str):
    for problem in load_suite(suite):
        return problem
    raise RuntimeError(f"Suite {suite} has no problems.")


def try_getting_commit(repo):
    if repo is not None:
        try:
            return rrperf.git.short_hash(repo)
        except Exception:
            pass
    return None


def get_commit(rundir: str = None, build_dir: Path = None) -> str:
    commit = try_getting_commit(build_dir)
    if commit is None:
        commit = try_getting_commit(rundir)
    if commit is None:
        commit = try_getting_commit(".")
    if commit is None:
        commit = try_getting_commit(Path(__file__).resolve().parent)
    if commit is None:
        commit = "NO_COMMIT"
    return commit


def get_work_dir(rundir: str = None, build_dir: Path = None) -> Path:
    """Return a new work directory path."""

    date = datetime.date.today().strftime("%Y-%m-%d")
    root = "."
    commit = get_commit(rundir, build_dir)

    if rundir is not None:
        root = Path(rundir)

    serial = len(list(Path(root).glob(f"{date}-{commit}-*")))
    return root / Path(f"{date}-{commit}-{serial:03d}")


def get_build_dir() -> Path:
    varname = "ROCROLLER_BUILD_DIR"
    if varname in os.environ:
        return Path(os.environ[varname])
    default = rrperf.git.top() / "shared" / "rocroller" / "build"
    if default.is_dir():
        return default

    raise RuntimeError(f"Build directory not found.  Set {varname} to override.")


def get_dataclass_id(obj):
    obj_dict = asdict(obj)
    for f in fields(obj):
        if not f.compare:
            del obj_dict[f.name]
    data_str = json.dumps(obj_dict, sort_keys=True)
    return sha1(data_str.encode()).hexdigest()


@functools.cache
def rocm_gfx():
    """Return GPU architecture (gfxXXXX) for local GPU device."""
    output = None
    try:
        output = subprocess.run(
            ["rocminfo"], capture_output=True, text=True, check=True
        ).stdout
    except subprocess.CalledProcessError:
        return None

    for line in output.splitlines():
        if line.startswith("  Name:"):
            _, arch, *_ = list(map(lambda x: x.strip(), line.split()))
            if arch.startswith("gfx"):
                return arch

    return None