From 6f3bc3a09ea73e1b1b5c2413201064e30f00d27a Mon Sep 17 00:00:00 2001
From: Roman Dzhabarov <rdzhabarov@google.com>
Date: Mon, 25 Jan 2021 15:30:49 -0800
Subject: [PATCH] Fix the bug when two passes with the same priority cannot be
 registered with the MlirOptimizationPassRegistry.

PiperOrigin-RevId: 353744082
Change-Id: Ibe9c35e5c23e5a2d10043bc4c3fa33f5e5fb0db1
---
 tensorflow/compiler/mlir/mlir_graph_optimization_pass.h  | 4 ++--
 .../compiler/mlir/mlir_graph_optimization_pass_test.cc   | 9 +++++++++
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/mlir/mlir_graph_optimization_pass.h b/tensorflow/compiler/mlir/mlir_graph_optimization_pass.h
index b01f6b57e64..ec26dfec1d2 100644
--- a/tensorflow/compiler/mlir/mlir_graph_optimization_pass.h
+++ b/tensorflow/compiler/mlir/mlir_graph_optimization_pass.h
@@ -63,7 +63,7 @@ class MlirOptimizationPassRegistry {
     }
   };
 
-  using Passes = std::set<PassRegistration, PriorityComparator>;
+  using Passes = std::multiset<PassRegistration, PriorityComparator>;
 
   // Returns the global registry of MLIR optimization passes.
   static MlirOptimizationPassRegistry& Global();
@@ -145,7 +145,7 @@ class MlirV1CompatOptimizationPassRegistry {
     }
   };
 
-  using Passes = std::set<PassRegistration, PriorityComparator>;
+  using Passes = std::multiset<PassRegistration, PriorityComparator>;
 
   // Returns the global registry of MLIR optimization passes.
   static MlirV1CompatOptimizationPassRegistry& Global();
diff --git a/tensorflow/compiler/mlir/mlir_graph_optimization_pass_test.cc b/tensorflow/compiler/mlir/mlir_graph_optimization_pass_test.cc
index 74992f67532..371719cb319 100644
--- a/tensorflow/compiler/mlir/mlir_graph_optimization_pass_test.cc
+++ b/tensorflow/compiler/mlir/mlir_graph_optimization_pass_test.cc
@@ -118,4 +118,13 @@ TEST_F(MlirGraphOptimizationPassTest, OptimizationPassFailsShadow) {
 #endif
 }
 
+TEST(MlirOptimizationPassRegistry, RegisterPassesWithTheSamePriority) {
+  MlirOptimizationPassRegistry::Global().Add(
+      0, std::make_unique<NiceMock<MockMlirOptimizationPass>>());
+  MlirOptimizationPassRegistry::Global().Add(
+      0, std::make_unique<NiceMock<MockMlirOptimizationPass>>());
+
+  EXPECT_EQ(MlirOptimizationPassRegistry::Global().passes().size(), 2);
+}
+
 }  // namespace tensorflow