[NewGVN] Fix caching for OpIsSafeForPhiOfOps (#98340)

The caching mechanism for 'OpIsSafeForPhiOfOps' is unsound. An operand
is deemed unsafe for PhiOfOps if it depends on a phi that resides in the
same block as the Phi block, i.e., where we are performing the PhiOfOps.
This is to avoid having to materialize the translated subexpressions. To
avoid redundant code walking, a cache is used to store these results.
Note, however, that since the safety is specific to the Phi block, we
cannot, in general, use the cached results for other blocks.

This patch addresses this by having a cache per block instead of a
single one for the entire function. closes #63335

Author: ManuelJBrito

llvm/lib/Transforms/Scalar/NewGVN.cpp

@@ -529,7 +529,11 @@ class NewGVN {
   // IR.
   SmallPtrSet<const Instruction *, 8> PHINodeUses;
 
-  DenseMap<const Value *, bool> OpSafeForPHIOfOps;
+  // The cached results, in general, are only valid for the specific block where
+  // they were computed. The unsigned part of the key is a unique block
+  // identifier
+  DenseMap<std::pair<const Value *, unsigned>, bool> OpSafeForPHIOfOps;
+  unsigned CacheIdx;
 
   // Map a temporary instruction we created to a parent block.
   DenseMap<const Value *, BasicBlock *> TempToBlock;
@@ -2596,19 +2600,19 @@ bool NewGVN::OpIsSafeForPHIOfOps(Value *V, const BasicBlock *PHIBlock,
     if (!isa<Instruction>(I))
       continue;
 
-    auto OISIt = OpSafeForPHIOfOps.find(I);
+    auto OISIt = OpSafeForPHIOfOps.find({I, CacheIdx});
     if (OISIt != OpSafeForPHIOfOps.end())
       return OISIt->second;
 
     // Keep walking until we either dominate the phi block, or hit a phi, or run
     // out of things to check.
     if (DT->properlyDominates(getBlockForValue(I), PHIBlock)) {
-      OpSafeForPHIOfOps.insert({I, true});
+      OpSafeForPHIOfOps.insert({{I, CacheIdx}, true});
       continue;
     }
     // PHI in the same block.
     if (isa<PHINode>(I) && getBlockForValue(I) == PHIBlock) {
-      OpSafeForPHIOfOps.insert({I, false});
+      OpSafeForPHIOfOps.insert({{I, CacheIdx}, false});
       return false;
     }
 
@@ -2627,10 +2631,10 @@ bool NewGVN::OpIsSafeForPHIOfOps(Value *V, const BasicBlock *PHIBlock,
       if (!isa<Instruction>(Op))
         continue;
       // Stop now if we find an unsafe operand.
-      auto OISIt = OpSafeForPHIOfOps.find(OrigI);
+      auto OISIt = OpSafeForPHIOfOps.find({OrigI, CacheIdx});
       if (OISIt != OpSafeForPHIOfOps.end()) {
         if (!OISIt->second) {
-          OpSafeForPHIOfOps.insert({I, false});
+          OpSafeForPHIOfOps.insert({{I, CacheIdx}, false});
           return false;
         }
         continue;
@@ -2640,7 +2644,7 @@ bool NewGVN::OpIsSafeForPHIOfOps(Value *V, const BasicBlock *PHIBlock,
       Worklist.push_back(cast<Instruction>(Op));
     }
   }
-  OpSafeForPHIOfOps.insert({V, true});
+  OpSafeForPHIOfOps.insert({{V, CacheIdx}, true});
   return true;
 }
 
@@ -3293,6 +3297,7 @@ void NewGVN::verifyIterationSettled(Function &F) {
   TouchedInstructions.set();
   TouchedInstructions.reset(0);
   OpSafeForPHIOfOps.clear();
+  CacheIdx = 0;
   iterateTouchedInstructions();
   DenseSet<std::pair<const CongruenceClass *, const CongruenceClass *>>
       EqualClasses;
@@ -3396,6 +3401,8 @@ void NewGVN::iterateTouchedInstructions() {
                             << " because it is unreachable\n");
           continue;
         }
+        // Use the appropriate cache for "OpIsSafeForPHIOfOps".
+        CacheIdx = RPOOrdering.lookup(DT->getNode(CurrBlock)) - 1;
         updateProcessedCount(CurrBlock);
       }
       // Reset after processing (because we may mark ourselves as touched when
@@ -3475,6 +3482,8 @@ bool NewGVN::runGVN() {
   LLVM_DEBUG(dbgs() << "Block " << getBlockName(&F.getEntryBlock())
                     << " marked reachable\n");
   ReachableBlocks.insert(&F.getEntryBlock());
+  // Use index corresponding to entry block.
+  CacheIdx = 0;
 
   iterateTouchedInstructions();
   verifyMemoryCongruency();

llvm/test/Transforms/NewGVN/cache-safe-phiofops.ll

@@ -0,0 +1,100 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=newgvn < %s | FileCheck %s
+
+define i32 @unsafe(i1 %arg, i32 %arg1) {
+; CHECK-LABEL: define i32 @unsafe(
+; CHECK-SAME: i1 [[ARG:%.*]], i32 [[ARG1:%.*]]) {
+; CHECK-NEXT:  [[BB:.*:]]
+; CHECK-NEXT:    br label %[[BB4:.*]]
+; CHECK:       [[BB4]]:
+; CHECK-NEXT:    br i1 [[ARG]], label %[[BB6:.*]], label %[[BB5:.*]]
+; CHECK:       [[BB5]]:
+; CHECK-NEXT:    br label %[[BB6]]
+; CHECK:       [[BB6]]:
+; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ [[ARG1]], %[[BB5]] ], [ 0, %[[BB4]] ]
+; CHECK-NEXT:    [[SREM:%.*]] = srem i32 [[ARG1]], [[PHI]]
+; CHECK-NEXT:    store i32 [[SREM]], ptr null, align 4
+; CHECK-NEXT:    br i1 [[ARG]], label %[[BB8:.*]], label %[[BB7:.*]]
+; CHECK:       [[BB7]]:
+; CHECK-NEXT:    br i1 true, label %[[BB8]], label %[[BB5]]
+; CHECK:       [[BB8]]:
+; CHECK-NEXT:    [[PHIOFOPS:%.*]] = phi i32 [ [[ARG1]], %[[BB6]] ], [ 0, %[[BB7]] ]
+; CHECK-NEXT:    [[PHI9:%.*]] = phi i32 [ 0, %[[BB7]] ], [ 1, %[[BB6]] ]
+; CHECK-NEXT:    store i32 [[PHIOFOPS]], ptr null, align 4
+; CHECK-NEXT:    br label %[[BB4]]
+;
+bb:
+  br label %bb4
+
+bb4:                                              ; preds = %bb8, %bb
+  br i1 %arg, label %bb6, label %bb5
+
+bb5:                                              ; preds = %bb7, %bb4
+  br label %bb6
+
+bb6:                                              ; preds = %bb5, %bb4
+  %phi = phi i32 [ %arg1, %bb5 ], [ 0, %bb4 ]
+  %or = or i32 %phi, %arg1
+  %srem = srem i32 %or, %phi
+  store i32 %srem, ptr null, align 4
+  br i1 %arg, label %bb8, label %bb7
+
+bb7:                                              ; preds = %bb6
+  br i1 true, label %bb8, label %bb5
+
+bb8:                                              ; preds = %bb7, %bb6
+  %phi9 = phi i32 [ 0, %bb7 ], [ 1, %bb6 ]
+  %mul = mul i32 %phi9, %or
+  store i32 %mul, ptr null, align 4
+  br label %bb4
+}
+
+define i32 @unsafe_load(i1 %arg) {
+; CHECK-LABEL: define i32 @unsafe_load(
+; CHECK-SAME: i1 [[ARG:%.*]]) {
+; CHECK-NEXT:  [[BB:.*:]]
+; CHECK-NEXT:    br label %[[BB1:.*]]
+; CHECK:       [[BB1]]:
+; CHECK-NEXT:    br i1 [[ARG]], label %[[BB3:.*]], label %[[BB2:.*]]
+; CHECK:       [[BB2]]:
+; CHECK-NEXT:    br label %[[BB3]]
+; CHECK:       [[BB3]]:
+; CHECK-NEXT:    [[PHI4:%.*]] = phi i32 [ 1, %[[BB2]] ], [ 0, %[[BB1]] ]
+; CHECK-NEXT:    [[LOAD:%.*]] = load i32, ptr null, align 4
+; CHECK-NEXT:    [[SREM:%.*]] = srem i32 [[LOAD]], [[PHI4]]
+; CHECK-NEXT:    store i32 [[SREM]], ptr null, align 4
+; CHECK-NEXT:    br i1 [[ARG]], label %[[BB6:.*]], label %[[BB5:.*]]
+; CHECK:       [[BB5]]:
+; CHECK-NEXT:    br i1 true, label %[[BB6]], label %[[BB2]]
+; CHECK:       [[BB6]]:
+; CHECK-NEXT:    [[PHIOFOPS:%.*]] = phi i32 [ [[LOAD]], %[[BB3]] ], [ 0, %[[BB5]] ]
+; CHECK-NEXT:    [[PHI7:%.*]] = phi i32 [ 0, %[[BB5]] ], [ 1, %[[BB3]] ]
+; CHECK-NEXT:    store i32 [[PHIOFOPS]], ptr null, align 4
+; CHECK-NEXT:    br label %[[BB1]]
+;
+bb:
+  br label %bb1
+
+bb1:                                              ; preds = %bb6, %bb
+  br i1 %arg, label %bb3, label %bb2
+
+bb2:                                              ; preds = %bb5, %bb1
+  %phi = phi i32 [ 1, %bb1 ], [ 0, %bb5 ]
+  br label %bb3
+
+bb3:                                              ; preds = %bb2, %bb1
+  %phi4 = phi i32 [ %phi, %bb2 ], [ 0, %bb1 ]
+  %load = load i32, ptr null, align 4
+  %srem = srem i32 %load, %phi4
+  store i32 %srem, ptr null, align 4
+  br i1 %arg, label %bb6, label %bb5
+
+bb5:                                              ; preds = %bb3
+  br i1 true, label %bb6, label %bb2
+
+bb6:                                              ; preds = %bb5, %bb3
+  %phi7 = phi i32 [ 0, %bb5 ], [ 1, %bb3 ]
+  %mul = mul i32 %phi7, %load
+  store i32 %mul, ptr null, align 4
+  br label %bb1
+}