|
13 | 13 | if TYPE_CHECKING: |
14 | 14 | from pathlib import Path |
15 | 15 |
|
16 | | - from ...utils.constants import EPName |
| 16 | + from ...utils.constants import EPName, EPNameOrAlias |
17 | 17 | from ..models.ihv_type import IHVType |
18 | 18 |
|
19 | 19 |
|
20 | 20 | logger = logging.getLogger(__name__) |
21 | 21 |
|
22 | 22 |
|
23 | | -def infer_ihv_from_ep_name(ep_name: EPName) -> IHVType: |
24 | | - """Infer IHVType from Execution Provider name. |
| 23 | +def infer_ihv_from_ep_name(ep_name: EPNameOrAlias) -> IHVType: |
| 24 | + """Infer IHVType from an Execution Provider name or alias. |
25 | 25 |
|
26 | | - Maps an execution provider name to its corresponding IHV type. |
27 | | - Supports multiple name variations for each provider. |
28 | | - Unknown EPs (e.g., CPUExecutionProvider, DmlExecutionProvider) resolve |
29 | | - to IHVType.MICROSOFT. |
| 26 | + Accepts either a canonical ``EPName`` or a shorthand ``EPAlias`` (e.g. |
| 27 | + ``"openvino"``); aliases are normalized to their canonical name before the |
| 28 | + exact lookup, which covers every member of the canonical set. Names that |
| 29 | + are neither a known EP nor a known alias resolve to ``IHVType.UNKNOWN`` |
| 30 | + rather than raising, so callers can treat inference as total. |
30 | 31 |
|
31 | 32 | Args: |
32 | | - ep_name: Execution Provider name (e.g., QNNExecutionProvider, OpenVINOExecutionProvider) |
| 33 | + ep_name: Execution Provider name or alias (see ``utils.constants``). |
33 | 34 |
|
34 | 35 | Returns: |
35 | | - IHVType: Inferred IHV type (QC, INTEL, AMD, NVIDIA, or MICROSOFT) |
| 36 | + IHVType: Inferred IHV type (QC, INTEL, AMD, NVIDIA, MICROSOFT, or |
| 37 | + UNKNOWN for unrecognized names). |
36 | 38 |
|
37 | 39 | Examples: |
38 | 40 | >>> infer_ihv_from_ep_name("QNNExecutionProvider") |
39 | 41 | <IHVType.QC: 'QC'> |
40 | | - >>> infer_ihv_from_ep_name("OpenVINOExecutionProvider") |
41 | | - <IHVType.INTEL: 'INTEL'> |
| 42 | + >>> infer_ihv_from_ep_name("openvino") |
| 43 | + <IHVType.INTEL: 'Intel'> |
42 | 44 | >>> infer_ihv_from_ep_name("VitisAIExecutionProvider") |
43 | 45 | <IHVType.AMD: 'AMD'> |
44 | 46 | >>> infer_ihv_from_ep_name("NvTensorRTRTXExecutionProvider") |
45 | 47 | <IHVType.NVIDIA: 'NVIDIA'> |
46 | 48 | >>> infer_ihv_from_ep_name("CPUExecutionProvider") |
47 | 49 | <IHVType.MICROSOFT: 'Microsoft'> |
| 50 | + >>> infer_ihv_from_ep_name("TotallyFakeEP") |
| 51 | + <IHVType.UNKNOWN: 'Unknown'> |
48 | 52 | """ |
| 53 | + from ...utils.constants import normalize_ep_name |
49 | 54 | from ..models.ihv_type import IHVType |
50 | 55 |
|
51 | | - ep_lower = ep_name.lower() |
52 | | - |
53 | | - # QNN / Qualcomm |
54 | | - if "qnn" in ep_lower or "qualcomm" in ep_lower: |
55 | | - return IHVType.QC |
56 | | - |
57 | | - # OpenVINO / Intel |
58 | | - if "openvino" in ep_lower or "intel" in ep_lower: |
59 | | - return IHVType.INTEL |
60 | | - |
61 | | - # VitisAI / MIGraphX / AMD / ACE (AMD) |
62 | | - amd_keywords = ("amd", "quark", "vitis", "ace", "migraphx") |
63 | | - if any(kw in ep_lower for kw in amd_keywords): |
64 | | - return IHVType.AMD |
65 | | - |
66 | | - # NVIDIA / TensorRT RTX |
67 | | - # This is intentionally a permissive substring fallback to cover common |
68 | | - # TensorRT naming variants. Callers should prefer canonical EP names. |
69 | | - nvidia_keywords = ("nvidia", "nvtensorrt", "trtrtx", "tensorrt", "rtx") |
70 | | - if any(kw in ep_lower for kw in nvidia_keywords): |
71 | | - return IHVType.NVIDIA |
72 | | - |
73 | | - # Default: Microsoft (e.g., CPUExecutionProvider, DmlExecutionProvider) |
74 | | - return IHVType.MICROSOFT |
| 56 | + ep_name_to_ihv: dict[EPName, IHVType] = { |
| 57 | + "QNNExecutionProvider": IHVType.QC, |
| 58 | + "OpenVINOExecutionProvider": IHVType.INTEL, |
| 59 | + "VitisAIExecutionProvider": IHVType.AMD, |
| 60 | + "MIGraphXExecutionProvider": IHVType.AMD, |
| 61 | + "NvTensorRTRTXExecutionProvider": IHVType.NVIDIA, |
| 62 | + "CUDAExecutionProvider": IHVType.NVIDIA, |
| 63 | + "CPUExecutionProvider": IHVType.MICROSOFT, |
| 64 | + "DmlExecutionProvider": IHVType.MICROSOFT, |
| 65 | + } |
| 66 | + |
| 67 | + canonical = normalize_ep_name(ep_name) |
| 68 | + return ep_name_to_ihv.get(canonical, IHVType.UNKNOWN) # type: ignore[arg-type] |
75 | 69 |
|
76 | 70 |
|
77 | 71 | def get_devices_with_rule_data(ep_name: EPName) -> list[str]: |
|
0 commit comments