experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fix the bug when two passes with the same priority cannot be registered with the MlirOptimizationPassRegistry.

Browse Source 
PiperOrigin-RevId: 353744082
Change-Id: Ibe9c35e5c23e5a2d10043bc4c3fa33f5e5fb0db1
This commit is contained in:
Roman Dzhabarov 2021-01-25 15:30:49 -08:00 committed by TensorFlower Gardener
parent 9729442a35
commit 6f3bc3a09e
2 changed files with 11 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/compiler/mlir/mlir_graph_optimization_pass.h
View File

@ -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();

9
tensorflow/compiler/mlir/mlir_graph_optimization_pass_test.cc
View File

@ -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