rocm-systems/projects/rocprofiler-compute/src/utils/mi_gpu_spec.py

import os
from dataclasses import dataclass, field
from typing import Any, Dict

import yaml

from utils.logger import console_debug, console_error, console_log, console_warning

# Constants for MI series
# NOTE: Currently supports MI50, MI100, MI200, MI300
MI50 = 0
MI100 = 1
MI200 = 2
MI300 = 3
MI350 = 4

MI_CONSTANS = {
    MI50: "mi50",
    MI100: "mi100",
    MI200: "mi200",
    MI300: "mi300",
    MI350: "mi350",
}


# ----------------------------
# Data Class handling to preserve the hierarchical gpu information
# ----------------------------


@dataclass
class MIGPUSpecs:
    _instance = None

    _gpu_series_dict = {}  # key: gpu arch
    _gpu_model_dict = {}  # key: gpu_arch
    _num_xcds_dict = {}  # key: gpu model
    _chip_id_dict = {}  # key: chip id (int)

    _initialized = False

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._initialize()
        return cls._instance

    @classmethod
    def _initialize(cls):
        if not cls._initialized:
            cls._parse_mi_gpu_spec()
            cls._initialized = True

    # ----------------------------
    # YAML Parsing and Data Handling
    # ----------------------------

    @classmethod
    def _load_yaml(cls, file_path: str) -> Dict[str, Any]:
        """
        Loads MI GPU YAML data /util into a Python dictionary.

        Args:
            file_path (str): The path to the YAML file.

        Returns:
            Dict[str, Any]: Parsed YAML data as a nested dictionary.
                            Exit with console error if an error occurs.
        """
        console_debug("[load_yaml]")
        try:
            with open(file_path, "r") as file:
                data = yaml.safe_load(file)
                return data
        except FileNotFoundError:
            console_error(f"Error: The file '{file_path}' was not found.")
        except yaml.YAMLError as exc:
            console_error(f"Error parsing YAML file '{file_path}': {exc}")
        except Exception as e:
            console_error(
                f"An unexpected error occurred while loading YAML file '{file_path}': {e}"
            )

    @classmethod
    def _parse_mi_gpu_spec(cls):
        """
        Parse out mi gpu data from yaml file and store in memory.
        MI GPUs
        |-- series
            |-- architecture (list)
                    |-- gpu model
                    |-- chip_ids
                    |-- partition_mode
        """

        current_dir = os.path.dirname(__file__)
        yaml_file_path = os.path.join(current_dir, "mi_gpu_spec.yaml")

        # Load the YAML data
        yaml_data = cls._load_yaml(yaml_file_path)

        for series in yaml_data["mi_gpu_spec"]:
            curr_gpu_series = series["gpu_series"]
            console_debug("[parse_mi_gpu_spec] Processing series: %s" % curr_gpu_series)
            for archs in series["gpu_archs"]:
                curr_gpu_arch = archs["gpu_arch"]
                cls._gpu_series_dict[curr_gpu_arch] = curr_gpu_series
                cls._gpu_model_dict[curr_gpu_arch] = []
                for models in archs["models"]:
                    curr_gpu_model = models["gpu_model"]
                    cls._gpu_model_dict[curr_gpu_arch].append(curr_gpu_model)
                    cls._num_xcds_dict[curr_gpu_model] = (
                        models.get("partition_mode", {})
                        .get("compute_partition_mode", {})
                        .get("num_xcds", {})
                    )
                    if "chip_ids" in models and "physical" in models["chip_ids"]:
                        cls._chip_id_dict[models["chip_ids"]["physical"]] = curr_gpu_model
                    if "chip_ids" in models and "virtual" in models["chip_ids"]:
                        cls._chip_id_dict[models["chip_ids"]["virtual"]] = curr_gpu_model

    @classmethod
    def get_gpu_series_dict(cls):
        if not cls._gpu_series_dict:
            console_error(
                "gpu_series_dict not yet populated, did you run parse_mi_gpu_spec()?"
            )
            return None
        return cls._gpu_series_dict

    @classmethod
    def get_gpu_series(cls, gpu_arch_):
        if not cls._gpu_series_dict:
            console_error(
                "gpu_series_dict not yet populated, did you run parse_mi_gpu_spec()?"
            )
            return None

        # Normalize the key by checking both the raw and lowercase versions
        gpu_series = cls._gpu_series_dict.get(gpu_arch_) or cls._gpu_series_dict.get(
            gpu_arch_.lower()
        )
        if gpu_series:
            return gpu_series.upper()

        console_warning(f"No matching gpu series found for gpu arch: {gpu_arch_}")
        return None

    @classmethod
    def get_gpu_model(cls, gpu_arch_, chip_id_):
        # Check that gpu_model_dict is populated first
        if not cls._gpu_model_dict:
            console_error(
                "gpu_model_dict not yet populated. Did you run parse_mi_gpu_spec()?"
            )
            return None

        gpu_arch_lower = gpu_arch_.lower()

        # Handle gfx942 with chip_id mapping
        if gpu_arch_lower not in ("gfx906", "gfx908", "gfx90a"):
            if chip_id_ and int(chip_id_) in cls._chip_id_dict:
                gpu_model = cls._chip_id_dict.get(int(chip_id_))
            else:
                console_warning(f"No gpu model found for chip id: {chip_id_}")
                return None

        # Otherwise use gpu_model_dict mapping for other mi architectures
        elif gpu_arch_lower in cls._gpu_model_dict:
            # NOTE: take the first element works for now
            gpu_model = cls._gpu_model_dict[gpu_arch_lower][0]
        else:
            console_warning(f"No gpu model found for gpu arch: {gpu_arch_lower}")
            return None

        if not gpu_model:
            console_warning(f"No gpu model found for gpu arch: {gpu_arch_lower}")
            return None

        return gpu_model.upper()

    @classmethod
    def get_num_xcds(cls, gpu_model_, compute_partition_):
        # Only gpu in and above mi 300 series have more than one XCDs
        if gpu_model_.lower() in ("mi50", "mi60", "mi100", "mi210", "mi250", "mi250x"):
            return 1

        if not cls._num_xcds_dict:
            console_error(
                "mi300_num_xcds_dict not yet populated, did you run parse_mi_gpu_spec()?"
            )
            return None

        gpu_model_lower = gpu_model_.lower()
        partition_lower = compute_partition_.lower()

        if gpu_model_lower not in cls._num_xcds_dict:
            return None

        model_dict = cls._num_xcds_dict[gpu_model_lower]
        if partition_lower not in model_dict:
            console_log(f"Unknown compute partition: {compute_partition_}")
            return None

        num_xcds = model_dict[partition_lower]
        if not num_xcds:
            console_warning(
                "Unknown compute partition found for %s / %s",
                compute_partition_,
                gpu_model_,
            )
            return None

        return num_xcds

    @classmethod
    def get_chip_id_dict(cls):
        if cls._chip_id_dict:
            return cls._chip_id_dict
        else:
            console_error()


# pre-initialize the instance when module loads

mi_gpu_specs = MIGPUSpecs()