Improve collectEphemeralValues and use it in CodeGenPrepare

[skatkov]

Summary

Improve collectEphemeralValues to find ephemeral induction variable
and be based on wouldInstructionBeTriviallyDead utility function.

Also use this function in CodeGenPrepare to collect and erase the same
set of ephemeral values.

[skatkov]

llvm/lib/CodeGen/CodeGenPrepare.cpp

729

will be removed after https://reviews.llvm.org/D154601 landed.

[skatkov]

ping, any feedback?

[skatkov]

ping.

I wonder whether it will easier to review if I split this patch into three one:

1. Use wouldInstructionBeTriviallyDead instead of checking side-effect and terminator only to extend the scope of collecting ephemeral values
2. Use CodeMetrics utility in CodeGenPrepare
3. Extend the algorithm for collecting ephemeral values by checking cycles (induction variable case)

I'd like to make a progress on this, so I need feedback.

[skatkov]

ping.

[apilipenko]

In D154584#4508666, @skatkov wrote:

I wonder whether it will easier to review if I split this patch into three one:

1. Use wouldInstructionBeTriviallyDead instead of checking side-effect and terminator only to extend the scope of collecting ephemeral values
2. Use CodeMetrics utility in CodeGenPrepare
3. Extend the algorithm for collecting ephemeral values by checking cycles (induction variable case)

Each of these steps is a reasonable improvement to what we have today. Given these improvements are independent (e.g. we would benefit from the more accurate collection of ephemeral values even without the CGP piece), I think it makes sense to split up and review them separately.

llvm/include/llvm/Analysis/CodeMetrics.h

 #include "llvm/Support/InstructionCost.h"
 
 namespace llvm {
 class AssumptionCache;
 class BasicBlock;
 class Loop;
 class Function;
 template <class T> class SmallPtrSetImpl;
+class TargetLibraryInfo;
 class TargetTransformInfo;
 class Value;
 
 /// Utility to calculate the size and a few similar metrics for a set
 /// of basic blocks.
 struct CodeMetrics {
   /// True if this function contains a call to setjmp or other functions
   /// with attribute "returns twice" without having the attribute itself.
struct CodeMetrics {
   /// Add information about a block to the current state.
   void analyzeBasicBlock(const BasicBlock *BB, const TargetTransformInfo &TTI,
                          const SmallPtrSetImpl<const Value *> &EphValues,
                          bool PrepareForLTO = false);
 
   /// Collect a loop's ephemeral values (those used only by an assume
   /// or similar intrinsics in the loop).
   static void collectEphemeralValues(const Loop *L, AssumptionCache *AC,
-                                     SmallPtrSetImpl<const Value *> &EphValues);
+                                     SmallPtrSetImpl<const Value *> &EphValues,
+                                     const TargetLibraryInfo *TLI = nullptr);
 
   /// Collect a functions's ephemeral values (those used only by an
   /// assume or similar intrinsics in the function).
   static void collectEphemeralValues(const Function *L, AssumptionCache *AC,
-                                     SmallPtrSetImpl<const Value *> &EphValues);
+                                     SmallPtrSetImpl<const Value *> &EphValues,
+                                     const TargetLibraryInfo *TLI = nullptr);
 };
 
 }
 
 #endif

llvm/lib/Analysis/CodeMetrics.cpp

 #include "llvm/Analysis/CodeMetrics.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/Function.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/InstructionCost.h"
+#include "llvm/Transforms/Utils/Local.h"
 
 #define DEBUG_TYPE "code-metrics"
 
 using namespace llvm;
 
-static void
-appendSpeculatableOperands(const Value *V,
+static void appendSpeculatableOperands(Value *V,
                                        SmallPtrSetImpl<const Value *> &Visited,
-                           SmallVectorImpl<const Value *> &Worklist) {
-  const User *U = dyn_cast<User>(V);
+                                       SmallVectorImpl<Value *> &Worklist,
+                                       const TargetLibraryInfo *TLI) {
+  User *U = dyn_cast<User>(V);
   if (!U)
     return;
 
-  for (const Value *Operand : U->operands())
+  for (Value *Operand : U->operands())
     if (Visited.insert(Operand).second)
-      if (const auto *I = dyn_cast<Instruction>(Operand))
-        if (!I->mayHaveSideEffects() && !I->isTerminator())
+      if (auto *I = dyn_cast<Instruction>(Operand))
+        if (wouldInstructionBeTriviallyDead(I, TLI))
           Worklist.push_back(I);
 }
 
-static void completeEphemeralValues(SmallPtrSetImpl<const Value *> &Visited,
-                                    SmallVectorImpl<const Value *> &Worklist,
-                                    SmallPtrSetImpl<const Value *> &EphValues) {
-  // Note: We don't speculate PHIs here, so we'll miss instruction chains kept
-  // alive only by ephemeral values.
+static bool collectUsersOfEphemeralCandidate(
+    Value *V, SmallPtrSetImpl<const Value *> &EphValues,
+    SmallPtrSetImpl<Value *> &UsersAndV, const TargetLibraryInfo *TLI) {
+  UsersAndV.insert(V);
+
+  SmallVector<Value *, 4> Worklist;
+  Worklist.push_back(V);
+  while (!Worklist.empty()) {
+    Value *Curr = Worklist.pop_back_val();
+    for (auto *U : Curr->users()) {
+      if (EphValues.count(U) || !UsersAndV.insert(U).second)
+        continue;
+      auto *I = dyn_cast<Instruction>(U);
+      if (!I || !wouldInstructionBeTriviallyDead(I, TLI))
+        return false;
+      Worklist.push_back(I);
+    }
+  }
+  return true;
+}
 
+static void completeEphemeralValues(SmallPtrSetImpl<const Value *> &Visited,
+                                    SmallVectorImpl<Value *> &Worklist,
+                                    SmallPtrSetImpl<const Value *> &EphValues,
+                                    const TargetLibraryInfo *TLI) {
   // Walk the worklist using an index but without caching the size so we can
   // append more entries as we process the worklist. This forms a queue without
   // quadratic behavior by just leaving processed nodes at the head of the
   // worklist forever.
+  SmallPtrSet<Value *, 4> Users;
   for (int i = 0; i < (int)Worklist.size(); ++i) {
-    const Value *V = Worklist[i];
+    Value *V = Worklist[i];
+    if (EphValues.count(V))
+      continue;
 
     assert(Visited.count(V) &&
            "Failed to add a worklist entry to our visited set!");
 
     // If all uses of this value are ephemeral, then so is this value.
-    if (!all_of(V->users(), [&](const User *U) { return EphValues.count(U); }))
-      continue;
-
-    EphValues.insert(V);
-    LLVM_DEBUG(dbgs() << "Ephemeral Value: " << *V << "\n");
+    Users.clear();
+    if (collectUsersOfEphemeralCandidate(V, EphValues, Users, TLI)) {
+      for (auto *EphV : Users) {
+        EphValues.insert(EphV);
+        LLVM_DEBUG(dbgs() << "Ephemeral Value: " << *EphV << "\n");
+        Visited.insert(EphV);
 
         // Append any more operands to consider.
-    appendSpeculatableOperands(V, Visited, Worklist);
+        appendSpeculatableOperands(EphV, Visited, Worklist, TLI);
+      }
+    }
   }
 }
 
 // Find all ephemeral values.
-void CodeMetrics::collectEphemeralValues(
-    const Loop *L, AssumptionCache *AC,
-    SmallPtrSetImpl<const Value *> &EphValues) {
+void
+CodeMetrics::collectEphemeralValues(const Loop *L, AssumptionCache *AC,
+                                    SmallPtrSetImpl<const Value *> &EphValues,
+                                    const TargetLibraryInfo *TLI) {
   SmallPtrSet<const Value *, 32> Visited;
-  SmallVector<const Value *, 16> Worklist;
+  SmallVector<Value *, 16> Worklist;
 
   for (auto &AssumeVH : AC->assumptions()) {
     if (!AssumeVH)
       continue;
     Instruction *I = cast<Instruction>(AssumeVH);
 
     // Filter out call sites outside of the loop so we don't do a function's
     // worth of work for each of its loops (and, in the common case, ephemeral
     // values in the loop are likely due to @llvm.assume calls in the loop).
     if (!L->contains(I->getParent()))
       continue;
 
     if (EphValues.insert(I).second)
-      appendSpeculatableOperands(I, Visited, Worklist);
+      appendSpeculatableOperands(I, Visited, Worklist, TLI);
   }
 
-  completeEphemeralValues(Visited, Worklist, EphValues);
+  completeEphemeralValues(Visited, Worklist, EphValues, TLI);
 }
 
-void CodeMetrics::collectEphemeralValues(
-    const Function *F, AssumptionCache *AC,
-    SmallPtrSetImpl<const Value *> &EphValues) {
+void
+CodeMetrics::collectEphemeralValues(const Function *F, AssumptionCache *AC,
+                                    SmallPtrSetImpl<const Value *> &EphValues,
+                                    const TargetLibraryInfo *TLI) {
   SmallPtrSet<const Value *, 32> Visited;
-  SmallVector<const Value *, 16> Worklist;
+  SmallVector<Value *, 16> Worklist;
 
   for (auto &AssumeVH : AC->assumptions()) {
     if (!AssumeVH)
       continue;
     Instruction *I = cast<Instruction>(AssumeVH);
     assert(I->getParent()->getParent() == F &&
            "Found assumption for the wrong function!");
 
     if (EphValues.insert(I).second)
-      appendSpeculatableOperands(I, Visited, Worklist);
+      appendSpeculatableOperands(I, Visited, Worklist, TLI);
   }
 
-  completeEphemeralValues(Visited, Worklist, EphValues);
+  completeEphemeralValues(Visited, Worklist, EphValues, TLI);
 }
 
 /// Fill in the current structure with information gleaned from the specified
 /// block.
 void CodeMetrics::analyzeBasicBlock(
     const BasicBlock *BB, const TargetTransformInfo &TTI,
     const SmallPtrSetImpl<const Value *> &EphValues, bool PrepareForLTO) {
   ++NumBlocks;

llvm/lib/CodeGen/CodeGenPrepare.cpp

 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/BypassSlowDivision.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/SimplifyLibCalls.h"
 #include "llvm/Transforms/Utils/SizeOpts.h"
+#include "llvm/Analysis/CodeMetrics.h"
+#include "llvm/Analysis/AssumptionCache.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <iterator>
 #include <limits>
 #include <memory>
 #include <optional>
 #include <utility>
void releaseMemory() override {
     BPI.reset();
     BFI.reset();
   }
 
   StringRef getPassName() const override { return "CodeGen Prepare"; }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     // FIXME: When we can selectively preserve passes, preserve the domtree.
+    AU.addRequired<AssumptionCacheTracker>();
     AU.addRequired<ProfileSummaryInfoWrapperPass>();
     AU.addRequired<TargetLibraryInfoWrapperPass>();
     AU.addRequired<TargetPassConfig>();
     AU.addRequired<TargetTransformInfoWrapperPass>();
     AU.addRequired<LoopInfoWrapperPass>();
     AU.addUsedIfAvailable<BasicBlockSectionsProfileReader>();
   }
 
private:
   // Get the DominatorTree, building if necessary.
   DominatorTree &getDT(Function &F) {
     if (!DT)
       DT = std::make_unique<DominatorTree>(F);
     return *DT;
   }
 
   void removeAllAssertingVHReferences(Value *V);
-  bool eliminateAssumptions(Function &F);
+  bool eliminateAssumptions(Function &F, AssumptionCache &AC);
   bool eliminateFallThrough(Function &F, DominatorTree *DT = nullptr);
   bool eliminateMostlyEmptyBlocks(Function &F);
   BasicBlock *findDestBlockOfMergeableEmptyBlock(BasicBlock *BB);
   bool canMergeBlocks(const BasicBlock *BB, const BasicBlock *DestBB) const;
   void eliminateMostlyEmptyBlock(BasicBlock *BB);
   bool isMergingEmptyBlockProfitable(BasicBlock *BB, BasicBlock *DestBB,
                                      bool isPreheader);
   bool makeBitReverse(Instruction &I);

 };
 
 } // end anonymous namespace
 
 char CodeGenPrepare::ID = 0;
 
 INITIALIZE_PASS_BEGIN(CodeGenPrepare, DEBUG_TYPE,
                       "Optimize for code generation", false, false)
+INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(BasicBlockSectionsProfileReader)
 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(ProfileSummaryInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_END(CodeGenPrepare, DEBUG_TYPE, "Optimize for code generation",
                     false, false)
while (BB != nullptr) {
         EverMadeChange |= bypassSlowDivision(BB, BypassWidths);
       BB = Next;
     }
   }
 
   // Get rid of @llvm.assume builtins before attempting to eliminate empty
   // blocks, since there might be blocks that only contain @llvm.assume calls
   // (plus arguments that we can get rid of).
-  EverMadeChange |= eliminateAssumptions(F);
+  auto &AC = getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
+  EverMadeChange |= eliminateAssumptions(F, AC);
 
   // Eliminate blocks that contain only PHI nodes and an
   // unconditional branch.
   EverMadeChange |= eliminateMostlyEmptyBlocks(F);
 
   ModifyDT ModifiedDT = ModifyDT::NotModifyDT;
   if (!DisableBranchOpts)
     EverMadeChange |= splitBranchCondition(F, ModifiedDT);

 #ifndef NDEBUG
   if (VerifyBFIUpdates)
     verifyBFIUpdates(F);
 #endif
 
   return EverMadeChange;
 }
 
-bool CodeGenPrepare::eliminateAssumptions(Function &F) {
+bool CodeGenPrepare::eliminateAssumptions(Function &F, AssumptionCache &AC) {
   bool MadeChange = false;
-  for (BasicBlock &BB : F) {
-    CurInstIterator = BB.begin();
-    while (CurInstIterator != BB.end()) {
-      Instruction *I = &*(CurInstIterator++);
-      if (auto *Assume = dyn_cast<AssumeInst>(I)) {
-        MadeChange = true;
-        Value *Operand = Assume->getOperand(0);
-        Assume->eraseFromParent();
+  SmallPtrSet<const Value *, 16> EphValues;
+  CodeMetrics::collectEphemeralValues(&F, &AC, EphValues, TLInfo);
 
-        resetIteratorIfInvalidatedWhileCalling(&BB, [&]() {
-          RecursivelyDeleteTriviallyDeadInstructions(Operand, TLInfo, nullptr);
-        });
-      }
+  for (const auto *V : EphValues) {
+    if (auto *CI = dyn_cast<Instruction>(V)) {
+      // const_cast is not a good idea, but here we know that Value is not
+      // constant but just an API collectEphemeralValues provides.
+      // Alternative would be to traverse all instructions and check whether it
+      // is in set but it is a redundant compiler time overhead.
+      auto *I = const_cast<Instruction *>(CI);
+      salvageDebugInfo(*I);
+      I->replaceAllUsesWith(PoisonValue::get(I->getType()));
+      I->eraseFromParent();
+      MadeChange = true;
     }
   }
+
   return MadeChange;
 }
 
 /// An instruction is about to be deleted, so remove all references to it in our
 /// GEP-tracking data strcutures.
 void CodeGenPrepare::removeAllAssertingVHReferences(Value *V) {
   LargeOffsetGEPMap.erase(V);
   NewGEPBases.erase(V);

llvm/test/Transforms/CodeGenPrepare/X86/delete-assume-dead-code.ll

 }
 
 define i32 @test2(i32 %N) nounwind {
 ; CHECK-LABEL: @test2(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[HEADER:%.*]]
 ; CHECK:       header:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_INC:%.*]], [[HEADER]] ]
-; CHECK-NEXT:    [[IV2:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[IV2_INC:%.*]], [[HEADER]] ]
 ; CHECK-NEXT:    [[IV_INC]] = add i32 [[IV]], 1
-; CHECK-NEXT:    [[IV2_INC]] = add i32 [[IV2]], 1
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[IV]], [[N:%.*]]
 ; CHECK-NEXT:    br i1 [[CMP]], label [[EXIT:%.*]], label [[HEADER]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret i32 [[IV]]
 ;
 entry:
   br label %header
 

 }
 
 define i32 @test3(i32 %N) nounwind {
 ; CHECK-LABEL: @test3(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[HEADER:%.*]]
 ; CHECK:       header:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_INC:%.*]], [[HEADER]] ]
-; CHECK-NEXT:    [[IV2:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[IV2_INC:%.*]], [[HEADER]] ]
 ; CHECK-NEXT:    [[IV_INC]] = add i32 [[IV]], 1
-; CHECK-NEXT:    [[IV2_INC]] = add i32 [[IV2]], 1
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[IV]], [[N:%.*]]
 ; CHECK-NEXT:    br i1 [[CMP]], label [[EXIT:%.*]], label [[HEADER]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret i32 [[IV]]
 ;
 entry:
   br label %header
 

 
 ; Two assumes case.
 define i32 @test4(i32 %N) nounwind {
 ; CHECK-LABEL: @test4(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[HEADER:%.*]]
 ; CHECK:       header:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_INC:%.*]], [[HEADER]] ]
-; CHECK-NEXT:    [[IV2:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[IV2_INC:%.*]], [[HEADER]] ]
 ; CHECK-NEXT:    [[IV_INC]] = add i32 [[IV]], 1
-; CHECK-NEXT:    [[IV2_INC]] = add i32 [[IV2]], 1
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[IV]], [[N:%.*]]
 ; CHECK-NEXT:    br i1 [[CMP]], label [[EXIT:%.*]], label [[HEADER]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret i32 [[IV]]
 ;
 entry:
   br label %header