[sil-inst-opt] Improve performance of InstModCallbacks by eliminating indirect call along default callback path.

include/swift/SILOptimizer/Analysis/SimplifyInstruction.h

@@ -25,6 +25,7 @@
 namespace swift {
 
 class SILInstruction;
+class InstModCallbacks;
 
 /// Try to simplify the specified instruction, performing local
 /// analysis of the operands of the instruction, without looking at its uses
@@ -46,11 +47,17 @@ SILValue simplifyInstruction(SILInstruction *I);
 ///
 /// NOTE: When OSSA is enabled this API assumes OSSA is properly formed and will
 /// insert compensating instructions.
-SILBasicBlock::iterator replaceAllSimplifiedUsesAndErase(
-    SILInstruction *I, SILValue result,
-    std::function<void(SILInstruction *)> eraseNotify = nullptr,
-    std::function<void(SILInstruction *)> newInstNotify = nullptr,
-    DeadEndBlocks *deadEndBlocks = nullptr);
+SILBasicBlock::iterator
+replaceAllSimplifiedUsesAndErase(SILInstruction *I, SILValue result,
+                                 InstModCallbacks &callbacks,
+                                 DeadEndBlocks *deadEndBlocks = nullptr);
+
+/// This is a low level routine that makes all uses of \p svi uses of \p
+/// newValue (ignoring end scope markers) and then deletes \p svi and all end
+/// scope markers. Then returns the next inst to process.
+SILBasicBlock::iterator replaceAllUsesAndErase(SingleValueInstruction *svi,
+                                               SILValue newValue,
+                                               InstModCallbacks &callbacks);
 
 /// Simplify invocations of builtin operations that may overflow.
 /// All such operations return a tuple (result, overflow_flag).

include/swift/SILOptimizer/Utils/CFGOptUtils.h

@@ -36,7 +36,7 @@ namespace swift {
 class DominanceInfo;
 class SILLoop;
 class SILLoopInfo;
-struct InstModCallbacks;
+class InstModCallbacks;
 
 /// Adds a new argument to an edge between a branch and a destination
 /// block. Allows for user injected callbacks via \p callbacks.
@@ -47,8 +47,7 @@ struct InstModCallbacks;
 /// \return The created branch. The old branch is deleted.
 /// The argument is appended at the end of the argument tuple.
 TermInst *addNewEdgeValueToBranch(TermInst *branch, SILBasicBlock *dest,
-                                  SILValue val,
-                                  const InstModCallbacks &callbacks);
+                                  SILValue val, InstModCallbacks &callbacks);
 
 /// Adds a new argument to an edge between a branch and a destination
 /// block.
@@ -60,7 +59,8 @@ TermInst *addNewEdgeValueToBranch(TermInst *branch, SILBasicBlock *dest,
 /// The argument is appended at the end of the argument tuple.
 inline TermInst *addNewEdgeValueToBranch(TermInst *branch, SILBasicBlock *dest,
                                          SILValue val) {
-  return addNewEdgeValueToBranch(branch, dest, val, InstModCallbacks());
+  InstModCallbacks callbacks;
+  return addNewEdgeValueToBranch(branch, dest, val, callbacks);
 }
 
 /// Changes the edge value between a branch and destination basic block

include/swift/SILOptimizer/Utils/CanonicalizeInstruction.h

@@ -29,6 +29,7 @@
 #include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/SILBasicBlock.h"
 #include "swift/SIL/SILInstruction.h"
+#include "swift/SILOptimizer/Utils/InstOptUtils.h"
 #include "llvm/Support/Debug.h"
 
 namespace swift {
@@ -40,10 +41,18 @@ struct CanonicalizeInstruction {
   static constexpr const char *defaultDebugType = "sil-canonicalize";
   const char *debugType = defaultDebugType;
   DeadEndBlocks &deadEndBlocks;
+  InstModCallbacks callbacks;
 
   CanonicalizeInstruction(const char *passDebugType,
                           DeadEndBlocks &deadEndBlocks)
-      : deadEndBlocks(deadEndBlocks) {
+      : deadEndBlocks(deadEndBlocks),
+        callbacks(
+            [&](SILInstruction *toDelete) { killInstruction(toDelete); },
+            [&](SILInstruction *newInst) { notifyNewInstruction(newInst); },
+            [&](SILValue oldValue, SILValue newValue) {
+              oldValue->replaceAllUsesWith(newValue);
+              notifyHasNewUsers(newValue);
+            }) {
 #ifndef NDEBUG
     if (llvm::DebugFlag && !llvm::isCurrentDebugType(debugType))
       debugType = passDebugType;

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -310,44 +310,97 @@ bool tryCheckedCastBrJumpThreading(
 
 /// A structure containing callbacks that are called when an instruction is
 /// removed or added.
-struct InstModCallbacks {
-  static const std::function<void(SILInstruction *)> defaultDeleteInst;
-  static const std::function<void(SILInstruction *)> defaultCreatedNewInst;
-  static const std::function<void(SILValue, SILValue)> defaultReplaceValueUsesWith;
-  static const std::function<void(SingleValueInstruction *, SILValue)>
-      defaultEraseAndRAUWSingleValueInst;
-
-  std::function<void(SILInstruction *)> deleteInst =
-      InstModCallbacks::defaultDeleteInst;
-  std::function<void(SILInstruction *)> createdNewInst =
-      InstModCallbacks::defaultCreatedNewInst;
-  std::function<void(SILValue, SILValue)>
-      replaceValueUsesWith =
-          InstModCallbacks::defaultReplaceValueUsesWith;
-  std::function<void(SingleValueInstruction *, SILValue)>
-      eraseAndRAUWSingleValueInst =
-          InstModCallbacks::defaultEraseAndRAUWSingleValueInst;
-
-  InstModCallbacks(decltype(deleteInst) deleteInst,
-                   decltype(createdNewInst) createdNewInst,
-                   decltype(replaceValueUsesWith) replaceValueUsesWith)
-      : deleteInst(deleteInst), createdNewInst(createdNewInst),
-        replaceValueUsesWith(replaceValueUsesWith),
-        eraseAndRAUWSingleValueInst(
-            InstModCallbacks::defaultEraseAndRAUWSingleValueInst) {}
+///
+/// PERFORMANCE NOTES: This code can be used in loops, so we want to make sure
+/// to not have overhead when the user does not specify a callback. To do that
+/// instead of defining a "default" std::function, we represent the "default"
+/// functions as nullptr. Then, in the helper function trampoline that actually
+/// gets called, we check if we have a nullptr and if we do, we perform the
+/// default operation inline. What is nice about this from a perf perspective is
+/// that in a loop this property should predict well since you have a single
+/// branch that is going to go the same way everytime.
+class InstModCallbacks {
+  /// A function that takes in an instruction and deletes the inst.
+  ///
+  /// Default implementation is instToDelete->eraseFromParent();
+  std::function<void(SILInstruction *instToDelete)> deleteInstFunc;
+
+  /// A function that is called to notify that a new function was created.
+  ///
+  /// Default implementation is a no-op, but we still mark madeChange.
+  std::function<void(SILInstruction *newlyCreatedInst)> createdNewInstFunc;
+
+  /// A function that first replaces all uses of oldValue with uses of newValue.
+  ///
+  /// Default implementation just calls oldValue->replaceAllUsesWith(newValue);
+  std::function<void(SILValue oldValue, SILValue newValue)>
+      replaceValueUsesWithFunc;
+
+  /// A function that first replaces all uses of oldValue with uses of newValue
+  /// and then erases oldValue.
+  ///
+  /// Default implementation just calls replaceValueUsesWithFunc and then
+  /// deleteInstFunc.
+  std::function<void(SingleValueInstruction *oldValue, SILValue newValue)>
+      eraseAndRAUWSingleValueInstFunc;
+
+  /// A boolean that tracks if any of our callbacks were ever called.
+  bool madeChange = false;
+
+public:
+  InstModCallbacks(decltype(deleteInstFunc) deleteInstFunc)
+      : deleteInstFunc(deleteInstFunc) {}
+
+  InstModCallbacks(decltype(deleteInstFunc) deleteInstFunc,
+                   decltype(createdNewInstFunc) createdNewInstFunc,
+                   decltype(replaceValueUsesWithFunc) replaceValueUsesWithFunc)
+      : deleteInstFunc(deleteInstFunc), createdNewInstFunc(createdNewInstFunc),
+        replaceValueUsesWithFunc(replaceValueUsesWithFunc) {}
 
   InstModCallbacks(
-      decltype(deleteInst) deleteInst, decltype(createdNewInst) createdNewInst,
-      decltype(replaceValueUsesWith) replaceValueUsesWith,
-      decltype(eraseAndRAUWSingleValueInst) eraseAndRAUWSingleValueInst)
-      : deleteInst(deleteInst), createdNewInst(createdNewInst),
-        replaceValueUsesWith(replaceValueUsesWith),
-        eraseAndRAUWSingleValueInst(eraseAndRAUWSingleValueInst) {}
+      decltype(deleteInstFunc) deleteInstFunc,
+      decltype(createdNewInstFunc) createdNewInstFunc,
+      decltype(replaceValueUsesWithFunc) replaceValueUsesWithFunc,
+      decltype(eraseAndRAUWSingleValueInstFunc) eraseAndRAUWSingleValueInstFunc)
+      : deleteInstFunc(deleteInstFunc), createdNewInstFunc(createdNewInstFunc),
+        replaceValueUsesWithFunc(replaceValueUsesWithFunc),
+        eraseAndRAUWSingleValueInstFunc(eraseAndRAUWSingleValueInstFunc) {}
 
   InstModCallbacks() = default;
   ~InstModCallbacks() = default;
   InstModCallbacks(const InstModCallbacks &) = default;
   InstModCallbacks(InstModCallbacks &&) = default;
+
+  void deleteInst(SILInstruction *instToDelete) {
+    madeChange = true;
+    if (deleteInstFunc)
+      return deleteInstFunc(instToDelete);
+    instToDelete->eraseFromParent();
+  }
+
+  void createdNewInst(SILInstruction *newlyCreatedInst) {
+    madeChange = true;
+    if (createdNewInstFunc)
+      createdNewInstFunc(newlyCreatedInst);
+  }
+
+  void replaceValueUsesWith(SILValue oldValue, SILValue newValue) {
+    madeChange = true;
+    if (replaceValueUsesWithFunc)
+      return replaceValueUsesWithFunc(oldValue, newValue);
+    oldValue->replaceAllUsesWith(newValue);
+  }
+
+  void eraseAndRAUWSingleValueInst(SingleValueInstruction *oldInst,
+                                   SILValue newValue) {
+    madeChange = true;
+    if (eraseAndRAUWSingleValueInstFunc)
+      return eraseAndRAUWSingleValueInstFunc(oldInst, newValue);
+    replaceValueUsesWith(oldInst, newValue);
+    deleteInst(oldInst);
+  }
+
+  bool getMadeChange() const { return madeChange; }
 };
 
 /// Get all consumed arguments of a partial_apply.

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -21,15 +21,15 @@
 
 #include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/SILModule.h"
+#include "swift/SILOptimizer/Utils/InstOptUtils.h"
 
 namespace swift {
 
 // Defined in BasicBlockUtils.h
 struct JointPostDominanceSetComputer;
 
 struct OwnershipFixupContext {
-  std::function<void(SILInstruction *)> eraseNotify;
-  std::function<void(SILInstruction *)> newInstNotify;
+  InstModCallbacks callbacks;
   DeadEndBlocks &deBlocks;
   JointPostDominanceSetComputer &jointPostDomSetComputer;
 

lib/SILOptimizer/Analysis/SimplifyInstruction.cpp

@@ -735,9 +735,9 @@ case BuiltinValueKind::id:
 //                           Top Level Entrypoints
 //===----------------------------------------------------------------------===//
 
-static SILBasicBlock::iterator
-replaceAllUsesAndEraseInner(SingleValueInstruction *svi, SILValue newValue,
-                            std::function<void(SILInstruction *)> eraseNotify) {
+SILBasicBlock::iterator
+swift::replaceAllUsesAndErase(SingleValueInstruction *svi, SILValue newValue,
+                              InstModCallbacks &callbacks) {
   assert(svi != newValue && "Cannot RAUW a value with itself");
   SILBasicBlock::iterator nextii = std::next(svi->getIterator());
 
@@ -749,18 +749,13 @@ replaceAllUsesAndEraseInner(SingleValueInstruction *svi, SILValue newValue,
     if (isEndOfScopeMarker(user)) {
       if (&*nextii == user)
         ++nextii;
-      if (eraseNotify)
-        eraseNotify(user);
-      else
-        user->eraseFromParent();
+      callbacks.deleteInst(user);
       continue;
     }
     use->set(newValue);
   }
-  if (eraseNotify)
-    eraseNotify(svi);
-  else
-    svi->eraseFromParent();
+
+  callbacks.deleteInst(svi);
 
   return nextii;
 }
@@ -775,21 +770,19 @@ replaceAllUsesAndEraseInner(SingleValueInstruction *svi, SILValue newValue,
 /// before this is called. It will perform fixups as necessary to preserve OSSA.
 ///
 /// Return an iterator to the next (nondeleted) instruction.
-SILBasicBlock::iterator swift::replaceAllSimplifiedUsesAndErase(
-    SILInstruction *i, SILValue result,
-    std::function<void(SILInstruction *)> eraseNotify,
-    std::function<void(SILInstruction *)> newInstNotify,
-    DeadEndBlocks *deadEndBlocks) {
+SILBasicBlock::iterator
+swift::replaceAllSimplifiedUsesAndErase(SILInstruction *i, SILValue result,
+                                        InstModCallbacks &callbacks,
+                                        DeadEndBlocks *deadEndBlocks) {
   auto *svi = cast<SingleValueInstruction>(i);
   assert(svi != result && "Cannot RAUW a value with itself");
 
   if (svi->getFunction()->hasOwnership()) {
     JointPostDominanceSetComputer computer(*deadEndBlocks);
-    OwnershipFixupContext ctx{eraseNotify, newInstNotify, *deadEndBlocks,
-                              computer};
+    OwnershipFixupContext ctx{callbacks, *deadEndBlocks, computer};
     return ctx.replaceAllUsesAndEraseFixingOwnership(svi, result);
   }
-  return replaceAllUsesAndEraseInner(svi, result, eraseNotify);
+  return replaceAllUsesAndErase(svi, result, callbacks);
 }
 
 /// Simplify invocations of builtin operations that may overflow.

lib/SILOptimizer/Mandatory/DefiniteInitialization.cpp

@@ -866,7 +866,8 @@ static void replaceValueMetatypeInstWithMetatypeArgument(
         valueMetatype->getLoc(), metatypeArgument,
         valueMetatype->getType());
   }
-  replaceAllSimplifiedUsesAndErase(valueMetatype, metatypeArgument);
+  InstModCallbacks callbacks;
+  replaceAllSimplifiedUsesAndErase(valueMetatype, metatypeArgument, callbacks);
 }
 
 void LifetimeChecker::handleLoadForTypeOfSelfUse(DIMemoryUse &Use) {

lib/SILOptimizer/Mandatory/OwnershipModelEliminator.cpp

@@ -435,10 +435,10 @@ static bool stripOwnership(SILFunction &func) {
     if (!value.hasValue())
       continue;
     if (SILValue newValue = simplifyInstruction(*value)) {
-      replaceAllSimplifiedUsesAndErase(*value, newValue,
-                                       [&](SILInstruction *instToErase) {
-                                         visitor.eraseInstruction(instToErase);
-                                       });
+      InstModCallbacks callbacks([&](SILInstruction *instToErase) {
+        visitor.eraseInstruction(instToErase);
+      });
+      replaceAllSimplifiedUsesAndErase(*value, newValue, callbacks);
       madeChange = true;
     }
   }

lib/SILOptimizer/Transforms/CSE.cpp

@@ -955,6 +955,7 @@ static bool isLazyPropertyGetter(ApplyInst *ai) {
 
 bool CSE::processNode(DominanceInfoNode *Node) {
   SILBasicBlock *BB = Node->getBlock();
+  InstModCallbacks callbacks;
   bool Changed = false;
 
   // See if any instructions in the block can be eliminated.  If so, do it.  If
@@ -981,8 +982,7 @@ bool CSE::processNode(DominanceInfoNode *Node) {
     if (SILValue V = simplifyInstruction(Inst)) {
       LLVM_DEBUG(llvm::dbgs()
                  << "SILCSE SIMPLIFY: " << *Inst << "  to: " << *V << '\n');
-      nextI = replaceAllSimplifiedUsesAndErase(Inst, V, nullptr, nullptr,
-                                               &DeadEndBBs);
+      nextI = replaceAllSimplifiedUsesAndErase(Inst, V, callbacks, &DeadEndBBs);
       Changed = true;
       ++NumSimplify;
       continue;
@@ -1412,9 +1412,8 @@ class SILCSE : public SILFunctionTransform {
     InstructionCloner Cloner(Fn);
     DeadEndBlocks DeadEndBBs(Fn);
     JointPostDominanceSetComputer Computer(DeadEndBBs);
-    OwnershipFixupContext FixupCtx{/* eraseNotify */ nullptr,
-                                   /* newInstNotify */ nullptr, DeadEndBBs,
-                                   Computer};
+    InstModCallbacks callbacks;
+    OwnershipFixupContext FixupCtx{callbacks, DeadEndBBs, Computer};
     CSE C(RunsOnHighLevelSil, SEA, FuncBuilder, OpenedArchetypesTracker, Cloner,
           DeadEndBBs, FixupCtx);
     bool Changed = false;

lib/SILOptimizer/Transforms/SimplifyCFG.cpp

@@ -1139,6 +1139,7 @@ bool SimplifyCFG::tryJumpThreading(BranchInst *BI) {
 /// result in exposing opportunities for CFG simplification.
 bool SimplifyCFG::simplifyBranchOperands(OperandValueArrayRef Operands) {
   bool Simplified = false;
+  InstModCallbacks callbacks;
   for (auto O = Operands.begin(), E = Operands.end(); O != E; ++O) {
     // All of our interesting simplifications are on single-value instructions
     // for now.
@@ -1149,7 +1150,7 @@ bool SimplifyCFG::simplifyBranchOperands(OperandValueArrayRef Operands) {
       // unreachable block. In this case it can reference itself as operand.
       if (Result && Result != I) {
         LLVM_DEBUG(llvm::dbgs() << "simplify branch operand " << *I);
-        replaceAllSimplifiedUsesAndErase(I, Result);
+        replaceAllSimplifiedUsesAndErase(I, Result, callbacks);
         Simplified = true;
       }
     }