Fix CSE across operator domains

Author: andife

onnxoptimizer/passes/cse_util.h

@@ -192,6 +192,7 @@ struct CSENodeHash {
     auto string_hasher = std::hash<std::string>();
     auto sym_hasher = std::hash<Symbol>();
     hash_combine(seed, std::hash<uint32_t>(), static_cast<uint32_t>(n->kind()),
+                 string_hasher, n->domain(), string_hasher, n->overload(),
                  size_t_hasher, inputs.size());
     for (const auto& input : inputs) {
       hash_combine(seed, string_hasher, input->uniqueName());
@@ -255,7 +256,9 @@ struct CSEEqual {
     SymbolCompare cmp;
     std::sort(attr_names_l.begin(), attr_names_l.end(), cmp);
     std::sort(attr_names_r.begin(), attr_names_r.end(), cmp);
-    if (lhs->kind() != rhs->kind() || inputs_l.size() != inputs_r.size() ||
+    if (lhs->kind() != rhs->kind() || lhs->domain() != rhs->domain() ||
+        lhs->overload() != rhs->overload() ||
+        inputs_l.size() != inputs_r.size() ||
         outputs_l.size() != outputs_r.size() || attr_names_l != attr_names_r) {
       return false;
     }

onnxoptimizer/test/optimizer_test.py

@@ -5017,6 +5017,42 @@ def test_eliminate_common_subexression_with_attribute1(self):  # type: () -> Non
 
         assert len(optimized_model.graph.node) == 6
 
+    def test_eliminate_common_subexpression_preserves_operator_domain(
+        self,
+    ):  # type: () -> None
+        X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1])
+        Y1 = helper.make_tensor_value_info("Y1", TensorProto.FLOAT, [1])
+        Y2 = helper.make_tensor_value_info("Y2", TensorProto.FLOAT, [1])
+        standard_add = helper.make_node("Add", ["X", "X"], ["Y1"])
+        local_add = helper.make_node("Add", ["X", "X"], ["Y2"], domain="local")
+        graph = helper.make_graph(
+            [standard_add, local_add], "test", [X], [Y1, Y2]
+        )
+        local_function = helper.make_function(
+            "local",
+            "Add",
+            ["A", "B"],
+            ["C"],
+            [helper.make_node("Sub", ["A", "B"], ["C"])],
+            [helper.make_opsetid("", LATEST_STABLE_OPSET_VERSION)],
+        )
+        model = helper.make_model(
+            graph,
+            functions=[local_function],
+            opset_imports=[
+                helper.make_opsetid("", LATEST_STABLE_OPSET_VERSION),
+                helper.make_opsetid("local", 1),
+            ],
+        )
+
+        optimized_model = onnxoptimizer.optimize(
+            model, ["eliminate_common_subexpression", "eliminate_deadend"]
+        )
+
+        assert len(optimized_model.graph.node) == 2
+        assert optimized_model.graph.node[0].domain == ""
+        assert optimized_model.graph.node[1].domain == "local"
+
     def _test_fuse_qkv_with_opset(self, opset_version):  # type: (int) -> None
         X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4096, 320])
         Y1 = helper.make_tensor_value_info("Y1", TensorProto.FLOAT, [1, 4096, 8, 40])