Skip to content

Commit c72e810

Browse files
gramalingamCopilot
andauthored
fix: invoke NameFixPass after graph-modifying operations in rewriter, version-converter, and optimizer (#2922)
## Problem TapeBuilder creates values with explicit names that can clash with existing graph value names when nodes are inserted via replace_nodes_and_values. NameAuthority silently registers duplicate names without conflict checking, causing ValueError when the clashing value is an initializer. ## Fix Invoke NameFixPass conditionally (only when changes were actually made) at the end of the three graph-modifying callers: - RewriteRuleSet.apply_to_model -- runs when count > 0 - _VersionConverter.visit_model -- runs when self._modified is True (set in replace_node) - FoldConstantsPass.call -- runs when self._modified is True Each call site includes an explanatory comment and the _modified check is encapsulated within the class that owns it. ## Tests Added name-clash regression tests to all three test suites: - NameClashAfterRewriteTest in onnxscript/rewriter/pattern_test.py - NameClashAfterConversionTest in onnxscript/version_converter/_version_converter_test.py - NameClashAfterFoldTest in onnxscript/optimizer/_constant_folding_test.py Each test triggers the operation, injects a duplicate value name to simulate what TapeBuilder can produce via NameAuthority, and asserts all value names are unique after the operation. ## Known limitation There is a pre-existing terminal NameFixPass in the optimizer pipeline (_optimizer.py), so optimize_ir may call NameFixPass up to 5 times on a 2-iteration run. This is acceptable overhead and is tracked as a follow-up to investigate deduplication. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 18280e0 commit c72e810

6 files changed

Lines changed: 202 additions & 0 deletions

File tree

onnxscript/optimizer/_constant_folding.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import onnx
2323
import onnx.reference.ops
2424
import onnx_ir as ir
25+
import onnx_ir.passes.common as ir_passes_common
2526

2627
import onnxscript.utils.utils as utils
2728
from onnxscript._internal.tape_builder import BuilderBase, TapeBuilder
@@ -1394,6 +1395,11 @@ def call(self, model: ir.Model) -> FoldConstantsResult:
13941395
for function in model.functions.values():
13951396
# TODO(rama): Should we specialize functions?
13961397
self.visit_function(function)
1398+
if self._modified:
1399+
# TapeBuilder may create values with names that clash with existing graph
1400+
# values when nodes are inserted via replace_nodes_and_values.
1401+
# NameFixPass ensures all value names are unique before returning.
1402+
ir_passes_common.NameFixPass()(model)
13971403
return FoldConstantsResult(model, self._modified, self._state.symbolic_value_map)
13981404

13991405

onnxscript/optimizer/_constant_folding_test.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -797,5 +797,61 @@ def test_initializer_as_graph_output_is_not_removed(self):
797797
self.assertIn("z", output_names)
798798

799799

800+
def _all_value_names_unique(model: ir.Model) -> bool:
801+
"""Return True if all named values in the top-level graph have unique names."""
802+
names = []
803+
for v in model.graph.inputs:
804+
if v.name:
805+
names.append(v.name)
806+
for v in model.graph.initializers.values():
807+
if v.name:
808+
names.append(v.name)
809+
for node in model.graph:
810+
for output in node.outputs:
811+
if output.name:
812+
names.append(output.name)
813+
return len(names) == len(set(names))
814+
815+
816+
class NameClashAfterFoldTest(unittest.TestCase):
817+
"""Tests that fold_constants calls NameFixPass to deduplicate value names.
818+
819+
TapeBuilder may assign names that collide with existing graph values when
820+
new nodes are inserted via replace_nodes_and_values. NameFixPass, invoked
821+
by FoldConstantsPass.call when the model was modified, resolves the
822+
duplicates.
823+
"""
824+
825+
def test_fold_constants_deduplicates_names(self):
826+
"""Duplicate value names present alongside a constant-fold are fixed."""
827+
model = ir.from_onnx_text(
828+
"""
829+
<ir_version: 7, opset_import: [ "" : 17]>
830+
agraph (float[N] x) => (float[N] z) {
831+
two = Constant <value_float=2.0> ()
832+
four = Add(two, two)
833+
extra = Relu(x)
834+
z = Mul(extra, four)
835+
}
836+
"""
837+
)
838+
839+
# Simulate the name clash that TapeBuilder can introduce: 'extra' (a
840+
# non-folded node that survives) is given the same name as 'four' (the
841+
# folded Add output) because NameAuthority does not check for conflicts
842+
# when registering pre-named values inserted by TapeBuilder.
843+
four_node = next(n for n in model.graph if n.op_type == "Add")
844+
extra_node = next(n for n in model.graph if n.op_type == "Relu")
845+
extra_node.outputs[0].name = four_node.outputs[0].name # inject clash
846+
847+
result = _constant_folding.fold_constants(model)
848+
849+
self.assertTrue(result.modified, "Folding must have modified the model")
850+
self.assertTrue(
851+
_all_value_names_unique(model),
852+
"All value names must be unique after fold_constants",
853+
)
854+
855+
800856
if __name__ == "__main__":
801857
unittest.main()

onnxscript/rewriter/_rewrite_rule.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
TypeVar,
1414
)
1515

16+
import onnx_ir.passes.common as ir_passes_common
17+
1618
import onnxscript.optimizer
1719
import onnxscript.rewriter._basics as _basics
1820
import onnxscript.rewriter._context as _context
@@ -835,6 +837,11 @@ def apply_to_model(
835837
)
836838
if self.remove_unused_nodes:
837839
onnxscript.optimizer.remove_unused_nodes(model)
840+
if count > 0:
841+
# TapeBuilder may create values with names that clash with existing graph
842+
# values when nodes are inserted via replace_nodes_and_values.
843+
# NameFixPass ensures all value names are unique before returning.
844+
ir_passes_common.NameFixPass()(model)
838845
return count
839846

840847
def __iter__(self):

onnxscript/rewriter/pattern_test.py

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -989,5 +989,70 @@ def test_pattern_builder_context(self):
989989
self.assertEqual(ops, ["Op1", "Op2", "Add", "Op3", "Mul"])
990990

991991

992+
def _all_value_names_unique(model: ir.Model) -> bool:
993+
"""Return True if all named values in the top-level graph have unique names."""
994+
names = []
995+
for v in model.graph.inputs:
996+
if v.name:
997+
names.append(v.name)
998+
for v in model.graph.initializers.values():
999+
if v.name:
1000+
names.append(v.name)
1001+
for node in model.graph:
1002+
for output in node.outputs:
1003+
if output.name:
1004+
names.append(output.name)
1005+
return len(names) == len(set(names))
1006+
1007+
1008+
class NameClashAfterRewriteTest(unittest.TestCase):
1009+
"""Tests that apply_to_model calls NameFixPass to deduplicate value names.
1010+
1011+
TapeBuilder may assign names that collide with existing graph values when
1012+
new nodes are inserted via replace_nodes_and_values. NameFixPass, invoked
1013+
by apply_to_model when at least one rewrite fires, resolves the duplicates.
1014+
"""
1015+
1016+
def test_apply_to_model_deduplicates_names(self):
1017+
"""Duplicate value names introduced alongside a rewrite are fixed."""
1018+
model_proto = onnx.parser.parse_model(
1019+
"""
1020+
<ir_version: 7, opset_import: [ "" : 17]>
1021+
agraph (float[N] x, float[N] y, float[N] p) => (float[N] z)
1022+
{
1023+
c1 = Constant<value_float = 1.0>()
1024+
t1 = Div(c1, x)
1025+
z1 = Mul(t1, y)
1026+
extra = Add(z1, p)
1027+
z = Identity(extra)
1028+
}
1029+
"""
1030+
)
1031+
model = ir.serde.deserialize_model(model_proto)
1032+
1033+
# Simulate the name clash that TapeBuilder can introduce: two values that
1034+
# survive the rewrite (not part of the matched pattern) end up sharing a
1035+
# name because NameAuthority does not check for conflicts when registering
1036+
# pre-named values from TapeBuilder.
1037+
extra_node = next(n for n in model.graph if n.op_type == "Add")
1038+
identity_node = next(n for n in model.graph if n.op_type == "Identity")
1039+
identity_node.outputs[0].name = extra_node.outputs[0].name # inject clash
1040+
1041+
def reciprocal_mul_pattern(op, x, y):
1042+
return (1 / x) * y
1043+
1044+
def div_replacement(op, x, y):
1045+
return op.Div(y, x)
1046+
1047+
rule = pattern.RewriteRule(reciprocal_mul_pattern, div_replacement)
1048+
count = rule.apply_to_model(model)
1049+
1050+
self.assertGreater(count, 0, "Rewrite rule must have fired to exercise the fix")
1051+
self.assertTrue(
1052+
_all_value_names_unique(model),
1053+
"All value names must be unique after apply_to_model",
1054+
)
1055+
1056+
9921057
if __name__ == "__main__":
9931058
unittest.main()

onnxscript/version_converter/_version_converter.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from typing import Callable, Sequence, Union
1111

1212
import onnx_ir.convenience as ir_convenience
13+
import onnx_ir.passes.common as ir_passes_common
1314

1415
import onnxscript.utils.metadata_merger as metadata_merger
1516
from onnxscript import ir
@@ -239,6 +240,7 @@ def groupnormalization_20_21(node: ir.Node, op):
239240
class _VersionConverter:
240241
def __init__(self, target_version: int):
241242
self._target_version = target_version
243+
self._modified: bool = False
242244
# Default metadata merger: no merging should be needed; keep the first value.
243245
self._default_metadata_merger: metadata_merger.MetadataMerger = (
244246
metadata_merger.MetadataMerger(
@@ -269,6 +271,7 @@ def replace_node(self, node: ir.Node, replacement, root: ir.Graph | ir.Function)
269271
ir_convenience.replace_nodes_and_values(
270272
root, node, [node], replacement.new_nodes, node.outputs, replacement.new_outputs
271273
)
274+
self._modified = True
272275

273276
def visit_attribute(self, attr: ir.Attr) -> None:
274277
if attr.is_ref():
@@ -341,6 +344,11 @@ def visit_model(self, model: ir.Model) -> None:
341344
self.visit_graph_or_function(function)
342345
_set_onnx_opset_version(function, self._target_version)
343346
_set_onnx_opset_version(model, self._target_version)
347+
if self._modified:
348+
# TapeBuilder may create values with names that clash with existing graph
349+
# values when nodes are inserted via replace_nodes_and_values.
350+
# NameFixPass ensures all value names are unique before returning.
351+
ir_passes_common.NameFixPass()(model)
344352

345353

346354
def convert_version(model: ir.Model, target_version: int) -> None:

onnxscript/version_converter/_version_converter_test.py

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -538,5 +538,65 @@ def test_version_convert_compatible(self):
538538
version_converter.convert_version(model, target_version=target_version)
539539

540540

541+
def _all_value_names_unique(model: ir.Model) -> bool:
542+
"""Return True if all named values in the top-level graph have unique names."""
543+
names = []
544+
for v in model.graph.inputs:
545+
if v.name:
546+
names.append(v.name)
547+
for v in model.graph.initializers.values():
548+
if v.name:
549+
names.append(v.name)
550+
for node in model.graph:
551+
for output in node.outputs:
552+
if output.name:
553+
names.append(output.name)
554+
return len(names) == len(set(names))
555+
556+
557+
class NameClashAfterConversionTest(unittest.TestCase):
558+
"""Tests that convert_version calls NameFixPass to deduplicate value names.
559+
560+
TapeBuilder may assign names that collide with existing graph values when
561+
new nodes are inserted via replace_nodes_and_values. NameFixPass, invoked
562+
inside _VersionConverter.visit_model when nodes were modified, resolves
563+
the duplicates.
564+
"""
565+
566+
def test_convert_version_deduplicates_names(self):
567+
"""Duplicate value names present after conversion are fixed by NameFixPass."""
568+
model = ir.from_onnx_text(
569+
"""
570+
<ir_version: 7, opset_import: [ "" : 18]>
571+
agraph (float[4, 512, 512] input_x, float[4, 1024, 1024] input_y) => (float[4, 1024, 1024] output)
572+
{
573+
shape_a = Constant<value: tensor = int64[5] {1, 4, 512, 512}>()
574+
reshape_x = Reshape (input_x, shape_a)
575+
shape_b = Constant<value: tensor = int64[5] {1, 4, 1024, 1024}>()
576+
reshape_y = Reshape (input_x, shape_b)
577+
gridsample = GridSample <mode = "bilinear"> (reshape_x, reshape_y)
578+
shape_c = Constant<value: tensor = int64[4] {4, 1024, 1024}>()
579+
output = Reshape (gridsample, shape_c)
580+
}
581+
"""
582+
)
583+
584+
# Simulate the name clash that TapeBuilder can introduce: two Constant
585+
# node outputs (not touched by the GridSample adapter) receive the same
586+
# name because NameAuthority does not check for conflicts when registering
587+
# pre-named values inserted by TapeBuilder.
588+
shape_a_output = model.graph.node(0).outputs[0]
589+
shape_c_output = model.graph.node(5).outputs[0]
590+
shape_c_output.name = shape_a_output.name # inject clash
591+
592+
version_converter.convert_version(model, target_version=20)
593+
594+
self.assertEqual(model.opset_imports[""], 20)
595+
self.assertTrue(
596+
_all_value_names_unique(model),
597+
"All value names must be unique after convert_version",
598+
)
599+
600+
541601
if __name__ == "__main__":
542602
unittest.main()

0 commit comments

Comments
 (0)