diff --git a/src/winml/modelkit/analyze/core/information_engine.py b/src/winml/modelkit/analyze/core/information_engine.py index a3cfe1e1f..1732253c1 100644 --- a/src/winml/modelkit/analyze/core/information_engine.py +++ b/src/winml/modelkit/analyze/core/information_engine.py @@ -117,19 +117,24 @@ def __init__( self._device = device # Load predefined information rules + from ..models.ihv_type import IHVType from ..utils import infer_ihv_from_ep_name from ..utils.rule_loader import RuleLoader self._rule_loader = RuleLoader(rules_dir=rules_dir) - # Infer IHV from EP name for per-IHV rule loading + # Infer IHV from EP name for per-IHV rule loading. An unrecognized EP + # resolves to IHVType.UNKNOWN, which we treat as "no IHV filter" so the + # loader falls back to loading all rules. infer_ihv_start = time.perf_counter() - try: - ihv_type = infer_ihv_from_ep_name(self._ep) - logger.info("Inferred IHV type %s from EP %s", ihv_type.value, self._ep) - except ValueError as e: - logger.warning("Could not infer IHV from EP %s: %s. Loading all rules.", self._ep, e) + inferred_ihv = infer_ihv_from_ep_name(self._ep) + ihv_type: IHVType | None + if inferred_ihv is IHVType.UNKNOWN: + logger.warning("Could not infer IHV from EP %s. Loading all rules.", self._ep) ihv_type = None + else: + logger.info("Inferred IHV type %s from EP %s", inferred_ihv.value, self._ep) + ihv_type = inferred_ihv infer_ihv_ms = int((time.perf_counter() - infer_ihv_start) * 1000) load_predefined_start = time.perf_counter() diff --git a/src/winml/modelkit/analyze/core/model_validators/batched_const_matmul_validator.py b/src/winml/modelkit/analyze/core/model_validators/batched_const_matmul_validator.py index 91a3114c1..0f3ad2b32 100644 --- a/src/winml/modelkit/analyze/core/model_validators/batched_const_matmul_validator.py +++ b/src/winml/modelkit/analyze/core/model_validators/batched_const_matmul_validator.py @@ -53,12 +53,9 @@ def _is_enabled(self) -> bool: ep = self.ep if not ep: return False - try: - from ...models.ihv_type import IHVType + from ...models.ihv_type import IHVType - return infer_ihv_from_ep_name(ep) == IHVType.INTEL - except Exception: # pragma: no cover - defensive - return False + return infer_ihv_from_ep_name(ep) == IHVType.INTEL def validate(self) -> Information | None: """Detect batched MatMul with a single constant rank>=3 operand.""" diff --git a/src/winml/modelkit/analyze/models/ihv_type.py b/src/winml/modelkit/analyze/models/ihv_type.py index e42503af8..f517c948f 100644 --- a/src/winml/modelkit/analyze/models/ihv_type.py +++ b/src/winml/modelkit/analyze/models/ihv_type.py @@ -15,3 +15,4 @@ class IHVType(str, Enum): AMD = "AMD" NVIDIA = "NVIDIA" MICROSOFT = "Microsoft" + UNKNOWN = "Unknown" diff --git a/src/winml/modelkit/analyze/utils/ep_utils.py b/src/winml/modelkit/analyze/utils/ep_utils.py index 5f29b2d99..e4ab29864 100644 --- a/src/winml/modelkit/analyze/utils/ep_utils.py +++ b/src/winml/modelkit/analyze/utils/ep_utils.py @@ -13,65 +13,59 @@ if TYPE_CHECKING: from pathlib import Path - from ...utils.constants import EPName + from ...utils.constants import EPName, EPNameOrAlias from ..models.ihv_type import IHVType logger = logging.getLogger(__name__) -def infer_ihv_from_ep_name(ep_name: EPName) -> IHVType: - """Infer IHVType from Execution Provider name. +def infer_ihv_from_ep_name(ep_name: EPNameOrAlias) -> IHVType: + """Infer IHVType from an Execution Provider name or alias. - Maps an execution provider name to its corresponding IHV type. - Supports multiple name variations for each provider. - Unknown EPs (e.g., CPUExecutionProvider, DmlExecutionProvider) resolve - to IHVType.MICROSOFT. + Accepts either a canonical ``EPName`` or a shorthand ``EPAlias`` (e.g. + ``"openvino"``); aliases are normalized to their canonical name before the + exact lookup, which covers every member of the canonical set. Names that + are neither a known EP nor a known alias resolve to ``IHVType.UNKNOWN`` + rather than raising, so callers can treat inference as total. Args: - ep_name: Execution Provider name (e.g., QNNExecutionProvider, OpenVINOExecutionProvider) + ep_name: Execution Provider name or alias (see ``utils.constants``). Returns: - IHVType: Inferred IHV type (QC, INTEL, AMD, NVIDIA, or MICROSOFT) + IHVType: Inferred IHV type (QC, INTEL, AMD, NVIDIA, MICROSOFT, or + UNKNOWN for unrecognized names). Examples: >>> infer_ihv_from_ep_name("QNNExecutionProvider") - >>> infer_ihv_from_ep_name("OpenVINOExecutionProvider") - + >>> infer_ihv_from_ep_name("openvino") + >>> infer_ihv_from_ep_name("VitisAIExecutionProvider") >>> infer_ihv_from_ep_name("NvTensorRTRTXExecutionProvider") >>> infer_ihv_from_ep_name("CPUExecutionProvider") + >>> infer_ihv_from_ep_name("TotallyFakeEP") + """ + from ...utils.constants import normalize_ep_name from ..models.ihv_type import IHVType - ep_lower = ep_name.lower() - - # QNN / Qualcomm - if "qnn" in ep_lower or "qualcomm" in ep_lower: - return IHVType.QC - - # OpenVINO / Intel - if "openvino" in ep_lower or "intel" in ep_lower: - return IHVType.INTEL - - # VitisAI / MIGraphX / AMD / ACE (AMD) - amd_keywords = ("amd", "quark", "vitis", "ace", "migraphx") - if any(kw in ep_lower for kw in amd_keywords): - return IHVType.AMD - - # NVIDIA / TensorRT RTX - # This is intentionally a permissive substring fallback to cover common - # TensorRT naming variants. Callers should prefer canonical EP names. - nvidia_keywords = ("nvidia", "nvtensorrt", "trtrtx", "tensorrt", "rtx") - if any(kw in ep_lower for kw in nvidia_keywords): - return IHVType.NVIDIA - - # Default: Microsoft (e.g., CPUExecutionProvider, DmlExecutionProvider) - return IHVType.MICROSOFT + ep_name_to_ihv: dict[EPName, IHVType] = { + "QNNExecutionProvider": IHVType.QC, + "OpenVINOExecutionProvider": IHVType.INTEL, + "VitisAIExecutionProvider": IHVType.AMD, + "MIGraphXExecutionProvider": IHVType.AMD, + "NvTensorRTRTXExecutionProvider": IHVType.NVIDIA, + "CUDAExecutionProvider": IHVType.NVIDIA, + "CPUExecutionProvider": IHVType.MICROSOFT, + "DmlExecutionProvider": IHVType.MICROSOFT, + } + + canonical = normalize_ep_name(ep_name) + return ep_name_to_ihv.get(canonical, IHVType.UNKNOWN) # type: ignore[arg-type] def get_devices_with_rule_data(ep_name: EPName) -> list[str]: diff --git a/src/winml/modelkit/compiler/utils.py b/src/winml/modelkit/compiler/utils.py index ea80ed480..c8867c8aa 100644 --- a/src/winml/modelkit/compiler/utils.py +++ b/src/winml/modelkit/compiler/utils.py @@ -12,27 +12,33 @@ if TYPE_CHECKING: from pathlib import Path - from ..utils.constants import EPAlias + from ..utils.constants import EPNameOrAlias # Canonical definition of ONNX QDQ operator types. # Import this constant instead of redefining {"QuantizeLinear", "DequantizeLinear"}. QDQ_OP_TYPES: frozenset[str] = frozenset({"QuantizeLinear", "DequantizeLinear"}) -def needs_format_conversion(model_path: Path, ep: EPAlias) -> bool: +def needs_format_conversion(model_path: Path, ep: EPNameOrAlias) -> bool: """Check if model's quant format is compatible with target EP. Minimal detection: checks for QLinear ops targeting QDQ-only EPs. FIXME: Expand to full EP-to-format compatibility matrix. """ from ..onnx import load_onnx + from ..utils.constants import normalize_ep_name model = load_onnx(model_path, load_weights=False, validate=False) op_types = {n.op_type for n in model.graph.node} has_qlinear = any(op.startswith("QLinear") for op in op_types) has_qdq = bool(op_types & QDQ_OP_TYPES) - if ep == "qnn" and has_qlinear and not has_qdq: # noqa: SIM103 + # Compare against the canonical EP name, not a single alias: one EP can have + # several aliases (e.g. nv_tensorrt_rtx / nvtensorrtrtx), so an alias-literal + # comparison would miss the others. + ep_canonical = normalize_ep_name(ep) + + if ep_canonical == "QNNExecutionProvider" and has_qlinear and not has_qdq: # noqa: SIM103 return True # QNN requires QDQ format # FIXME: add more EP rules as needed return False diff --git a/src/winml/modelkit/models/auto.py b/src/winml/modelkit/models/auto.py index 4849f17e6..4b5eb0050 100644 --- a/src/winml/modelkit/models/auto.py +++ b/src/winml/modelkit/models/auto.py @@ -363,19 +363,22 @@ def from_pretrained( if task is not None: from .winml.composite_model import COMPOSITE_MODEL_REGISTRY - _known_composite_tasks = {t for (_, t) in COMPOSITE_MODEL_REGISTRY} - if task in _known_composite_tasks: - from transformers import AutoConfig - - _hf_cfg = AutoConfig.from_pretrained(model_id, trust_remote_code=trust_remote_code) - _model_type = getattr(_hf_cfg, "model_type", None) - else: - _model_type = None - # Explicit override wins so a variant composite (e.g. # "qwen3_transformer_only") can be selected over the native type. + _model_type: str | None if model_type is not None: _model_type = model_type + else: + _known_composite_tasks = {t for (_, t) in COMPOSITE_MODEL_REGISTRY} + if task in _known_composite_tasks: + from transformers import AutoConfig + + _hf_cfg = AutoConfig.from_pretrained( + model_id, trust_remote_code=trust_remote_code + ) + _model_type = getattr(_hf_cfg, "model_type", None) + else: + _model_type = None if _model_type is not None and (_model_type, task) in COMPOSITE_MODEL_REGISTRY: from .winml.composite_model import WinMLCompositeModel diff --git a/src/winml/modelkit/serve/app.py b/src/winml/modelkit/serve/app.py index b38f96f6b..a666c237f 100644 --- a/src/winml/modelkit/serve/app.py +++ b/src/winml/modelkit/serve/app.py @@ -490,7 +490,7 @@ async def get_mcp_schema() -> dict[str, Any]: # ------------------------------------------------------------------ @app.post("/v1/ep", tags=["management"], summary="Switch execution provider") async def switch_ep(request: EpSwitchRequest) -> dict[str, Any]: - # Pydantic already validates ep against the EPAlias Literal (rejects + # Pydantic already validates ep against the EPNameOrAlias Literal (rejects # unknown values with a 422 at parse time), so no extra check needed. ep = request.ep mgr = _get_mgr() diff --git a/src/winml/modelkit/serve/schema.py b/src/winml/modelkit/serve/schema.py index 09ff2232f..2f55bc968 100644 --- a/src/winml/modelkit/serve/schema.py +++ b/src/winml/modelkit/serve/schema.py @@ -14,7 +14,7 @@ from pydantic import BaseModel, Field -from ..utils.constants import EPAlias, EPNameOrAlias +from ..utils.constants import EPNameOrAlias # --------------------------------------------------------------------------- @@ -25,7 +25,9 @@ class EpSwitchRequest(BaseModel): """POST /v1/ep — switch execution provider.""" - ep: EPAlias = Field(..., description="EP short name: cpu, dml, qnn, openvino") + ep: EPNameOrAlias = Field( + ..., description="EP name or short alias (e.g. cpu, dml, qnn, QNNExecutionProvider)" + ) class PredictJsonRequest(BaseModel): diff --git a/tests/unit/analyze/core/test_output_aggregator.py b/tests/unit/analyze/core/test_output_aggregator.py index 014fb6fb6..d9803ade8 100644 --- a/tests/unit/analyze/core/test_output_aggregator.py +++ b/tests/unit/analyze/core/test_output_aggregator.py @@ -510,7 +510,7 @@ def test_full_workflow_multiple_ihv(self, sample_metadata: ModelStats) -> None: result=RuntimeTestResult(compile=True, run=True), ), ], - "ACEExecutionProvider": [ + "VitisAIExecutionProvider": [ PatternRuntime( pattern_id="OP/ai.onnx/Add", result=RuntimeTestResult(compile=False, run=False), @@ -521,7 +521,7 @@ def test_full_workflow_multiple_ihv(self, sample_metadata: ModelStats) -> None: information_list = { "QNNExecutionProvider": [], "OpenVINOExecutionProvider": [], - "ACEExecutionProvider": [ + "VitisAIExecutionProvider": [ Information( explanation="Add not supported", pattern_id="OP/ai.onnx/Add", diff --git a/tests/unit/analyze/test_analyzer.py b/tests/unit/analyze/test_analyzer.py index d4de6fd0a..dcfa735f2 100644 --- a/tests/unit/analyze/test_analyzer.py +++ b/tests/unit/analyze/test_analyzer.py @@ -23,7 +23,6 @@ ModelStats, ONNXStaticAnalyzer, SupportLevel, - infer_ihv_from_ep_name, ) from winml.modelkit.analyze.analyzer import _build_runtime_debug_details_summary from winml.modelkit.analyze.models.runtime_checks import PatternRuntime, RuntimeTestResult @@ -864,34 +863,6 @@ def test_init_custom_config(self) -> None: assert analyzer.config.enable_information is True assert analyzer.config.max_memory_mb == 4096 - def test_map_ep_to_ihv_qnn(self) -> None: - """Test EP to IHV mapping for QNN.""" - assert infer_ihv_from_ep_name("QNNExecutionProvider") == IHVType.QC - assert infer_ihv_from_ep_name("qnnexecutionprovider") == IHVType.QC - assert infer_ihv_from_ep_name("QualcommProvider") == IHVType.QC - - def test_map_ep_to_ihv_openvino(self) -> None: - """Test EP to IHV mapping for OpenVINO.""" - assert infer_ihv_from_ep_name("OpenVINOExecutionProvider") == IHVType.INTEL - assert infer_ihv_from_ep_name("openvino") == IHVType.INTEL - assert infer_ihv_from_ep_name("IntelProvider") == IHVType.INTEL - - def test_map_ep_to_ihv_vitisai(self) -> None: - """Test EP to IHV mapping for VitisAI.""" - assert infer_ihv_from_ep_name("VitisAIExecutionProvider") == IHVType.AMD - assert infer_ihv_from_ep_name("vitis") == IHVType.AMD - assert infer_ihv_from_ep_name("AMDProvider") == IHVType.AMD - - def test_map_ep_to_ihv_nvidia(self) -> None: - """Test EP to IHV mapping for NvTensorRTRTX.""" - assert infer_ihv_from_ep_name("NvTensorRTRTXExecutionProvider") == IHVType.NVIDIA - assert infer_ihv_from_ep_name("nvtensorrtx") == IHVType.NVIDIA - assert infer_ihv_from_ep_name("TensorRTProvider") == IHVType.NVIDIA - - def test_map_ep_to_ihv_invalid(self) -> None: - """Test EP to IHV mapping with unrecognized EP resolves to MICROSOFT.""" - assert infer_ihv_from_ep_name("InvalidEP") == IHVType.MICROSOFT - def test_analyze_file_not_found(self) -> None: """Test analyze with non-existent file.""" analyzer = ONNXStaticAnalyzer() diff --git a/tests/unit/analyze/test_has_rule_data.py b/tests/unit/analyze/test_has_rule_data.py index 209e8fc9c..5e8539b26 100644 --- a/tests/unit/analyze/test_has_rule_data.py +++ b/tests/unit/analyze/test_has_rule_data.py @@ -27,6 +27,20 @@ class TestInferIHVFromEPName: """Tests for infer_ihv_from_ep_name().""" + def test_all_known_eps_resolve(self) -> None: + """Every canonical EPName maps to a valid IHVType (map covers the Literal).""" + from winml.modelkit.analyze.models.ihv_type import IHVType + from winml.modelkit.utils.constants import EP_NAMES + + for ep in EP_NAMES: + assert isinstance(infer_ihv_from_ep_name(ep), IHVType) + + def test_unknown_ep_resolves_to_unknown(self) -> None: + """Unknown EP names resolve to IHVType.UNKNOWN rather than raising.""" + from winml.modelkit.analyze.models.ihv_type import IHVType + + assert infer_ihv_from_ep_name("TotallyFakeEP") == IHVType.UNKNOWN + def test_qnn(self) -> None: from winml.modelkit.analyze.models.ihv_type import IHVType @@ -48,18 +62,12 @@ def test_migraphx_maps_to_amd(self) -> None: assert infer_ihv_from_ep_name("MIGraphXExecutionProvider") == IHVType.AMD - def test_case_insensitive(self) -> None: - from winml.modelkit.analyze.models.ihv_type import IHVType - - assert infer_ihv_from_ep_name("qnnexecutionprovider") == IHVType.QC - assert infer_ihv_from_ep_name("OPENVINOEXECUTIONPROVIDER") == IHVType.INTEL - assert infer_ihv_from_ep_name("vitisaiexecutionprovider") == IHVType.AMD - assert infer_ihv_from_ep_name("nvtensorrtxexecutionprovider") == IHVType.NVIDIA - - def test_unknown_ep_resolves_to_microsoft(self) -> None: + def test_alias_resolves(self) -> None: + """Shorthand aliases are normalized before lookup (EPNameOrAlias).""" from winml.modelkit.analyze.models.ihv_type import IHVType - assert infer_ihv_from_ep_name("TotallyFakeEP") == IHVType.MICROSOFT + assert infer_ihv_from_ep_name("openvino") == IHVType.INTEL + assert infer_ihv_from_ep_name("qnn") == IHVType.QC def test_cpu_ep_resolves_to_microsoft(self) -> None: """CPUExecutionProvider is a Microsoft EP — should resolve to MICROSOFT.""" @@ -79,11 +87,11 @@ def test_nvidia_ep_maps_to_nvidia(self) -> None: assert infer_ihv_from_ep_name("NvTensorRTRTXExecutionProvider") == IHVType.NVIDIA - def test_trtrtx_ep_maps_to_nvidia(self) -> None: - """TrtRTXExecutionProvider should map to IHVType.NVIDIA.""" + def test_cuda_ep_maps_to_nvidia(self) -> None: + """CUDAExecutionProvider should map to IHVType.NVIDIA.""" from winml.modelkit.analyze.models.ihv_type import IHVType - assert infer_ihv_from_ep_name("TrtRTXExecutionProvider") == IHVType.NVIDIA + assert infer_ihv_from_ep_name("CUDAExecutionProvider") == IHVType.NVIDIA class TestHasRuleDataForEP: diff --git a/tests/unit/compiler/test_utils.py b/tests/unit/compiler/test_utils.py index 7fbe3d3a0..078cd559f 100644 --- a/tests/unit/compiler/test_utils.py +++ b/tests/unit/compiler/test_utils.py @@ -86,6 +86,13 @@ def test_qlinear_for_qnn(self, tmp_path: Path) -> None: onnx.save(model, str(path)) assert needs_format_conversion(path, "qnn") is True + def test_qlinear_for_qnn_canonical_name(self, tmp_path: Path) -> None: + """Canonical EP name must be recognized, not just the alias.""" + model = _make_simple_model(["QLinearConv", "Relu"]) + path = tmp_path / "qlinear.onnx" + onnx.save(model, str(path)) + assert needs_format_conversion(path, "QNNExecutionProvider") is True + def test_qdq_for_qnn(self, tmp_path: Path) -> None: model = _make_simple_model(["QuantizeLinear", "DequantizeLinear"]) path = tmp_path / "qdq.onnx"