[ONNX] improve regex matching in onnx-importer name sanitization (#3955)

Instead of adding unsupported characters on a case-by-case basis, we should replace anything that isn't alphanumeric, `_`, or `.`.
llvm · Jan 14, 2025 · 62eb38b · 62eb38b
1 parent 11efcda
commit 62eb38b
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 1 deletion.
diff --git a/python/torch_mlir/extras/onnx_importer.py b/python/torch_mlir/extras/onnx_importer.py
@@ -742,7 +742,7 @@ def _sanitize_name(self, name):
 
         # Remove characters that are invalid in MLIR identifier names.
         # https://mlir.llvm.org/docs/LangRef/#identifiers-and-keywords
-        return re.sub("[:/-]", "_", name)
+        return re.sub("[^\w\.]", "_", name)
 
     def tensor_proto_to_attr(self, tp: onnx.TensorProto) -> Attribute:
         tensor_type = self.tensor_proto_to_builtin_type(tp)

diff --git a/test/python/onnx_importer/BadName.onnx b/test/python/onnx_importer/BadName.onnx
diff --git a/test/python/onnx_importer/BadName.runlit b/test/python/onnx_importer/BadName.runlit
@@ -0,0 +1,5 @@
+# The original constant name : "abz_.(1, 2)[$something, %anotherthing]"
+
+# RUN: %PYTHON -m torch_mlir.tools.import_onnx %S/BadName.onnx | FileCheck %s
+
+# CHECK: torch.operator "onnx.Constant"() {torch.onnx.value = dense_resource<_abz_._1__2___something___anotherthing_>