[GVN] Propagate simple equalities from assumes within the tail of the block

[reames]

Summary

This extends the existing logic for propagating constant expressions in an analogous manner for what we do across basic blocks. The core point is that we chose some order of operands, and canonicalize uses towards that one.

The heuristic used is inspired by the one used across blocks; in a follow up change, I'd like to common them so that the cross block version uses the slightly stronger ordering herein. If preferred, I can weaken this one to match the cross block, then strengthen both in a single patch.

As noted by the TODOs in the code, there's a good amount of room for improving the existing code and making it more powerful. Unless there's strong objection, I'd like to land this first, then iterate. I am willing to commit to the findLeader change, and *maybe* the general restructure of propagateEqualities (depending on how much that snowballs - the tests are *super fragile*).

[lebedev.ri]

lib/Transforms/Scalar/GVN.cpp

1448

s/const/???/

[reames]

Remove copied comment which doesn't make sense in new context

[fhahn]

LGTM, with a few nits/suggestions for comments. Also, it would be great if you could also add a test case with floats.

lib/Transforms/Scalar/GVN.cpp

1390

could just be something like the snippet below, which might be short & explicit enough to have inline.

any_of(I.users(), [B](User *U) {
return isa<Instruction>(U) &&  cast<Instruction>(U)->getParent() == B;
});

1438–1442

We only replace the uses that come after in the current BB, right? I think it would be better to say that, instead of talking about replacing all dominated uses

1439

It might be worth mentioning what we mean by oldest here (seen earlier during value numbering)

1450

Shouldn't that be %0? Did I miss where %v is coming from in the example?

1458–1465

Typo Hueristically

1479

Could you add a test case for this case?

1482

Again I think mentioning that we only replacing uses in the current BB would be slightly clearer.

[reames]

In D66977#1653969, @fhahn wrote:

LGTM, with a few nits/suggestions for comments. Also, it would be great if you could also add a test case with floats.

Landed w/comments applied and one additional fix noticed when writing your requested floating point test case.

It turned out the code I was extending had a bug where it didn't account for the fact that +0.0 compares equal to -0.0. I went ahead and added a bailout since the none of the existing tests were effected.

The extra bit seemed obvious. If anyone thinks it warrants re-review, let me know and I'll revert.

include/llvm/Transforms/Scalar/GVN.h

struct LeaderTableEntry {
     LeaderTableEntry *Next;
   };
   DenseMap<uint32_t, LeaderTableEntry> LeaderTable;
   BumpPtrAllocator TableAllocator;
 
   // Block-local map of equivalent values to their leader, does not
   // propagate to any successors. Entries added mid-block are applied
   // to the remaining instructions in the block.
-  SmallMapVector<Value *, Constant *, 4> ReplaceWithConstMap;
+  SmallMapVector<Value *, Value *, 4> ReplaceOperandsWithMap;
   SmallVector<Instruction *, 8> InstrsToErase;
 
   // Map the block to reversed postorder traversal number. It is used to
   // find back edge easily.
   DenseMap<AssertingVH<BasicBlock>, uint32_t> BlockRPONumber;
 
   // This is set 'true' initially and also when new blocks have been added to
   // the function being analyzed. This boolean is used to control the updating
private:
   bool performScalarPREInsertion(Instruction *Instr, BasicBlock *Pred,
                                  BasicBlock *Curr, unsigned int ValNo);
   Value *findLeader(const BasicBlock *BB, uint32_t num);
   void cleanupGlobalSets();
   void fillImplicitControlFlowInfo(BasicBlock *BB);
   void verifyRemoved(const Instruction *I) const;
   bool splitCriticalEdges();
   BasicBlock *splitCriticalEdges(BasicBlock *Pred, BasicBlock *Succ);
-  bool replaceOperandsWithConsts(Instruction *I) const;
+  bool replaceOperandsForInBlockEquality(Instruction *I) const;
   bool propagateEquality(Value *LHS, Value *RHS, const BasicBlockEdge &Root,
                          bool DominatesByEdge);
   bool processFoldableCondBr(BranchInst *BI);
   void addDeadBlock(BasicBlock *BB);
   void assignValNumForDeadCode();
   void assignBlockRPONumber(Function &F);
 };
 

lib/Transforms/Scalar/GVN.cpp

bool GVN::processNonLocalLoad(LoadInst *LI) {
 
   // Step 4: Eliminate partial redundancy.
   if (!EnablePRE || !EnableLoadPRE)
     return false;
 
   return PerformLoadPRE(LI, ValuesPerBlock, UnavailableBlocks);
 }
 
+static bool hasUsersIn(Value *V, BasicBlock *BB) {
+  for (User *U : V->users())
+    if (isa<Instruction>(U) &&
+        cast<Instruction>(U)->getParent() == BB)
+      return true;
+  return false;
+}
+
 bool GVN::processAssumeIntrinsic(IntrinsicInst *IntrinsicI) {
   assert(IntrinsicI->getIntrinsicID() == Intrinsic::assume &&
          "This function can only be called with llvm.assume intrinsic");
   Value *V = IntrinsicI->getArgOperand(0);
 
   if (ConstantInt *Cond = dyn_cast<ConstantInt>(V)) {
     if (Cond->isZero()) {
       Type *Int8Ty = Type::getInt8Ty(V->getContext());
for (BasicBlock *Successor : successors(IntrinsicI->getParent())) {
     // This property is only true in dominated successors, propagateEquality
     // will check dominance for us.
     Changed |= propagateEquality(V, True, Edge, false);
   }
 
   // We can replace assume value with true, which covers cases like this:
   // call void @llvm.assume(i1 %cmp)
   // br i1 %cmp, label %bb1, label %bb2 ; will change %cmp to true
-  ReplaceWithConstMap[V] = True;
+  ReplaceOperandsWithMap[V] = True;
 
-  // If one of *cmp *eq operand is const, adding it to map will cover this:
+  // If we find an equality fact, canonicalize all dominated uses to one
+  // of the two values.  We heuristically choice the "oldest" of the two.
+  // 
+  // Key case to cover are:
+  // 1) 
   // %cmp = fcmp oeq float 3.000000e+00, %0 ; const on lhs could happen
   // call void @llvm.assume(i1 %cmp)
   // ret float %0 ; will change it to ret float 3.000000e+00
+  // 2)
+  // %load = load float, float* %addr
+  // %cmp = fcmp oeq float %load, %0 ; const on lhs could happen
+  // call void @llvm.assume(i1 %cmp)
+  // ret float %load ; will change it to ret float %v
   if (auto *CmpI = dyn_cast<CmpInst>(V)) {
     if (CmpI->getPredicate() == CmpInst::Predicate::ICMP_EQ ||
         CmpI->getPredicate() == CmpInst::Predicate::FCMP_OEQ ||
         (CmpI->getPredicate() == CmpInst::Predicate::FCMP_UEQ &&
          CmpI->getFastMathFlags().noNaNs())) {
       Value *CmpLHS = CmpI->getOperand(0);
       Value *CmpRHS = CmpI->getOperand(1);
-      if (isa<Constant>(CmpLHS))
+      // Hueristically pick the better replacement -- the choice of heuristic
+      // isn't terribly important here, but the fact we canonicalize on some
+      // replacement is for exposing other simplifications.
+      // TODO: pull this out as a helper function and reuse w/existing
+      // (slightly different) logic.
+      if (isa<Constant>(CmpLHS) && !isa<Constant>(CmpRHS))
+        std::swap(CmpLHS, CmpRHS);
+      if (!isa<Instruction>(CmpLHS) && isa<Instruction>(CmpRHS))
         std::swap(CmpLHS, CmpRHS);
-      auto *RHSConst = dyn_cast<Constant>(CmpRHS);
+      if ((isa<Argument>(CmpLHS) && isa<Argument>(CmpRHS)) ||
+          (isa<Instruction>(CmpLHS) && isa<Instruction>(CmpRHS))) {
+        // Move the 'oldest' value to the right-hand side, using the value
+        // number as a proxy for age.
+        uint32_t LVN = VN.lookupOrAdd(CmpLHS);
+        uint32_t RVN = VN.lookupOrAdd(CmpRHS);
+        if (LVN < RVN)
+          std::swap(CmpLHS, CmpRHS);
+      }
+
+      // Handle degenerate case where we either haven't pruned a dead path or a
+      // removed a trivial assume yet.
+      if (isa<Constant>(CmpLHS) && isa<Constant>(CmpRHS))
+        return Changed;
+
+      LLVM_DEBUG(dbgs() << "Replacing dominated uses of "
+                 << *CmpLHS << " with "
+                 << *CmpRHS << "\n");
+      
 
-      // If only one operand is constant.
-      if (RHSConst != nullptr && !isa<Constant>(CmpLHS))
-        ReplaceWithConstMap[CmpLHS] = RHSConst;
+      // Setup the replacement map - this handles uses within the same block
+      if (hasUsersIn(CmpLHS, IntrinsicI->getParent()))
+        ReplaceOperandsWithMap[CmpLHS] = CmpRHS;
+
+      // NOTE: The non-block local cases are handled by the call to
+      // propagateEquality above; this block is just about handling the block
+      // local cases.  TODO: There's a bunch of logic in propagateEqualiy which
+      // isn't duplicated for the block local case, can we share it somehow?
     }
   }
   return Changed;
 }
 
 static void patchAndReplaceAllUsesWith(Instruction *I, Value *Repl) {
   patchReplacementInstruction(I, Repl);
   I->replaceAllUsesWith(Repl);
void GVN::assignBlockRPONumber(Function &F) {
   BlockRPONumber.clear();
   uint32_t NextBlockNumber = 1;
   ReversePostOrderTraversal<Function *> RPOT(&F);
   for (BasicBlock *BB : RPOT)
     BlockRPONumber[BB] = NextBlockNumber++;
   InvalidBlockRPONumbers = false;
 }
 
-// Tries to replace instruction with const, using information from
-// ReplaceWithConstMap.
-bool GVN::replaceOperandsWithConsts(Instruction *Instr) const {
+bool GVN::replaceOperandsForInBlockEquality(Instruction *Instr) const {
+  // TODO: We can remove the separate ReplaceOperandsWithMap data structure in
+  // favor of putting equalitys into the leader table and using findLeader
+  // here. 
   bool Changed = false;
   for (unsigned OpNum = 0; OpNum < Instr->getNumOperands(); ++OpNum) {
-    Value *Operand = Instr->getOperand(OpNum);
-    auto it = ReplaceWithConstMap.find(Operand);
-    if (it != ReplaceWithConstMap.end()) {
-      assert(!isa<Constant>(Operand) &&
-             "Replacing constants with constants is invalid");
+    Value *Operand = Instr->getOperand(OpNum); 
+    auto it = ReplaceOperandsWithMap.find(Operand);
+    if (it != ReplaceOperandsWithMap.end()) {
       LLVM_DEBUG(dbgs() << "GVN replacing: " << *Operand << " with "
                         << *it->second << " in instruction " << *Instr << '\n');
       Instr->setOperand(OpNum, it->second);
       Changed = true;
     }
   }
   return Changed;
 }
bool GVN::processBlock(BasicBlock *BB) {
   // FIXME: Kill off InstrsToErase by doing erasing eagerly in a helper function
   // (and incrementing BI before processing an instruction).
   assert(InstrsToErase.empty() &&
          "We expect InstrsToErase to be empty across iterations");
   if (DeadBlocks.count(BB))
     return false;
 
   // Clearing map before every BB because it can be used only for single BB.
-  ReplaceWithConstMap.clear();
+  ReplaceOperandsWithMap.clear();
   bool ChangedFunction = false;
 
   for (BasicBlock::iterator BI = BB->begin(), BE = BB->end();
        BI != BE;) {
-    if (!ReplaceWithConstMap.empty())
-      ChangedFunction |= replaceOperandsWithConsts(&*BI);
+    if (!ReplaceOperandsWithMap.empty())
+      ChangedFunction |= replaceOperandsForInBlockEquality(&*BI);
     ChangedFunction |= processInstruction(&*BI);
 
     if (InstrsToErase.empty()) {
       ++BI;
       continue;
     }
 
     // If we need some instructions deleted, do it now.

test/Transforms/GVN/equality-assume.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -gvn < %s | FileCheck %s
 
 define i32 @test(i32* %p, i32 %v) {
 ; CHECK-LABEL: @test(
 ; CHECK-NEXT:    [[LOAD:%.*]] = load i32, i32* [[P:%.*]]
 ; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[LOAD]], [[V:%.*]]
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[C]])
-; CHECK-NEXT:    ret i32 [[LOAD]]
+; CHECK-NEXT:    ret i32 [[V]]
 ;
   %load = load i32, i32* %p
   %c = icmp eq i32 %load, %v
   call void @llvm.assume(i1 %c)
   ret i32 %load
 }
 
 define i32 @reverse(i32* %p, i32 %v) {
;
   ret i32 %v
 }
 
 define i32 @test2(i32* %p, i32 %v) {
 ; CHECK-LABEL: @test2(
 ; CHECK-NEXT:    [[LOAD:%.*]] = load i32, i32* [[P:%.*]]
 ; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[LOAD]], [[V:%.*]]
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[C]])
-; CHECK-NEXT:    ret i32 [[LOAD]]
+; CHECK-NEXT:    ret i32 [[V]]
 ;
   %load = load i32, i32* %p
   %c = icmp eq i32 %load, %v
   call void @llvm.assume(i1 %c)
   %load2 = load i32, i32* %p
   ret i32 %load2
 }