Avoid mutating stream partition op infos while testing. (iree-org#11465)

benvanik · web-flow · commit d52c5f327d4c · 2022-12-07T18:46:35.000-08:00
Fixes iree-org#11451.
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning/ReferencePartitioning.cpp b/compiler/src/iree/compiler/Dialect/Stream/Analysis/Partitioning/ReferencePartitioning.cpp
@@ -125,7 +125,9 @@ PartitionSet partitionStreamableOpsReference(
     // dependency chain down the use-def chain to a partition.
     llvm::BitVector consumers(builders.size(), /*t=*/false);
     for (auto user : op.getUsers()) {
-      auto &userInfo = opInfos[user];
+      auto userInfoIt = opInfos.find(user);
+      if (userInfoIt == opInfos.end()) continue;
+      auto &userInfo = userInfoIt->second;
       LLVM_DEBUG({
         llvm::dbgs() << "Testing user:\n";
         user->print(llvm::dbgs(), *asmState);
@@ -332,7 +334,9 @@ PartitionSet partitionRegionConcurrencyReference(
     // We prune the set based on whether the users are part of a transitive
     // dependency chain down the use-def chain to a wave.
     for (auto user : op.getUsers()) {
-      auto &userInfo = opInfos[user];
+      auto userInfoIt = opInfos.find(user);
+      if (userInfoIt == opInfos.end()) continue;
+      auto &userInfo = userInfoIt->second;
       LLVM_DEBUG({
         llvm::dbgs() << "Testing user:\n";
         user->print(llvm::dbgs(), *asmState);
