[ConstantFolding] Accept FoldedOps cache as parameter

[nikic]

Summary

InstCombine internally caches constant operands that it folds. ConstantFolding also internally caches folded operands, but does so only on the "second level": The direct operands of an instruction/constant are not cached, but the operands of the operands are. Presumably that's an oversight...

This patch allows passing in the folded operand cache as an extra parameter, and changes the ConstantFoldConstant API to also look up from it / write to it by itself. That way all users handle caching correctly, and the cache can be shared across calls (done only in InstCombine for now).

In particular this will also avoid duplicate constant folding between ConstantFoldInstruction and operand constant folding in InstCombine. I haven't run into this as a perf problem myself, but recently ran across D18155, which did identify this as a performance problem for some types of code.

[efriedma]

The point of the cache in ConstantFolding is really to avoid bad behavior for constant expression trees with a "diamond" shape; in that context, it doesn't really matter what level it starts at, since constants have a limited number of operands. But sure, if we have the cache anyway, we might as well use it.

I'd like to see some sort of performance measurement here, given we're complicating the ConstantFoldConstant API for the sake of performance.

lib/Analysis/ConstantFolding.cpp

1192

Do you have any idea why we're checking for ConstantVector specifically, instead of ConstantAggregate?

1195

Might as well use DenseMap::lookup() here.

[nikic]

lib/Analysis/ConstantFolding.cpp

1192

The ConstantVector handling was added lateron in https://github.com/llvm/llvm-project/commit/d536f2328ededb3aae6563c721c6134c735f1918, so the reason might be just "it hasn't been implemented".

Maybe there is also a compile-time concern here? A ConstantVector will usually have relatively few operands, while a ConstantArray might have thousands.

[efriedma]

lib/Analysis/ConstantFolding.cpp

1192

"it hasn't been implemented" makes sense, I guess.

[nikic]

Going to close this one, as I actually see a 0.2% regression with it (http://llvm-compile-time-tracker.com/compare.php?from=b236b4cb430674a98917ad1b2fa9f7f9838e66f1&to=a0dcbfc3080e6dd91c7b4253dd83e9eea4eefd03&stat=instructions). Presumably, for the cases where we don't have a lot of unfoldable constant expressions (which is the common case), we end up not actually benefiting from the cache, but still paying (a bit) for it.

include/llvm/Analysis/ConstantFolding.h

 // DataLayout information. These functions cannot go in VMCore due to library
 // dependency issues.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_CONSTANTFOLDING_H
 #define LLVM_ANALYSIS_CONSTANTFOLDING_H
 
+#include "llvm/ADT/DenseMap.h"
+
 namespace llvm {
 class APInt;
 template <typename T> class ArrayRef;
 class CallBase;
 class Constant;
 class ConstantExpr;
 class ConstantVector;
 class DataLayout;

 /// ConstantFoldInstruction - Try to constant fold the specified instruction.
 /// If successful, the constant result is returned, if not, null is returned.
 /// Note that this fails if not all of the operands are constant.  Otherwise,
 /// this function can only fail when attempting to fold instructions like loads
 /// and stores, which have no constant expression form.
 Constant *ConstantFoldInstruction(Instruction *I, const DataLayout &DL,
                                   const TargetLibraryInfo *TLI = nullptr);
 
+/// ConstantFoldInstruction overload that allows persisting the folded
+/// expression cache across calls.
+Constant *
+ConstantFoldInstruction(Instruction *I, const DataLayout &DL,
+                        const TargetLibraryInfo *TLI,
+                        SmallDenseMap<Constant *, Constant *> &FoldedOps);
+
 /// ConstantFoldConstant - Fold the constant using the specified DataLayout.
 /// This function always returns a non-null constant: Either the folding result,
 /// or the original constant if further folding is not possible.
 Constant *ConstantFoldConstant(const Constant *C, const DataLayout &DL,
                                const TargetLibraryInfo *TLI = nullptr);
 
+/// ConstantFoldConstant overload that allows persisting the folded expression
+/// cache across calls.
+Constant *
+ConstantFoldConstant(const Constant *C, const DataLayout &DL,
+                     const TargetLibraryInfo *TLI,
+                     SmallDenseMap<Constant *, Constant *> &FoldedOps);
+
 /// ConstantFoldInstOperands - Attempt to constant fold an instruction with the
 /// specified operands.  If successful, the constant result is returned, if not,
 /// null is returned.  Note that this function can fail when attempting to
 /// fold instructions like loads and stores, which have no constant expression
 /// form.
 ///
 Constant *ConstantFoldInstOperands(Instruction *I, ArrayRef<Constant *> Ops,
                                    const DataLayout &DL,

lib/Analysis/ConstantFolding.cpp

 
 //===----------------------------------------------------------------------===//
 // Constant Folding public APIs
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 Constant *
-ConstantFoldConstantImpl(const Constant *C, const DataLayout &DL,
+ConstantFoldConstantUncached(const Constant *C, const DataLayout &DL,
                              const TargetLibraryInfo *TLI,
                              SmallDenseMap<Constant *, Constant *> &FoldedOps) {
-  if (!isa<ConstantVector>(C) && !isa<ConstantExpr>(C))
-    return const_cast<Constant *>(C);
+  assert(isa<ConstantVector>(C) || isa<ConstantExpr>(C));
 
   SmallVector<Constant *, 8> Ops;
   for (const Use &OldU : C->operands()) {
-    Constant *OldC = cast<Constant>(&OldU);
-    Constant *NewC = OldC;
-    // Recursively fold the ConstantExpr's operands. If we have already folded
-    // a ConstantExpr, we don't have to process it again.
-    if (isa<ConstantVector>(OldC) || isa<ConstantExpr>(OldC)) {
-      auto It = FoldedOps.find(OldC);
-      if (It == FoldedOps.end()) {
-        NewC = ConstantFoldConstantImpl(OldC, DL, TLI, FoldedOps);
-        FoldedOps.insert({OldC, NewC});
-      } else {
-        NewC = It->second;
-      }
-    }
-    Ops.push_back(NewC);
+    Constant *C = cast<Constant>(&OldU);
+    Ops.push_back(ConstantFoldConstant(C, DL, TLI, FoldedOps));
   }
 
   if (auto *CE = dyn_cast<ConstantExpr>(C)) {
     if (CE->isCompare())
       return ConstantFoldCompareInstOperands(CE->getPredicate(), Ops[0], Ops[1],
                                              DL, TLI);
 
     return ConstantFoldInstOperandsImpl(CE, CE->getOpcode(), Ops, DL, TLI);
   }
 
   assert(isa<ConstantVector>(C));
   return ConstantVector::get(Ops);
 }
 
 } // end anonymous namespace
 
 Constant *llvm::ConstantFoldInstruction(Instruction *I, const DataLayout &DL,
                                         const TargetLibraryInfo *TLI) {
+  SmallDenseMap<Constant *, Constant *> FoldedOps;
+  return ConstantFoldInstruction(I, DL, TLI, FoldedOps);
+}
+
+Constant *llvm::ConstantFoldInstruction(
+    Instruction *I, const DataLayout &DL, const TargetLibraryInfo *TLI,
+    SmallDenseMap<Constant *, Constant *> &FoldedOps) {
   // Handle PHI nodes quickly here...
   if (auto *PN = dyn_cast<PHINode>(I)) {
     Constant *CommonValue = nullptr;
 
-    SmallDenseMap<Constant *, Constant *> FoldedOps;
     for (Value *Incoming : PN->incoming_values()) {
       // If the incoming value is undef then skip it.  Note that while we could
       // skip the value if it is equal to the phi node itself we choose not to
       // because that would break the rule that constant folding only applies if
       // all operands are constants.
       if (isa<UndefValue>(Incoming))
         continue;
       // If the incoming value is not a constant, then give up.
       auto *C = dyn_cast<Constant>(Incoming);
       if (!C)
         return nullptr;
       // Fold the PHI's operands.
-      C = ConstantFoldConstantImpl(C, DL, TLI, FoldedOps);
+      C = ConstantFoldConstant(C, DL, TLI, FoldedOps);
       // If the incoming value is a different constant to
       // the one we saw previously, then give up.
       if (CommonValue && C != CommonValue)
         return nullptr;
       CommonValue = C;
     }
 
     // If we reach here, all incoming values are the same constant or undef.
     return CommonValue ? CommonValue : UndefValue::get(PN->getType());
   }
 
   // Scan the operand list, checking to see if they are all constants, if so,
   // hand off to ConstantFoldInstOperandsImpl.
   if (!all_of(I->operands(), [](Use &U) { return isa<Constant>(U); }))
     return nullptr;
 
-  SmallDenseMap<Constant *, Constant *> FoldedOps;
   SmallVector<Constant *, 8> Ops;
   for (const Use &OpU : I->operands()) {
     auto *Op = cast<Constant>(&OpU);
     // Fold the Instruction's operands.
-    Op = ConstantFoldConstantImpl(Op, DL, TLI, FoldedOps);
+    Op = ConstantFoldConstant(Op, DL, TLI, FoldedOps);
     Ops.push_back(Op);
   }
 
   if (const auto *CI = dyn_cast<CmpInst>(I))
     return ConstantFoldCompareInstOperands(CI->getPredicate(), Ops[0], Ops[1],
                                            DL, TLI);
 
   if (const auto *LI = dyn_cast<LoadInst>(I))
return ConstantExpr::getExtractValue(
                                     EVI->getIndices());
   }
 
   return ConstantFoldInstOperands(I, Ops, DL, TLI);
 }
 
 Constant *llvm::ConstantFoldConstant(const Constant *C, const DataLayout &DL,
                                      const TargetLibraryInfo *TLI) {
+  if (!isa<ConstantVector>(C) && !isa<ConstantExpr>(C))
+    return const_cast<Constant *>(C);
+
   SmallDenseMap<Constant *, Constant *> FoldedOps;
-  return ConstantFoldConstantImpl(C, DL, TLI, FoldedOps);
+  return ConstantFoldConstantUncached(C, DL, TLI, FoldedOps);
+}
+
+Constant *
+llvm::ConstantFoldConstant(const Constant *C, const DataLayout &DL,
+                           const TargetLibraryInfo *TLI,
+                           SmallDenseMap<Constant *, Constant *> &FoldedOps) {
+  if (!isa<ConstantVector>(C) && !isa<ConstantExpr>(C))
+    return const_cast<Constant *>(C);
+
+  auto It = FoldedOps.find(C);
+  if (It != FoldedOps.end())
+    return It->second;
+
+  Constant *NewC = ConstantFoldConstantUncached(C, DL, TLI, FoldedOps);
+  FoldedOps.insert({ const_cast<Constant *>(C), NewC });
+  return NewC;
 }
 
 Constant *llvm::ConstantFoldInstOperands(Instruction *I,
                                          ArrayRef<Constant *> Ops,
                                          const DataLayout &DL,
                                          const TargetLibraryInfo *TLI) {
   return ConstantFoldInstOperandsImpl(I, I->getOpcode(), Ops, DL, TLI);
 }

lib/Transforms/InstCombine/InstructionCombining.cpp

static bool AddReachableCodeToWorklist(BasicBlock *BB, const DataLayout &DL,
                                        SmallPtrSetImpl<BasicBlock *> &Visited,
                                        InstCombineWorklist &ICWorklist,
                                        const TargetLibraryInfo *TLI) {
   bool MadeIRChange = false;
   SmallVector<BasicBlock*, 256> Worklist;
   Worklist.push_back(BB);
 
   SmallVector<Instruction*, 128> InstrsForInstCombineWorklist;
-  DenseMap<Constant *, Constant *> FoldedConstants;
+  SmallDenseMap<Constant *, Constant *> FoldedConstants;
 
   do {
     BB = Worklist.pop_back_val();
 
     // We have now visited this block!  If we've already been here, ignore it.
     if (!Visited.insert(BB).second)
       continue;
 
     for (BasicBlock::iterator BBI = BB->begin(), E = BB->end(); BBI != E; ) {
       Instruction *Inst = &*BBI++;
 
       // ConstantProp instruction if trivially constant.
       if (!Inst->use_empty() &&
           (Inst->getNumOperands() == 0 || isa<Constant>(Inst->getOperand(0))))
-        if (Constant *C = ConstantFoldInstruction(Inst, DL, TLI)) {
+        if (Constant *C = ConstantFoldInstruction(Inst, DL, TLI,
+                                                  FoldedConstants)) {
           LLVM_DEBUG(dbgs() << "IC: ConstFold to: " << *C << " from: " << *Inst
                             << '\n');
           Inst->replaceAllUsesWith(C);
           ++NumConstProp;
           if (isInstructionTriviallyDead(Inst, TLI))
             Inst->eraseFromParent();
           MadeIRChange = true;
           continue;
         }
 
       // See if we can constant fold its operands.
       for (Use &U : Inst->operands()) {
         if (!isa<ConstantVector>(U) && !isa<ConstantExpr>(U))
           continue;
 
-        auto *C = cast<Constant>(U);
-        Constant *&FoldRes = FoldedConstants[C];
-        if (!FoldRes)
-          FoldRes = ConstantFoldConstant(C, DL, TLI);
-
+        Constant *C = cast<Constant>(U);
+        Constant *FoldRes = ConstantFoldConstant(C, DL, TLI, FoldedConstants);
         if (FoldRes != C) {
           LLVM_DEBUG(dbgs() << "IC: ConstFold operand of: " << *Inst
                             << "\n    Old = " << *C
                             << "\n    New = " << *FoldRes << '\n');
           U = FoldRes;
           MadeIRChange = true;
         }
       }