xuchen-amd
385 lignes
14 KiB
Python
385 lignes
14 KiB
Python
##############################################################################
|
|
# MIT License
|
|
#
|
|
# Copyright (c) 2025 Advanced Micro Devices, Inc. All Rights Reserved.
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
# copies of the Software, and to permit persons to whom the Software is
|
|
# furnished to do so, subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in
|
|
# all copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
|
|
# THE SOFTWARE.
|
|
|
|
##############################################################################
|
|
|
|
from pathlib import Path
|
|
from typing import Any, Optional
|
|
|
|
import yaml
|
|
|
|
from utils.logger import console_debug, console_error, console_warning
|
|
|
|
# ----------------------------
|
|
# Data Class handling to preserve the hierarchical gpu information
|
|
# ----------------------------
|
|
|
|
|
|
class MIGPUSpecs:
|
|
_instance: Optional["MIGPUSpecs"] = None
|
|
|
|
_gpu_series_dict: dict[str, str] = {} # key: gpu_arch
|
|
_gpu_model_dict: dict[str, list[str]] = {} # key: gpu_arch
|
|
_num_xcds_dict: dict[str, dict[str, int]] = {} # key: gpu_model
|
|
_chip_id_dict: dict[int, str] = {} # key: chip_id (int)
|
|
_perfmon_config: dict[str, Any] = {} # key: gpu_arch
|
|
|
|
# key: gpu_arch, used for gpu archs containing only one gpu model and
|
|
# thus one compute partition
|
|
_gpu_arch_to_compute_partition_dict: dict[str, dict[str, int]] = {}
|
|
|
|
_all_gpu_models: list[str] = []
|
|
|
|
_initialized = False
|
|
|
|
def __new__(cls) -> "MIGPUSpecs":
|
|
if cls._instance is None:
|
|
cls._instance = super().__new__(cls)
|
|
cls._initialize()
|
|
return cls._instance
|
|
|
|
@classmethod
|
|
def _initialize(cls) -> None:
|
|
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("mi_gpu_spec", "[load_yaml]")
|
|
with open(file_path) as file:
|
|
data = yaml.safe_load(file)
|
|
return data or {}
|
|
|
|
@classmethod
|
|
def _parse_mi_gpu_spec(cls) -> None:
|
|
"""
|
|
Parse out mi gpu data from yaml file and store in memory.
|
|
MI GPUs
|
|
|-- series
|
|
|-- architecture (list)
|
|
|-- perfmon_config
|
|
|-- gpu model
|
|
|-- chip_ids
|
|
| -- physical
|
|
| -- virtual
|
|
|-- partition_mode
|
|
| -- compute partition mode
|
|
| -- memory partition mode
|
|
"""
|
|
|
|
current_dir = Path(__file__).parent
|
|
yaml_file_path = current_dir / "mi_gpu_spec.yaml"
|
|
|
|
# Load the YAML data
|
|
yaml_data = cls._load_yaml(yaml_file_path)
|
|
|
|
for series in yaml_data.get("mi_gpu_spec", []):
|
|
curr_gpu_series = series.get("gpu_series")
|
|
if not curr_gpu_series:
|
|
continue
|
|
|
|
console_debug(
|
|
"mi_gpu_spec",
|
|
f"[parse_mi_gpu_spec] Processing series: {curr_gpu_series}",
|
|
)
|
|
|
|
for archs in series.get("gpu_archs", []):
|
|
curr_gpu_arch = archs.get("gpu_arch")
|
|
|
|
cls._gpu_series_dict[curr_gpu_arch] = curr_gpu_series
|
|
cls._perfmon_config[curr_gpu_arch] = archs.get("perfmon_config", {})
|
|
cls._gpu_model_dict[curr_gpu_arch] = []
|
|
|
|
for models in archs.get("models", []):
|
|
curr_gpu_model = models["gpu_model"]
|
|
|
|
cls._all_gpu_models.append(curr_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
|
|
)
|
|
|
|
# detect gpu arch to compute partition relationships
|
|
cls._populate_gpu_arch_to_compute_partition_dict()
|
|
|
|
@classmethod
|
|
def _populate_gpu_arch_to_compute_partition_dict(cls) -> None:
|
|
"""
|
|
This creates a mapping of gpu_arch -> compute_partition for architectures
|
|
where there's only one model (and therefore one partition configuration).
|
|
"""
|
|
for gpu_arch, gpu_models in cls._gpu_model_dict.items():
|
|
if len(gpu_models) == 1:
|
|
single_model = gpu_models[0]
|
|
compute_partition = cls._num_xcds_dict.get(single_model)
|
|
|
|
if compute_partition is not None:
|
|
cls._gpu_arch_to_compute_partition_dict[gpu_arch] = (
|
|
compute_partition
|
|
)
|
|
console_debug(
|
|
f"[populate_single_arch_partition_dict] Single model "
|
|
f"arch found: {gpu_arch} -> {single_model} (partition:"
|
|
f" {compute_partition})"
|
|
)
|
|
|
|
@classmethod
|
|
def get_gpu_series_dict(cls) -> dict[str, str]:
|
|
if not cls._gpu_series_dict:
|
|
console_error(
|
|
"gpu_series_dict not yet populated, did you run parse_mi_gpu_spec()?",
|
|
exit=False,
|
|
)
|
|
return cls._gpu_series_dict
|
|
|
|
@classmethod
|
|
def get_gpu_series(cls, gpu_arch: str) -> Optional[str]:
|
|
if not cls._gpu_series_dict:
|
|
console_error(
|
|
"gpu_series_dict not yet populated, did you run parse_mi_gpu_spec()?",
|
|
exit=False,
|
|
)
|
|
|
|
# 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_perfmon_config(cls, gpu_arch: str) -> dict[Any, Any]:
|
|
# Check that gpu_model_dict is populated first
|
|
if not cls._perfmon_config:
|
|
console_error(
|
|
"gpu_model_dict not yet populated. Did you run parse_mi_gpu_spec()?"
|
|
)
|
|
|
|
return cls._perfmon_config.get(gpu_arch.lower(), {})
|
|
|
|
@classmethod
|
|
def get_gpu_model(
|
|
cls, gpu_arch: Optional[str], chip_id: Optional[str] = None
|
|
) -> Optional[str]:
|
|
if not gpu_arch and not chip_id:
|
|
return None
|
|
|
|
# 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 ("gfx908", "gfx90a"):
|
|
if chip_id and chip_id.isdigit():
|
|
chip_id_int = int(chip_id)
|
|
if chip_id_int in cls._chip_id_dict:
|
|
gpu_model = cls._chip_id_dict[chip_id_int]
|
|
else:
|
|
console_warning(f"No gpu model found for chip id: {chip_id}")
|
|
return None
|
|
else:
|
|
console_warning(f"No valid chip id provided: {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_models = cls._gpu_model_dict[gpu_arch_lower]
|
|
if gpu_models:
|
|
gpu_model = gpu_models[0]
|
|
else:
|
|
console_warning(f"No gpu models found for gpu arch: {gpu_arch_lower}")
|
|
return None
|
|
else:
|
|
console_warning(f"No gpu model found for gpu arch: {gpu_arch_lower}")
|
|
return None
|
|
|
|
return gpu_model.upper() if gpu_model else None
|
|
|
|
@classmethod
|
|
def set_default_gpu_settings(
|
|
cls,
|
|
gpu_arch: Optional[str],
|
|
gpu_model: Optional[str],
|
|
compute_partition: Optional[str],
|
|
) -> int:
|
|
"""
|
|
Set default GPU settings when model is unknown or cannot be
|
|
determined. NOTE: This is a fallback to gfx942 settings -
|
|
consider making this architecture-specific.
|
|
"""
|
|
|
|
DEFAULT_COMPUTE_PARTITION = "SPX"
|
|
DEFAULT_NUM_XCD = 8
|
|
console_warning(
|
|
"Unable to determine xcd count from:\n\t"
|
|
f'GPU arch: "{gpu_arch}", model: "{gpu_model}",\n\t'
|
|
f'partition: "{compute_partition}"'
|
|
)
|
|
console_warning(
|
|
f"Applying default gfx942 settings:\n"
|
|
f"\t- Compute partition: {DEFAULT_COMPUTE_PARTITION}\n"
|
|
f"\t- Number of XCDs: {DEFAULT_NUM_XCD}"
|
|
)
|
|
|
|
return DEFAULT_NUM_XCD
|
|
|
|
@classmethod
|
|
def get_num_xcds(
|
|
cls,
|
|
gpu_arch: Optional[str] = None,
|
|
gpu_model: Optional[str] = None,
|
|
compute_partition: Optional[str] = None,
|
|
) -> int:
|
|
"""
|
|
Retrieve the number of XCDs based on GPU architecture, model,
|
|
and compute partition.
|
|
|
|
Priority order:
|
|
1. Legacy GPU check (returns 1 XCD for older architectures/models)
|
|
2. Architecture-based lookup (preferred)
|
|
3. Model + partition-based lookup (fallback)
|
|
4. Default settings (last resort)
|
|
"""
|
|
# Constants for legacy GPUs that don't support compute partitions
|
|
LEGACY_ARCHS = {"gfx908", "gfx90a"}
|
|
LEGACY_MODELS = {"mi50", "mi60", "mi100", "mi210", "mi250", "mi250x"}
|
|
|
|
# Normalize inputs to lowercase for consistent comparison
|
|
gpu_arch_norm = gpu_arch.lower().strip() if gpu_arch else ""
|
|
gpu_model_norm = gpu_model.lower().strip() if gpu_model else ""
|
|
partition_norm = compute_partition.lower().strip() if compute_partition else ""
|
|
|
|
# 1. Return 1 XCDs for archs/models not supporting compute partition
|
|
# NOTE: gpu arch is enough to verify this logic, gpu model is used as a backup.
|
|
if gpu_arch_norm in LEGACY_ARCHS or gpu_model_norm in LEGACY_MODELS:
|
|
return 1
|
|
|
|
# 2. Try architecture-based lookup first (preferred method)
|
|
if gpu_arch_norm and gpu_arch_norm in cls._gpu_arch_to_compute_partition_dict:
|
|
arch_dict = cls._gpu_arch_to_compute_partition_dict[gpu_arch_norm]
|
|
if partition_norm and partition_norm in arch_dict:
|
|
num_xcds = arch_dict[partition_norm]
|
|
if num_xcds is not None:
|
|
return num_xcds
|
|
else:
|
|
console_warning(
|
|
f"No compute partition data found for "
|
|
f"architecture: {gpu_arch.upper() if gpu_arch else None}"
|
|
)
|
|
|
|
# 3. Fall back to model + partition-based lookup
|
|
if gpu_model_norm:
|
|
# Validate XCD dictionary is populated
|
|
if not cls._num_xcds_dict:
|
|
console_error(
|
|
"mi300_num_xcds_dict not populated. "
|
|
"Did you run parse_mi_gpu_spec()?"
|
|
)
|
|
elif gpu_model_norm not in cls._num_xcds_dict:
|
|
console_warning(
|
|
f"Unknown gpu model provided for num xcds lookup: {gpu_model}."
|
|
)
|
|
else:
|
|
model_dict = cls._num_xcds_dict[gpu_model_norm]
|
|
|
|
if not partition_norm:
|
|
console_warning(
|
|
"No compute partition provided for model-based lookup"
|
|
)
|
|
elif partition_norm not in model_dict:
|
|
console_warning(
|
|
f"Unknown compute partition: "
|
|
f"{compute_partition} for model: {gpu_model}"
|
|
)
|
|
else:
|
|
num_xcds = model_dict[partition_norm]
|
|
if num_xcds is not None:
|
|
return num_xcds
|
|
else:
|
|
console_warning(
|
|
"Unknown compute partition found "
|
|
f"for {compute_partition} / {gpu_model}"
|
|
)
|
|
|
|
else:
|
|
console_warning("No gpu model provided for num xcds lookup.")
|
|
|
|
# 4. Last resort: use default settings
|
|
return cls.set_default_gpu_settings(gpu_arch, gpu_model, compute_partition)
|
|
|
|
@classmethod
|
|
def get_chip_id_dict(cls) -> dict[int, str]:
|
|
return cls._chip_id_dict
|
|
|
|
@classmethod
|
|
def get_num_xcds_dict(cls) -> dict[str, dict[str, int]]:
|
|
return cls._num_xcds_dict
|
|
|
|
@classmethod
|
|
def get_gpu_arch_to_compute_partition_dict(cls) -> dict[str, dict[str, int]]:
|
|
return cls._gpu_arch_to_compute_partition_dict
|
|
|
|
@classmethod
|
|
def get_all_gpu_models(cls) -> list:
|
|
return cls._all_gpu_models
|
|
|
|
|
|
# pre-initialize the instance when module loads
|
|
mi_gpu_specs = MIGPUSpecs()
|