From f1fa3e23cef8c7a987646ecf0d33c3ce096cd4cc Mon Sep 17 00:00:00 2001 From: Ziliang Zhang Date: Tue, 28 Apr 2026 16:31:51 +0800 Subject: [PATCH 1/3] [onnx_importer] Disambiguate empty string: optional none vs tensor name The NodeImporter cached torch.constant.none under _nv_map[""], matching ONNX's convention that an empty string in Node.input denotes an omitted optional input. Some producers (e.g. Microsoft SkipSimplifiedLayerNormalization) also bind real intermediate results to outputs whose names are the empty string. Each such output overwrote _nv_map[""], so later nodes that use "" for omitted optionals (e.g. GroupQueryAttention's trailing inputs) incorrectly received those tensor SSA values instead of torch.constant.none. Behavior changes: - Cache the shared none value under _OPTIONAL_NONE_CACHE_KEY instead of "". - When resolving node inputs, treat input_name == "" as omitted optional: append get_none() and an empty onnx.TypeProto without indexing _nv_map[""]. - Register outputs named "" under unique keys __torch_mlir_onnx_importer_anon_ so multiple anonymous outputs do not overwrite each other. Adds test/python/onnx_importer/test_empty_string_optional_inputs.py: minimal Identity -> custom op graph where optional inputs are "" and must import as %none operands, not tensor values stored under "". Symptom fixed: GroupQueryAttention previously showed duplicated operands such as (%10#2, %10#2, %10#2) instead of (%none, %none, %none) for optional slots. --- python/torch_mlir/extras/onnx_importer.py | 27 ++++++-- .../test_empty_string_optional_inputs.py | 68 +++++++++++++++++++ 2 files changed, 91 insertions(+), 4 deletions(-) create mode 100644 test/python/onnx_importer/test_empty_string_optional_inputs.py diff --git a/python/torch_mlir/extras/onnx_importer.py b/python/torch_mlir/extras/onnx_importer.py index c6923e82df0c..df30c88a9170 100644 --- a/python/torch_mlir/extras/onnx_importer.py +++ b/python/torch_mlir/extras/onnx_importer.py @@ -77,6 +77,10 @@ func as func_dialect, ) +# Cache key for the shared torch.constant.none used for omitted ONNX optional inputs. +# Must not collide with ONNX tensor names (including the empty string). +_OPTIONAL_NONE_CACHE_KEY = "__torch_mlir_onnx_importer_optional_none__" + @dataclass class Config: @@ -239,6 +243,7 @@ class NodeImporter: "_p", "_b", "_nv_map", + "_anon_output_counter", ] def __init__( @@ -259,6 +264,7 @@ def __init__( self._p = parent_op self._b = block self._nv_map: Dict[str, Value] = {} + self._anon_output_counter = 0 @classmethod def define_function( @@ -366,8 +372,8 @@ def import_all(self, func=True): Operation.create(name="torch.operator_terminator", operands=outputs) def get_none(self): - if "" in self._nv_map: - return self._nv_map[""] + if _OPTIONAL_NONE_CACHE_KEY in self._nv_map: + return self._nv_map[_OPTIONAL_NONE_CACHE_KEY] with InsertionPoint(self._b), Location.name("onnx_importer.none"): nne = Operation.create( @@ -376,7 +382,7 @@ def get_none(self): operands=[], attributes={}, ).results[0] - self._nv_map[""] = nne + self._nv_map[_OPTIONAL_NONE_CACHE_KEY] = nne return nne def import_node(self, node: onnx.NodeProto): @@ -396,6 +402,12 @@ def import_node(self, node: onnx.NodeProto): input_values = [] input_type_protos = [] for input_name in node.input: + # ONNX uses the empty string for omitted optional inputs; it must not + # be confused with _nv_map[""], which may hold a real tensor named "". + if input_name == "": + input_values.append(self.get_none()) + input_type_protos.append(onnx.TypeProto()) + continue try: input_values.append(self._nv_map[input_name]) # Missing optional arguments will have empty types @@ -447,7 +459,14 @@ def import_node(self, node: onnx.NodeProto): self.import_regions(node.attribute, custom_op) for output_name, output_value in zip(output_names, custom_op.results): - self._nv_map[output_name] = output_value + if output_name == "": + key = ( + f"__torch_mlir_onnx_importer_anon_{self._anon_output_counter}" + ) + self._anon_output_counter += 1 + self._nv_map[key] = output_value + else: + self._nv_map[output_name] = output_value def import_attributes(self, onnx_attrs: List[onnx.AttributeProto]): attrs = {} diff --git a/test/python/onnx_importer/test_empty_string_optional_inputs.py b/test/python/onnx_importer/test_empty_string_optional_inputs.py new file mode 100644 index 000000000000..6f7680478001 --- /dev/null +++ b/test/python/onnx_importer/test_empty_string_optional_inputs.py @@ -0,0 +1,68 @@ +# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +# RUN: %PYTHON %s + +"""Regression for NodeImporter: ONNX input name '' means omitted optional. + +The importer must not conflate that with _nv_map[\"\"] when an earlier node binds +a real tensor to the empty-string output name (see onnx_importer empty-string +collision fix). +""" + +import unittest + +import onnx +from onnx import TensorProto, helper + +from _torch_mlir_config import configure_context, ir, onnx_importer + + +def _minimal_collision_model() -> onnx.ModelProto: + """Identity writes to output \"\"; second node lists '', '' as omitted inputs.""" + inp = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 2]) + out = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 2]) + n1 = helper.make_node("Identity", ["x"], [""]) + n2 = helper.make_node( + "ReproEmptyStringCollision", + ["x", "", ""], + ["y"], + domain="zmc.repro", + ) + graph = helper.make_graph([n1, n2], "g", [inp], [out]) + return helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 21), + helper.make_opsetid("zmc.repro", 1), + ], + ) + + +class EmptyStringOptionalInputsTest(unittest.TestCase): + def test_optional_slots_use_constant_none_not_prior_tensor(self): + model = _minimal_collision_model() + ctx = ir.Context() + configure_context(ctx) + mi = onnx_importer.ModelInfo(model) + m = mi.create_module(context=ctx).operation + onnx_importer.NodeImporter.define_function(mi.main_graph, m).import_all() + asm = m.get_asm() + lines = [ln.strip() for ln in asm.splitlines() if "ReproEmptyStringCollision" in ln] + self.assertEqual( + len(lines), + 1, + msg="expected exactly one onnx.ReproEmptyStringCollision operator line", + ) + line = lines[0] + # Correct: trailing optionals are torch.constant.none uses (printed as %none, %none). + self.assertGreaterEqual( + line.count("%none"), + 2, + msg=f"expected at least two %none operands for omitted inputs, got:\n{line}", + ) + + +if __name__ == "__main__": + unittest.main() From 4cfcec4c66880d51e4b33c004d71b12b1b54e807 Mon Sep 17 00:00:00 2001 From: Ziliang Zhang Date: Thu, 7 May 2026 17:34:22 +0800 Subject: [PATCH 2/3] fix --- .../onnx_importer/test_empty_string_optional_inputs.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/test/python/onnx_importer/test_empty_string_optional_inputs.py b/test/python/onnx_importer/test_empty_string_optional_inputs.py index 6f7680478001..c401526805a3 100644 --- a/test/python/onnx_importer/test_empty_string_optional_inputs.py +++ b/test/python/onnx_importer/test_empty_string_optional_inputs.py @@ -6,7 +6,7 @@ """Regression for NodeImporter: ONNX input name '' means omitted optional. -The importer must not conflate that with _nv_map[\"\"] when an earlier node binds +The importer must not conflate that with _nv_map[""] when an earlier node binds a real tensor to the empty-string output name (see onnx_importer empty-string collision fix). """ @@ -20,7 +20,7 @@ def _minimal_collision_model() -> onnx.ModelProto: - """Identity writes to output \"\"; second node lists '', '' as omitted inputs.""" + """Identity writes to output ""; second node lists '', '' as omitted inputs.""" inp = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 2]) out = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 2]) n1 = helper.make_node("Identity", ["x"], [""]) @@ -49,7 +49,9 @@ def test_optional_slots_use_constant_none_not_prior_tensor(self): m = mi.create_module(context=ctx).operation onnx_importer.NodeImporter.define_function(mi.main_graph, m).import_all() asm = m.get_asm() - lines = [ln.strip() for ln in asm.splitlines() if "ReproEmptyStringCollision" in ln] + lines = [ + ln.strip() for ln in asm.splitlines() if "ReproEmptyStringCollision" in ln + ] self.assertEqual( len(lines), 1, From 425c0aa22927b805b6cabb8dd49d2cf31c4e7616 Mon Sep 17 00:00:00 2001 From: Ziliang Zhang Date: Thu, 7 May 2026 17:51:51 +0800 Subject: [PATCH 3/3] fix --- python/torch_mlir/extras/onnx_importer.py | 22 ++++++---------------- 1 file changed, 6 insertions(+), 16 deletions(-) diff --git a/python/torch_mlir/extras/onnx_importer.py b/python/torch_mlir/extras/onnx_importer.py index df30c88a9170..a4f95c3ee778 100644 --- a/python/torch_mlir/extras/onnx_importer.py +++ b/python/torch_mlir/extras/onnx_importer.py @@ -77,10 +77,6 @@ func as func_dialect, ) -# Cache key for the shared torch.constant.none used for omitted ONNX optional inputs. -# Must not collide with ONNX tensor names (including the empty string). -_OPTIONAL_NONE_CACHE_KEY = "__torch_mlir_onnx_importer_optional_none__" - @dataclass class Config: @@ -243,7 +239,7 @@ class NodeImporter: "_p", "_b", "_nv_map", - "_anon_output_counter", + "_none_value", ] def __init__( @@ -264,7 +260,7 @@ def __init__( self._p = parent_op self._b = block self._nv_map: Dict[str, Value] = {} - self._anon_output_counter = 0 + self._none_value: Optional[Value] = None @classmethod def define_function( @@ -372,8 +368,8 @@ def import_all(self, func=True): Operation.create(name="torch.operator_terminator", operands=outputs) def get_none(self): - if _OPTIONAL_NONE_CACHE_KEY in self._nv_map: - return self._nv_map[_OPTIONAL_NONE_CACHE_KEY] + if self._none_value is not None: + return self._none_value with InsertionPoint(self._b), Location.name("onnx_importer.none"): nne = Operation.create( @@ -382,7 +378,7 @@ def get_none(self): operands=[], attributes={}, ).results[0] - self._nv_map[_OPTIONAL_NONE_CACHE_KEY] = nne + self._none_value = nne return nne def import_node(self, node: onnx.NodeProto): @@ -459,13 +455,7 @@ def import_node(self, node: onnx.NodeProto): self.import_regions(node.attribute, custom_op) for output_name, output_value in zip(output_names, custom_op.results): - if output_name == "": - key = ( - f"__torch_mlir_onnx_importer_anon_{self._anon_output_counter}" - ) - self._anon_output_counter += 1 - self._nv_map[key] = output_value - else: + if output_name != "": self._nv_map[output_name] = output_value def import_attributes(self, onnx_attrs: List[onnx.AttributeProto]):