From 6062ccd8cc3ccacd6923f2ea44959792ea249ebb Mon Sep 17 00:00:00 2001 From: hualxie Date: Thu, 13 Aug 2026 10:54:41 +0800 Subject: [PATCH 1/2] feat(optim): show custom operator domains Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/winml/modelkit/analyze/optim_output.py | 7 +++- src/winml/modelkit/optim/analysis.py | 16 +++++-- tests/unit/analyze/test_optim_output.py | 22 +++++++--- tests/unit/commands/test_optimize_cli.py | 24 ++++++++++- tests/unit/optim/test_analysis.py | 49 ++++++++++++++++++++++ 5 files changed, 105 insertions(+), 13 deletions(-) diff --git a/src/winml/modelkit/analyze/optim_output.py b/src/winml/modelkit/analyze/optim_output.py index 3cd868d45..b7a64e96a 100644 --- a/src/winml/modelkit/analyze/optim_output.py +++ b/src/winml/modelkit/analyze/optim_output.py @@ -127,11 +127,14 @@ def support_counts(self) -> dict[SupportLevel, int]: @staticmethod def _node_ref_dict(ref: NodeRef) -> dict[str, object]: """Return the stable JSON representation of one graph-delta node.""" - return { + data: dict[str, object] = { "op_type": ref.op_type, "name": ref.name, "outputs": list(ref.outputs), } + if ref.domain and ref.domain != "ai.onnx": + data["domain"] = ref.domain + return data def to_dict(self) -> dict[str, object]: """Return actionable graph-delta and target-support evidence.""" @@ -239,7 +242,7 @@ def _check_one( support, reason = _lookup_support(ref, support_by_output) result.operators.append( ProducedOperatorSupport( - op_type=ref.op_type, + op_type=ref.qualified_op_type(), label=ref.label(), change=change, support=support, diff --git a/src/winml/modelkit/optim/analysis.py b/src/winml/modelkit/optim/analysis.py index af69d1e9e..1d21e9c88 100644 --- a/src/winml/modelkit/optim/analysis.py +++ b/src/winml/modelkit/optim/analysis.py @@ -56,11 +56,19 @@ class NodeRef: op_type: The node's operator type (e.g. ``"MatMul"``). name: The node's name (may be empty — ONNX names are optional). outputs: The node's output tensor names. + domain: The node's operator domain. Empty means the default ONNX domain. """ op_type: str name: str outputs: tuple[str, ...] + domain: str = "" + + def qualified_op_type(self) -> str: + """Return the operator type, qualified when it uses a custom domain.""" + if self.domain and self.domain != "ai.onnx": + return f"{self.domain}::{self.op_type}" + return self.op_type def label(self) -> str: """Return a human-readable identifier for this node. @@ -69,7 +77,7 @@ def label(self) -> str: since output names are unique within a graph. """ ident = self.name or (self.outputs[0] if self.outputs else "?") - return f"{self.op_type} '{ident}'" + return f"{self.qualified_op_type()} '{ident}'" @dataclass @@ -136,7 +144,7 @@ def op_histogram(self, kind: str) -> list[tuple[str, int]]: "added": self.added_nodes, "modified": self.modified_nodes, }[kind] - return Counter(n.op_type for n in nodes).most_common() + return Counter(n.qualified_op_type() for n in nodes).most_common() # ============================================================================= @@ -163,7 +171,7 @@ def _node_identity(node: NodeProto) -> tuple[Any, ...]: """ if len(node.output) > 0: return tuple(node.output) - return ("\0no-output", node.op_type, node.name, tuple(node.input)) + return ("\0no-output", node.domain, node.op_type, node.name, tuple(node.input)) def _collect_nodes( @@ -186,7 +194,7 @@ def _collect_nodes( key = (cur_scope, _node_identity(node)) table[key] = ( node.SerializeToString(), - NodeRef(node.op_type, node.name, tuple(node.output)), + NodeRef(node.op_type, node.name, tuple(node.output), node.domain), ) for attr in node.attribute: if attr.type == AttributeProto.GRAPH: diff --git a/tests/unit/analyze/test_optim_output.py b/tests/unit/analyze/test_optim_output.py index 04cc9fd52..948d3475d 100644 --- a/tests/unit/analyze/test_optim_output.py +++ b/tests/unit/analyze/test_optim_output.py @@ -175,10 +175,15 @@ def test_to_dict_includes_graph_delta_and_target_support(self) -> None: description="Replace static Split with Slice.", pipe_name="algebraic", removed_nodes=[NodeRef("Split", "split", ("a", "b"))], - added_nodes=[NodeRef("Slice", "slice_0", ("a",))], + added_nodes=[NodeRef("Slice", "slice_0", ("a",), "com.microsoft")], modified_initializers=["starts"], operators=[ - ProducedOperatorSupport("Slice", "Slice 'slice_0'", "added", SupportLevel.SUPPORTED) + ProducedOperatorSupport( + "com.microsoft::Slice", + "com.microsoft::Slice 'slice_0'", + "added", + SupportLevel.SUPPORTED, + ) ], ) @@ -192,7 +197,14 @@ def test_to_dict_includes_graph_delta_and_target_support(self) -> None: "support_counts": {"supported": 1}, "graph_delta": { "removed_nodes": [{"op_type": "Split", "name": "split", "outputs": ["a", "b"]}], - "added_nodes": [{"op_type": "Slice", "name": "slice_0", "outputs": ["a"]}], + "added_nodes": [ + { + "op_type": "Slice", + "name": "slice_0", + "outputs": ["a"], + "domain": "com.microsoft", + } + ], "modified_nodes": [], "removed_initializers": [], "added_initializers": [], @@ -200,8 +212,8 @@ def test_to_dict_includes_graph_delta_and_target_support(self) -> None: }, "operators": [ { - "op_type": "Slice", - "label": "Slice 'slice_0'", + "op_type": "com.microsoft::Slice", + "label": "com.microsoft::Slice 'slice_0'", "change": "added", "support": "supported", } diff --git a/tests/unit/commands/test_optimize_cli.py b/tests/unit/commands/test_optimize_cli.py index 314c91b74..cab204edd 100644 --- a/tests/unit/commands/test_optimize_cli.py +++ b/tests/unit/commands/test_optimize_cli.py @@ -276,7 +276,9 @@ def test_device_target_forwarded_to_optimizer( _ANALYZE_MODEL = "winml.modelkit.optim.analyze_model" -def _make_finding(name: str = "clamp-constant-values") -> MagicMock: +def _make_finding( + name: str = "clamp-constant-values", node_domain: str = "" +) -> MagicMock: """Build a stand-in CapabilityFinding for renderer tests.""" from winml.modelkit.optim import CapabilityFinding, NodeRef @@ -288,7 +290,7 @@ def _make_finding(name: str = "clamp-constant-values") -> MagicMock: description="clamp things", pipe_name="surgery", modified_initializers=["BIG"], - removed_nodes=[NodeRef("MatMul", "mm", ("mm",))], + removed_nodes=[NodeRef("MatMul", "mm", ("mm",), node_domain)], ) @@ -335,6 +337,24 @@ def test_check_optim_lists_applicable_flag(self, runner: CliRunner, tmp_path: Pa assert "--enable-matmul-add-fusion" in result.output assert "1 applicable optimization" in result.output + def test_check_optim_shows_custom_operator_domain( + self, runner: CliRunner, tmp_path: Path + ) -> None: + model_file = tmp_path / "model.onnx" + model_file.touch() + + with ( + patch(_LOAD_ONNX, return_value=_make_mock_model()), + patch( + _ANALYZE_MODEL, + return_value=[_make_finding(node_domain="com.microsoft")], + ), + ): + result = runner.invoke(optimize, ["-m", str(model_file), "--check-optim"]) + + assert result.exit_code == 0, result.output + assert "com.microsoft::MatMul" in result.output + def test_check_optim_no_findings_message(self, runner: CliRunner, tmp_path: Path) -> None: model_file = tmp_path / "model.onnx" model_file.touch() diff --git a/tests/unit/optim/test_analysis.py b/tests/unit/optim/test_analysis.py index 82cb1a512..f98004541 100644 --- a/tests/unit/optim/test_analysis.py +++ b/tests/unit/optim/test_analysis.py @@ -137,6 +137,29 @@ def test_identical_graphs_have_no_diff(self) -> None: removed, added, modified = _diff_nodes(table_a, table_b) assert not removed and not added and not modified + def test_collected_node_preserves_custom_domain(self) -> None: + graph = helper.make_graph( + [ + helper.make_node( + "Gelu", + ["x"], + ["y"], + name="gelu", + domain="com.microsoft", + ) + ], + "custom_domain", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1])], + ) + table: dict = {} + + _collect_nodes(graph, (), table) + + ref = next(iter(table.values()))[1] + assert ref.domain == "com.microsoft" + assert ref.qualified_op_type() == "com.microsoft::Gelu" + def test_subgraph_nodes_are_collected(self) -> None: """Nodes inside a control-flow subgraph are included in the table.""" then_graph = helper.make_graph( @@ -505,6 +528,32 @@ def test_node_ref_label_uses_name_then_output(self) -> None: assert NodeRef("MatMul", "mm", ("mm_out",)).label() == "MatMul 'mm'" assert NodeRef("Add", "", ("y",)).label() == "Add 'y'" + def test_node_ref_label_qualifies_custom_domains(self) -> None: + assert ( + NodeRef("Gelu", "gelu", ("y",), "com.microsoft").label() + == "com.microsoft::Gelu 'gelu'" + ) + assert NodeRef("Add", "add", ("y",), "ai.onnx").label() == "Add 'add'" + + def test_op_histogram_distinguishes_custom_domains(self) -> None: + finding = CapabilityFinding( + name="x", + python_name="x", + enable_flag="--enable-x", + category="misc", + description="", + pipe_name="p", + added_nodes=[ + NodeRef("Gelu", "standard", ("a",)), + NodeRef("Gelu", "contrib", ("b",), "com.microsoft"), + ], + ) + + assert finding.op_histogram("added") == [ + ("Gelu", 1), + ("com.microsoft::Gelu", 1), + ] + # ============================================================================= # END-TO-END ANALYSIS From ec7c2814912865fe288b7ab73227232e6fc6f045 Mon Sep 17 00:00:00 2001 From: hualxie Date: Fri, 14 Aug 2026 09:43:54 +0800 Subject: [PATCH 2/2] test(analyze): cover default ONNX domain output Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- tests/unit/analyze/test_optim_output.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/analyze/test_optim_output.py b/tests/unit/analyze/test_optim_output.py index 948d3475d..38725400f 100644 --- a/tests/unit/analyze/test_optim_output.py +++ b/tests/unit/analyze/test_optim_output.py @@ -174,7 +174,7 @@ def test_to_dict_includes_graph_delta_and_target_support(self) -> None: category="rewrite", description="Replace static Split with Slice.", pipe_name="algebraic", - removed_nodes=[NodeRef("Split", "split", ("a", "b"))], + removed_nodes=[NodeRef("Split", "split", ("a", "b"), "ai.onnx")], added_nodes=[NodeRef("Slice", "slice_0", ("a",), "com.microsoft")], modified_initializers=["starts"], operators=[