[sil-combine] Use canonicalizeOSSALifetimes to eliminate unnecessary copies inserted due to RAUWing

include/swift/SIL/SILInstructionWorklist.h

@@ -128,6 +128,30 @@ class SILInstructionWorklist : SILInstructionWorklistBase {
     }
   }
 
+  /// All operands of \p instruction to the worklist when performing 2 stage
+  /// instruction deletion. Meant to be used right before deleting an
+  /// instruction in callbacks like InstModCallback::onNotifyWillBeDeleted().
+  void addOperandsToWorklist(SILInstruction &instruction) {
+    assert(!instruction.hasUsesOfAnyResult() &&
+           "Cannot erase instruction that is used!");
+
+    // Make sure that we reprocess all operands now that we reduced their
+    // use counts.
+    if (instruction.getNumOperands() < 8) {
+      for (auto &operand : instruction.getAllOperands()) {
+        if (auto *operandInstruction =
+                operand.get()->getDefiningInstruction()) {
+          withDebugStream([&](llvm::raw_ostream &stream,
+                              StringRef loggingName) {
+            stream << loggingName << ": add op " << *operandInstruction << '\n'
+                   << " from erased inst to worklist\n";
+          });
+          add(operandInstruction);
+        }
+      }
+    }
+  }
+
   /// When an instruction has been simplified, add all of its users to the
   /// worklist, since additional simplifications of its users may have been
   /// exposed.
@@ -296,29 +320,15 @@ class SILInstructionWorklist : SILInstructionWorklistBase {
   }
 
   void eraseSingleInstFromFunction(SILInstruction &instruction,
-                                   bool addOperandsToWorklist) {
+                                   bool shouldAddOperandsToWorklist) {
     withDebugStream([&](llvm::raw_ostream &stream, StringRef loggingName) {
       stream << loggingName << ": ERASE " << instruction << '\n';
     });
 
-    assert(!instruction.hasUsesOfAnyResult() &&
-           "Cannot erase instruction that is used!");
+    // If we are asked to add operands to the worklist, do so now.
+    if (shouldAddOperandsToWorklist)
+      addOperandsToWorklist(instruction);
 
-    // Make sure that we reprocess all operands now that we reduced their
-    // use counts.
-    if (instruction.getNumOperands() < 8 && addOperandsToWorklist) {
-      for (auto &operand : instruction.getAllOperands()) {
-        if (auto *operandInstruction =
-                operand.get()->getDefiningInstruction()) {
-          withDebugStream([&](llvm::raw_ostream &stream,
-                              StringRef loggingName) {
-            stream << loggingName << ": add op " << *operandInstruction << '\n'
-                   << " from erased inst to worklist\n";
-          });
-          add(operandInstruction);
-        }
-      }
-    }
     erase(&instruction);
     instruction.eraseFromParent();
   }

include/swift/SILOptimizer/Utils/CanonicalOSSALifetime.h

@@ -398,6 +398,10 @@ class CanonicalizeOSSALifetime {
   /// lifetime, call this API again on the value defined by the new copy.
   bool canonicalizeValueLifetime(SILValue def);
 
+  /// Return the inst mod callbacks struct used by this CanonicalizeOSSALifetime
+  /// to pass to other APIs that need to compose with CanonicalizeOSSALifetime.
+  InstModCallbacks getInstModCallbacks() const { return instModCallbacks; }
+
 protected:
   void recordDebugValue(DebugValueInst *dvi) {
     debugValues.insert(dvi);
@@ -438,6 +442,29 @@ class CanonicalizeOSSALifetime {
   void injectPoison();
 };
 
+/// Canonicalize the passed in set of defs, eliminating in one batch any that
+/// are not needed given the canonical lifetime of the various underlying owned
+/// value introducers.
+///
+/// On success, returns the invalidation kind that the caller must use to
+/// invalidate analyses. Currently it will only ever return
+/// SILAnalysis::InvalidationKind::Instructions or None.
+///
+/// NOTE: This routine is guaranteed to not invalidate
+/// NonLocalAccessBlockAnalysis, so callers should lock it before invalidating
+/// instructions. E.x.:
+///
+///    accessBlockAnalysis->lockInvalidation();
+///    pass->invalidateAnalysis(SILAnalysis::InvalidationKind::Instructions);
+///    accessBlockAnalysis->unlockInvalidation();
+///
+/// NOTE: We assume that all \p inputDefs is a set (that is it has no duplicate
+/// elements) and that all values have been generated by running a copy through
+/// CanonicalizeOSSALifetime::getCanonicalCopiedDef(copy)).
+SILAnalysis::InvalidationKind
+canonicalizeOSSALifetimes(CanonicalizeOSSALifetime &canonicalizeLifetime,
+                          ArrayRef<SILValue> inputDefs);
+
 } // end namespace swift
 
 #endif

lib/SILOptimizer/SILCombiner/SILCombine.cpp

@@ -21,14 +21,17 @@
 #define DEBUG_TYPE "sil-combine"
 
 #include "SILCombiner.h"
+#include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SIL/DebugUtils.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILVisitor.h"
-#include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SILOptimizer/Analysis/AliasAnalysis.h"
+#include "swift/SILOptimizer/Analysis/DominanceAnalysis.h"
+#include "swift/SILOptimizer/Analysis/NonLocalAccessBlockAnalysis.h"
 #include "swift/SILOptimizer/Analysis/SimplifyInstruction.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
+#include "swift/SILOptimizer/Utils/CanonicalOSSALifetime.h"
 #include "swift/SILOptimizer/Utils/CanonicalizeInstruction.h"
 #include "swift/SILOptimizer/Utils/DebugOptUtils.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
@@ -299,10 +302,39 @@ bool SILCombiner::doOneIteration(SILFunction &F, unsigned Iteration) {
     }
 
     // Our tracking list has been accumulating instructions created by the
-    // SILBuilder during this iteration. Go through the tracking list and add
-    // its contents to the worklist and then clear said list in preparation for
-    // the next iteration.
+    // SILBuilder during this iteration. In order to finish this round of
+    // SILCombine, go through the tracking list and add its contents to the
+    // worklist and then clear said list in preparation for the next
+    // iteration. We canonicalize any copies that we created in order to
+    // eliminate unnecessary copies introduced by RAUWing when ownership is
+    // enabled.
+    //
+    // NOTE: It is ok if copy propagation results in MadeChanges being set to
+    // true. This is because we only add elements to the tracking list if we
+    // actually made a change to the IR, so MadeChanges should already be true
+    // at this point.
     auto &TrackingList = *Builder.getTrackingList();
+    if (TrackingList.size() && Builder.hasOwnership()) {
+      SmallSetVector<SILValue, 16> defsToCanonicalize;
+      for (auto *trackedInst : TrackingList) {
+        if (auto *cvi = dyn_cast<CopyValueInst>(trackedInst)) {
+          defsToCanonicalize.insert(
+              CanonicalizeOSSALifetime::getCanonicalCopiedDef(cvi));
+        }
+      }
+      if (defsToCanonicalize.size()) {
+        CanonicalizeOSSALifetime canonicalizer(
+            false /*prune debug*/, false /*canonicalize borrows*/,
+            false /*poison refs*/, NLABA, DA, instModCallbacks);
+        auto analysisInvalidation = canonicalizeOSSALifetimes(
+            canonicalizer, defsToCanonicalize.getArrayRef());
+        if (bool(analysisInvalidation)) {
+          NLABA->lockInvalidation();
+          parentTransform->invalidateAnalysis(analysisInvalidation);
+          NLABA->unlockInvalidation();
+        }
+      }
+    }
     for (SILInstruction *I : TrackingList) {
       LLVM_DEBUG(llvm::dbgs() << "SC: add " << *I
                               << " from tracking list to worklist\n");
@@ -359,12 +391,13 @@ class SILCombine : public SILFunctionTransform {
     auto *DA = PM->getAnalysis<DominanceAnalysis>();
     auto *PCA = PM->getAnalysis<ProtocolConformanceAnalysis>();
     auto *CHA = PM->getAnalysis<ClassHierarchyAnalysis>();
+    auto *NLABA = PM->getAnalysis<NonLocalAccessBlockAnalysis>();
 
     SILOptFunctionBuilder FuncBuilder(*this);
     // Create a SILBuilder with a tracking list for newly added
     // instructions, which we will periodically move to our worklist.
     SILBuilder B(*getFunction(), &TrackingList);
-    SILCombiner Combiner(FuncBuilder, B, AA, DA, PCA, CHA,
+    SILCombiner Combiner(this, FuncBuilder, B, AA, DA, PCA, CHA, NLABA,
                          getOptions().RemoveRuntimeAsserts);
     bool Changed = Combiner.runOnFunction(*getFunction());
     assert(TrackingList.empty() &&

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -29,6 +29,7 @@
 #include "swift/SIL/SILValue.h"
 #include "swift/SIL/SILVisitor.h"
 #include "swift/SILOptimizer/Analysis/ClassHierarchyAnalysis.h"
+#include "swift/SILOptimizer/Analysis/NonLocalAccessBlockAnalysis.h"
 #include "swift/SILOptimizer/Analysis/ProtocolConformanceAnalysis.h"
 #include "swift/SILOptimizer/Utils/CastOptimizer.h"
 #include "swift/SILOptimizer/Utils/Existential.h"
@@ -46,6 +47,7 @@ class AliasAnalysis;
 /// the worklist.
 class SILCombiner :
     public SILInstructionVisitor<SILCombiner, SILInstruction *> {
+  SILFunctionTransform *parentTransform;
 
   AliasAnalysis *AA;
 
@@ -59,6 +61,10 @@ class SILCombiner :
   /// conforming class.
   ClassHierarchyAnalysis *CHA;
 
+  /// Non local access block analysis that we use when canonicalize object
+  /// lifetimes in OSSA.
+  NonLocalAccessBlockAnalysis *NLABA;
+
   /// Worklist containing all of the instructions primed for simplification.
   SmallSILInstructionWorklist<256> Worklist;
 
@@ -97,39 +103,47 @@ class SILCombiner :
   OwnershipFixupContext ownershipFixupContext;
 
 public:
-  SILCombiner(SILOptFunctionBuilder &FuncBuilder, SILBuilder &B,
+  SILCombiner(SILFunctionTransform *parentTransform,
+              SILOptFunctionBuilder &FuncBuilder, SILBuilder &B,
               AliasAnalysis *AA, DominanceAnalysis *DA,
               ProtocolConformanceAnalysis *PCA, ClassHierarchyAnalysis *CHA,
-              bool removeCondFails)
-      : AA(AA), DA(DA), PCA(PCA), CHA(CHA), Worklist("SC"),
-        deadEndBlocks(&B.getFunction()), MadeChange(false),
-        RemoveCondFails(removeCondFails), Iteration(0), Builder(B),
-        CastOpt(
-            FuncBuilder, nullptr /*SILBuilderContext*/,
-            /* ReplaceValueUsesAction */
-            [&](SILValue Original, SILValue Replacement) {
-              replaceValueUsesWith(Original, Replacement);
-            },
-            /* ReplaceInstUsesAction */
-            [&](SingleValueInstruction *I, ValueBase *V) {
-              replaceInstUsesWith(*I, V);
-            },
-            /* EraseAction */
-            [&](SILInstruction *I) { eraseInstFromFunction(*I); }),
-        instModCallbacks(),
-        deBlocks(&B.getFunction()),
+              NonLocalAccessBlockAnalysis *NLABA, bool removeCondFails)
+      : parentTransform(parentTransform), AA(AA), DA(DA), PCA(PCA), CHA(CHA),
+        NLABA(NLABA), Worklist("SC"), deadEndBlocks(&B.getFunction()),
+        MadeChange(false), RemoveCondFails(removeCondFails), Iteration(0),
+        Builder(B), CastOpt(
+                        FuncBuilder, nullptr /*SILBuilderContext*/,
+                        /* ReplaceValueUsesAction */
+                        [&](SILValue Original, SILValue Replacement) {
+                          replaceValueUsesWith(Original, Replacement);
+                        },
+                        /* ReplaceInstUsesAction */
+                        [&](SingleValueInstruction *I, ValueBase *V) {
+                          replaceInstUsesWith(*I, V);
+                        },
+                        /* EraseAction */
+                        [&](SILInstruction *I) { eraseInstFromFunction(*I); }),
+        instModCallbacks(), deBlocks(&B.getFunction()),
         ownershipFixupContext(instModCallbacks, deBlocks) {
-    instModCallbacks = InstModCallbacks()
-      .onDelete([&](SILInstruction *instToDelete) {
-              eraseInstFromFunction(*instToDelete);
-      })
-      .onCreateNewInst([&](SILInstruction *newlyCreatedInst) {
-        Worklist.add(newlyCreatedInst);
-      })
-      .onSetUseValue([&](Operand *use, SILValue newValue) {
-        use->set(newValue);
-        Worklist.add(use->getUser());
-      });
+    instModCallbacks =
+        InstModCallbacks()
+            .onDelete([&](SILInstruction *instToDelete) {
+              // We allow for users in SILCombine to perform 2 stage deletion,
+              // so we need to split the erasing of instructions from adding
+              // operands to the worklist.
+              eraseInstFromFunction(
+                  *instToDelete, false /*do not add operands to the worklist*/);
+            })
+            .onNotifyWillBeDeleted([&](SILInstruction *instThatWillBeDeleted) {
+              Worklist.addOperandsToWorklist(*instThatWillBeDeleted);
+            })
+            .onCreateNewInst([&](SILInstruction *newlyCreatedInst) {
+              Worklist.add(newlyCreatedInst);
+            })
+            .onSetUseValue([&](Operand *use, SILValue newValue) {
+              use->set(newValue);
+              Worklist.add(use->getUser());
+            });
   }
 
   bool runOnFunction(SILFunction &F);

lib/SILOptimizer/Utils/CanonicalOSSALifetime.cpp

@@ -1104,7 +1104,8 @@ void CanonicalizeOSSALifetime::rewriteCopies() {
   // Remove the leftover copy_value and destroy_value instructions.
   if (!instsToDelete.empty()) {
     recursivelyDeleteTriviallyDeadInstructions(instsToDelete.takeVector(),
-                                               /*force=*/true);
+                                               /*force=*/true,
+                                               instModCallbacks);
     setChanged();
   }
 }
@@ -1220,3 +1221,64 @@ SWIFT_ASSERT_ONLY_DECL(
       llvm::dbgs() << "  " << *blockAndInst.getSecond();
     }
   })
+
+//===----------------------------------------------------------------------===//
+//                    MARK: Canonicalize Lifetimes Utility
+//===----------------------------------------------------------------------===//
+
+SILAnalysis::InvalidationKind
+swift::canonicalizeOSSALifetimes(CanonicalizeOSSALifetime &canonicalizer,
+                                 ArrayRef<SILValue> copiedDefs) {
+  auto isCopyDead = [](CopyValueInst *copy, bool pruneDebug) -> bool {
+    for (Operand *use : copy->getUses()) {
+      auto *user = use->getUser();
+      if (isa<DestroyValueInst>(user)) {
+        continue;
+      }
+      if (pruneDebug && isa<DebugValueInst>(user)) {
+        continue;
+      }
+      return false;
+    }
+    return true;
+  };
+
+  // Cleanup dead copies. If getCanonicalCopiedDef returns a copy (because the
+  // copy's source operand is unrecgonized), then the copy is itself treated
+  // like a def and may be dead after canonicalization.
+  SmallVector<SILInstruction *, 4> deadCopies;
+  for (auto def : copiedDefs) {
+    // Canonicalized this def. If we did not perform any canonicalization, just
+    // continue.
+    if (!canonicalizer.canonicalizeValueLifetime(def))
+      continue;
+
+    // Otherwise, see if we have a dead copy.
+    if (auto *copy = dyn_cast<CopyValueInst>(def)) {
+      if (isCopyDead(copy, false /*prune debug*/)) {
+        deadCopies.push_back(copy);
+      }
+    }
+
+    // Canonicalize any new outer copy.
+    if (SILValue outerCopy = canonicalizer.createdOuterCopy()) {
+      SILValue outerDef = canonicalizer.getCanonicalCopiedDef(outerCopy);
+      canonicalizer.canonicalizeValueLifetime(outerDef);
+    }
+
+    // TODO: also canonicalize any lifetime.persistentCopies like separate owned
+    // live ranges.
+  }
+
+  if (!canonicalizer.hasChanged() && deadCopies.empty()) {
+    return SILAnalysis::InvalidationKind::Nothing;
+  }
+
+  // We use our canonicalizers inst mod callbacks.
+  InstructionDeleter deleter(canonicalizer.getInstModCallbacks());
+  for (auto *copy : deadCopies) {
+    deleter.recursivelyDeleteUsersIfDead(copy);
+  }
+
+  return SILAnalysis::InvalidationKind::Instructions;
+}