SILOptimizer: Add a new TempRValue optimization pass

This is a separate optimization that detects short-lived temporaries that can be eliminated.
This is necessary now that SILGen no longer performs basic RValue forwarding in some cases.

SR-5508: Performance regression in benchmarks caused by removing SILGen peephole for LoadExpr in +0 context

Author: eeckstein

include/swift/SILOptimizer/PassManager/Passes.def

@@ -226,6 +226,8 @@ IRGEN_PASS(LoadableByAddress, "loadable-address",
      "SIL Large Loadable type by-address lowering.")
 PASS(RemovePins, "remove-pins",
      "Remove SIL pin/unpin pairs")
+PASS(TempRValueOpt, "temp-rvalue-opt",
+     "Remove short-lived immutable temporary copies")
 PASS(SideEffectsDumper, "side-effects-dump",
      "Print Side-Effect Information for all Functions")
 PASS(SILCleanup, "cleanup",

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -306,6 +306,9 @@ static void addPerfEarlyModulePassPipeline(SILPassPipelinePlan &P) {
   P.addDeadFunctionElimination();
   // Start by cloning functions from stdlib.
   P.addSILLinker();
+
+  // Cleanup after SILGen: remove trivial copies to temporaries.
+  P.addTempRValueOpt();
 }
 
 static void addHighLevelEarlyLoopOptPipeline(SILPassPipelinePlan &P) {

lib/SILOptimizer/Transforms/CopyForwarding.cpp

@@ -61,6 +61,7 @@
 #include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILVisitor.h"
+#include "swift/SILOptimizer/Analysis/AliasAnalysis.h"
 #include "swift/SILOptimizer/Analysis/DominanceAnalysis.h"
 #include "swift/SILOptimizer/Analysis/PostOrderAnalysis.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
@@ -1324,10 +1325,244 @@ class CopyForwardingPass : public SILFunctionTransform
       invalidateAnalysis(SILAnalysis::InvalidationKind::CallsAndInstructions);
       }
   }
+};
+
+/// Temporary RValue Optimization
+///
+/// Peephole optimization to eliminate short-lived immutable temporary copies.
+/// This handles a common pattern generated by SILGen where temporary RValues
+/// are emitted as copies...
+///
+///   %temp = alloc_stack $T
+///   copy_addr %src to [initialization] %temp : $*T
+///   // no writes to %src and %temp
+///   destroy_addr %temp : $*T
+///   dealloc_stack %temp : $*T
+///
+/// This differs from the copy forwarding algorithm because it handles
+/// copy source and dest lifetimes that are unavoidably overlappying. Instead,
+/// it finds cases in which it is easy to determine that the source is
+/// unmodified during the copy destination's lifetime. Thus, the destination can
+/// be viewed as a short-lived "rvalue".
+class TempRValueOptPass : public SILFunctionTransform
+{
+  AliasAnalysis *AA = nullptr;
+
+  static bool collectLoads(SILInstruction *CurrentInst, SILInstruction *addr,
+                           llvm::SmallPtrSetImpl<SILInstruction *> &loadInsts);
+
+  bool checkNoSourceModification(CopyAddrInst *copyInst,
+                       const llvm::SmallPtrSetImpl<SILInstruction *> &useInsts);
 
+  bool tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst);
+
+  void run() override;
 };
+
+/// The main entry point of the pass.
+void TempRValueOptPass::run() {
+  DEBUG(llvm::dbgs() << "Copy Peephole in Func " << getFunction()->getName()
+        << "\n");
+
+  AA = PM->getAnalysis<AliasAnalysis>();
+  bool Changed = false;
+
+  // Find all copy_addr instructions.
+  for (auto &BB : *getFunction()) {
+    auto II = BB.begin();
+    while (II != BB.end()) {
+      auto *CopyInst = dyn_cast<CopyAddrInst>(&*II);
+
+      if (CopyInst) {
+        // In case of success, this may delete instructions, but not the
+        // CopyInst itself.
+        Changed |= tryOptimizeCopyIntoTemp(CopyInst);
+      }
+
+      // Increment the instruction iterator here. We can't do it at the begin of
+      // the loop because the instruction after CopyInst might be deleted in
+      // in tryOptimizeCopyIntoTemp. We can't do it at the end of the loop
+      // because the CopyInst might be deleted in the following code.
+      ++II;
+
+      // Remove identity copies which are a result of this optimization.
+      if (CopyInst && CopyInst->getSrc() == CopyInst->getDest() &&
+          // Identity copies cannot take the source. This check is just here
+          // to be on the safe side.
+          !CopyInst->isTakeOfSrc()) {
+        // This is either the CopyInst which just got optimized or it is a
+        // follow-up from an earlier iteration, where another copy_addr copied
+        // the temporary back to the source location.
+        CopyInst->eraseFromParent();
+      }
+    }
+  }
+
+  if (Changed) {
+    invalidateAnalysis(SILAnalysis::InvalidationKind::Instructions);
+  }
+}
+
+/// Transitively explore all data flow uses of the given \p address until
+/// reaching a load or returning false.
+bool TempRValueOptPass::
+collectLoads(SILInstruction *user, SILInstruction *address,
+             llvm::SmallPtrSetImpl<SILInstruction *> &loadInsts) {
+  // All normal uses (loads) must be in the initialization block.
+  // (The destroy and dealloc are commonly in a different block though.)
+  if (user->getParent() != address->getParent())
+    return false;
+
+  // Only allow uses that cannot destroy their operand. We need to be sure
+  // that replacing all this temporary's uses with the copy source doesn't
+  // destroy the source. This way, we know that the destroy_addr instructions
+  // that we recorded cover all the temporary's lifetime termination points.
+  //
+  // Currently we whitelist address projections and loads.
+  //
+  // TODO: handle non-destructive projections of enums
+  // (unchecked_take_enum_data_addr of Optional is nondestructive.)
+  switch (user->getKind()) {
+    default:
+      DEBUG(llvm::dbgs() << "  Temp use may write/destroy its source" << *user);
+      return false;
+
+    case ValueKind::StructElementAddrInst:
+    case ValueKind::TupleElementAddrInst:
+      // Transitively look through projections on stack addresses.
+      for (auto *useOper : user->getUses()) {
+        if (!collectLoads(useOper->getUser(), user, loadInsts))
+          return false;
+      }
+      return true;
+
+    case ValueKind::LoadInst:
+    case ValueKind::LoadBorrowInst:
+      // Loads are the end of the data flow chain. The users of the load can't
+      // access the temporary storage.
+      loadInsts.insert(user);
+      return true;
+
+    case ValueKind::CopyAddrInst: {
+      // copy_addr which read from the temporary are like loads.
+      // TODO: Handle copy_addr [take]. But this doesn't seem to be important.
+      auto *copyFromTmp = cast<CopyAddrInst>(user);
+      if (copyFromTmp->getDest() == address || copyFromTmp->isTakeOfSrc()) {
+        DEBUG(llvm::dbgs() << "  Temp written or taken" << *user);
+        return false;
+      }
+      loadInsts.insert(user);
+      return true;
+    }
+  }
+}
+
+/// Checks if the copy's source can be modified within the temporary's lifetime.
+///
+/// Unfortunately, we cannot simply use the destroy points as the lifetime end,
+/// because they can be in a different basic block (that's what SILGen
+/// generates). Instead we guarantee that all normal uses are within the block
+/// of the temporary and look for the last use, which effectively ends the
+/// lifetime.
+bool TempRValueOptPass::checkNoSourceModification(CopyAddrInst *copyInst,
+                     const llvm::SmallPtrSetImpl<SILInstruction *> &useInsts) {
+  unsigned NumLoadsFound = 0;
+  auto iter = std::next(copyInst->getIterator());
+  // We already checked that the useful lifetime of the temporary ends in
+  // the initialization block.
+  auto iterEnd = copyInst->getParent()->end();
+  for (; iter != iterEnd; ++iter) {
+    SILInstruction *I = &*iter;
+
+    if (useInsts.count(I))
+      NumLoadsFound++;
+
+    // If this is the last use of the temp we are ok. After this point,
+    // modifications to the source don't matter anymore.
+    if (NumLoadsFound == useInsts.size())
+      return true;
+
+    if (AA->mayWriteToMemory(I, copyInst->getSrc())) {
+      DEBUG(llvm::dbgs() << "  Source modified by" << *iter);
+      return false;
+    }
+  }
+  // For some reason, not all normal uses have been seen between the copy and
+  // the end of the initialization block. We should never reach here.
+  return false;
+}
+
+/// Tries to perform the temporary rvalue copy elimination for \p copyInst
+bool TempRValueOptPass::tryOptimizeCopyIntoTemp(CopyAddrInst *copyInst) {
+  if (copyInst->isTakeOfSrc() || !copyInst->isInitializationOfDest())
+    return false;
+
+  auto *tempObj = dyn_cast<AllocStackInst>(copyInst->getDest());
+  if (!tempObj)
+    return false;
+
+  assert(tempObj != copyInst->getSrc() &&
+           "can't initialize temporary with itself");
+
+  // Scan all uses of the temporary storage (tempObj) to verify they all refer
+  // to the value initialized by this copy. It is sufficient to check that the
+  // only users that modify memory are the copy_addr [initialization] and
+  // destroy_addr.
+  llvm::SmallPtrSet<SILInstruction *, 8> loadInsts;
+  for (auto *useOper : tempObj->getUses()) {
+    SILInstruction *user = useOper->getUser();
+
+    if (user == copyInst)
+      continue;
+
+    // Destroys and deallocations are allowed to be in a different block.
+    if (isa<DestroyAddrInst>(user) || isa<DeallocStackInst>(user))
+      continue;
+
+    if (!collectLoads(user, tempObj, loadInsts))
+      return false;
+  }
+
+  // Check if the source is modified within the lifetime of the temporary.
+  if (!checkNoSourceModification(copyInst, loadInsts))
+    return false;
+
+  DEBUG(llvm::dbgs() << "  Success: replace temp" << *tempObj);
+
+  // Do a "replaceAllUses" by either deleting the users or replacing them with
+  // the source address. Note: we must not delete the original copyInst because
+  // it would crash the instruction iteration in run(). Instead the copyInst
+  // gets identical Src and Dest operands.
+  while (!tempObj->use_empty()) {
+    Operand *use = *tempObj->use_begin();
+    SILInstruction *user = use->getUser();
+    switch (user->getKind()) {
+      case ValueKind::DestroyAddrInst:
+      case ValueKind::DeallocStackInst:
+        user->eraseFromParent();
+        break;
+      case ValueKind::CopyAddrInst:
+      case ValueKind::StructElementAddrInst:
+      case ValueKind::TupleElementAddrInst:
+      case ValueKind::LoadInst:
+      case ValueKind::LoadBorrowInst:
+        use->set(copyInst->getSrc());
+        break;
+
+      default:
+        llvm_unreachable("unhandled instruction");
+    }
+  }
+  tempObj->eraseFromParent();
+  return true;
+}
+
 } // end anonymous namespace
 
 SILTransform *swift::createCopyForwarding() {
   return new CopyForwardingPass();
 }
+
+SILTransform *swift::createTempRValueOpt() {
+  return new TempRValueOptPass();
+}