Merge pull request swiftlang#34996 from meg-gupta/ossarle

Enable RLE on OSSA

Author: meg-gupta

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -20,6 +20,7 @@
 #ifndef SWIFT_SILOPTIMIZER_UTILS_INSTOPTUTILS_H
 #define SWIFT_SILOPTIMIZER_UTILS_INSTOPTUTILS_H
 
+#include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SILOptimizer/Analysis/ARCAnalysis.h"
@@ -688,6 +689,20 @@ SILBasicBlock::iterator replaceAllUsesAndErase(SingleValueInstruction *svi,
 SILBasicBlock::iterator replaceSingleUse(Operand *use, SILValue newValue,
                                          InstModCallbacks &callbacks);
 
+/// Creates a copy of \p value and inserts additional control equivalent copy
+/// and destroy at leaking blocks to adjust ownership and make available for use
+/// at \p inBlock.
+SILValue
+makeCopiedValueAvailable(SILValue value, SILBasicBlock *inBlock,
+                         JointPostDominanceSetComputer *jointPostDomComputer);
+
+/// Given a newly created @owned value \p value without any uses, this utility
+/// inserts control equivalent copy and destroy at leaking blocks to adjust
+/// ownership and make \p value available for use at \p inBlock.
+SILValue
+makeNewValueAvailable(SILValue value, SILBasicBlock *inBlock,
+                      JointPostDominanceSetComputer *jointPostDomComputer);
+
 } // end namespace swift
 
 #endif

include/swift/SILOptimizer/Utils/LoadStoreOptUtils.h

@@ -255,10 +255,27 @@ class LSValue : public LSBase {
   ///
   /// In the case where we have a single value this can be materialized by
   /// applying Path to the Base.
-  SILValue materialize(SILInstruction *Inst) {
+  SILValue materialize(SILInstruction *Inst,
+                       JointPostDominanceSetComputer *jointPostDomComputer) {
     if (CoveringValue)
       return SILValue();
-    return Path.getValue().createExtract(Base, Inst, true);
+    auto Val = Base;
+    auto InsertPt = getInsertAfterPoint(Base).getValue();
+    SILBuilderWithScope Builder(InsertPt);
+    if (Inst->getFunction()->hasOwnership() && !Path.getValue().empty()) {
+      // We have to create a @guaranteed scope with begin_borrow in order to
+      // create a struct_extract in OSSA
+      Val = Builder.emitBeginBorrowOperation(InsertPt->getLoc(), Base);
+    }
+    auto Res = Path.getValue().createExtract(Val, &*InsertPt, true);
+    if (Val != Base) {
+      Res = makeCopiedValueAvailable(Res, Inst->getParentBlock(),
+                                     jointPostDomComputer);
+      Builder.emitEndBorrowOperation(InsertPt->getLoc(), Val);
+      // Insert a destroy on the Base
+      SILBuilderWithScope(Inst).emitDestroyValueOperation(Inst->getLoc(), Base);
+    }
+    return Res;
   }
 
   void print(llvm::raw_ostream &os) override {
@@ -279,9 +296,11 @@ class LSValue : public LSBase {
   /// location holds. This may involve extracting and aggregating available
   /// values.
   static void reduceInner(LSLocation &B, SILModule *M, LSLocationValueMap &Vals,
-                          SILInstruction *InsertPt);
+                          SILInstruction *InsertPt,
+                          JointPostDominanceSetComputer *jointPostDomComputer);
   static SILValue reduce(LSLocation &B, SILModule *M, LSLocationValueMap &Vals,
-                         SILInstruction *InsertPt);
+                         SILInstruction *InsertPt,
+                         JointPostDominanceSetComputer *jointPostDomComputer);
 };
 
 static inline llvm::hash_code hash_value(const LSValue &V) {

lib/SIL/Utils/BasicBlockUtils.cpp

@@ -463,6 +463,13 @@ void JointPostDominanceSetComputer::findJointPostDominatingSet(
     for (auto *predBlock : block->getPredecessorBlocks()) {
       if (initialBlocks.count(predBlock)) {
         reachableInputBlocks.push_back(predBlock);
+        for (auto *succBlock : predBlock->getSuccessorBlocks()) {
+          if (visitedBlocks.count(succBlock))
+            continue;
+          if (deadEndBlocks.isDeadEnd(succBlock))
+            continue;
+          blocksThatLeakIfNeverVisited.insert(succBlock);
+        }
       }
       if (visitedBlocks.insert(predBlock).second)
         worklist.push_back(predBlock);

lib/SIL/Utils/Projection.cpp

@@ -384,7 +384,8 @@ Optional<ProjectionPath> ProjectionPath::getProjectionPath(SILValue Start,
     //
     // TODO: migrate users to getProjectionPath to the AccessPath utility to
     // avoid this hack.
-    if (!isa<EndCOWMutationInst>(Iter) && !isa<BeginAccessInst>(Iter)) {
+    if (!isa<EndCOWMutationInst>(Iter) && !isa<BeginAccessInst>(Iter) &&
+        !isa<BeginBorrowInst>(Iter)) {
       Projection AP(Iter);
       if (!AP.isValid())
         break;

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -385,6 +385,8 @@ void addFunctionPasses(SILPassPipelinePlan &P,
   } else {
     P.addRedundantLoadElimination();
   }
+  // Optimize copies created during RLE.
+  P.addSemanticARCOpts();
 
   P.addCOWOpts();
   P.addPerformanceConstantPropagation();

lib/SILOptimizer/Transforms/RedundantLoadElimination.cpp

@@ -155,6 +155,8 @@ static bool isRLEInertInstruction(SILInstruction *Inst) {
   case SILInstructionKind::EndAccessInst:
   case SILInstructionKind::SetDeallocatingInst:
   case SILInstructionKind::DeallocRefInst:
+  case SILInstructionKind::BeginBorrowInst:
+  case SILInstructionKind::EndBorrowInst:
     return true;
   default:
     return false;
@@ -481,14 +483,17 @@ class RLEContext {
   /// If set, RLE ignores loads from that array type.
   NominalTypeDecl *ArrayType;
 
+  JointPostDominanceSetComputer &jointPostDomComputer;
+
 #ifndef NDEBUG
   SILPrintContext printCtx;
 #endif
 
 public:
   RLEContext(SILFunction *F, SILPassManager *PM, AliasAnalysis *AA,
              TypeExpansionAnalysis *TE, PostOrderFunctionInfo *PO,
-             EpilogueARCFunctionInfo *EAFI, bool disableArrayLoads);
+             EpilogueARCFunctionInfo *EAFI, bool disableArrayLoads,
+             JointPostDominanceSetComputer &computer);
 
   RLEContext(const RLEContext &) = delete;
   RLEContext(RLEContext &&) = delete;
@@ -536,6 +541,11 @@ class RLEContext {
   /// Return the BlockState for the basic block this basic block belongs to.
   BlockState &getBlockState(SILBasicBlock *B) { return BBToLocState[B]; }
 
+  /// Return the initialized jointPostDomComputer
+  JointPostDominanceSetComputer *getJointPostDomSetComputer() {
+    return &jointPostDomComputer;
+  }
+
   /// Get the bit representing the LSLocation in the LocationVault.
   unsigned getLocationBit(const LSLocation &L);
 
@@ -679,14 +689,17 @@ SILValue BlockState::reduceValuesAtEndOfBlock(RLEContext &Ctx, LSLocation &L) {
   // we do not have a concrete value in the current basic block.
   ValueTableMap &OTM = getForwardValOut();
   for (unsigned i = 0; i < Locs.size(); ++i) {
-    Values[Locs[i]] = Ctx.getValue(OTM[Ctx.getLocationBit(Locs[i])]);
+    auto Val = Ctx.getValue(OTM[Ctx.getLocationBit(Locs[i])]);
+    auto AvailVal = makeCopiedValueAvailable(Val.getBase(), BB,
+                                             Ctx.getJointPostDomSetComputer());
+    Values[Locs[i]] = LSValue(AvailVal, Val.getPath().getValue());
   }
 
   // Second, reduce the available values into a single SILValue we can use to
   // forward.
-  SILValue TheForwardingValue;
-  TheForwardingValue =
-      LSValue::reduce(L, &BB->getModule(), Values, BB->getTerminator());
+  SILValue TheForwardingValue =
+      LSValue::reduce(L, &BB->getModule(), Values, BB->getTerminator(),
+                      Ctx.getJointPostDomSetComputer());
   /// Return the forwarding value.
   return TheForwardingValue;
 }
@@ -711,7 +724,9 @@ bool BlockState::setupRLE(RLEContext &Ctx, SILInstruction *I, SILValue Mem) {
 
   // Reduce the available values into a single SILValue we can use to forward.
   SILModule *Mod = &I->getModule();
-  SILValue TheForwardingValue = LSValue::reduce(L, Mod, Values, I);
+  SILValue TheForwardingValue =
+      LSValue::reduce(L, Mod, Values, I, Ctx.getJointPostDomSetComputer());
+
   if (!TheForwardingValue)
     return false;
 
@@ -849,11 +864,11 @@ void BlockState::processWrite(RLEContext &Ctx, SILInstruction *I, SILValue Mem,
     return;
   }
 
+  auto *Fn = I->getFunction();
   // Expand the given location and val into individual fields and process
   // them as separate writes.
   LSLocationList Locs;
-  LSLocation::expand(L, &I->getModule(),
-                     TypeExpansionContext(*I->getFunction()), Locs,
+  LSLocation::expand(L, &I->getModule(), TypeExpansionContext(*Fn), Locs,
                      Ctx.getTE());
 
   if (isComputeAvailSetMax(Kind)) {
@@ -873,8 +888,8 @@ void BlockState::processWrite(RLEContext &Ctx, SILInstruction *I, SILValue Mem,
 
   // Are we computing available value or performing RLE?
   LSValueList Vals;
-  LSValue::expand(Val, &I->getModule(), TypeExpansionContext(*I->getFunction()),
-                  Vals, Ctx.getTE());
+  LSValue::expand(Val, &I->getModule(), TypeExpansionContext(*Fn), Vals,
+                  Ctx.getTE());
   if (isComputeAvailValue(Kind) || isPerformingRLE(Kind)) {
     for (unsigned i = 0; i < Locs.size(); ++i) {
       updateForwardSetAndValForWrite(Ctx, Ctx.getLocationBit(Locs[i]),
@@ -903,11 +918,11 @@ void BlockState::processRead(RLEContext &Ctx, SILInstruction *I, SILValue Mem,
   if (!L.isValid())
     return;
 
+  auto *Fn = I->getFunction();
   // Expand the given LSLocation and Val into individual fields and process
   // them as separate reads.
   LSLocationList Locs;
-  LSLocation::expand(L, &I->getModule(),
-                     TypeExpansionContext(*I->getFunction()), Locs,
+  LSLocation::expand(L, &I->getModule(), TypeExpansionContext(*Fn), Locs,
                      Ctx.getTE());
 
   if (isComputeAvailSetMax(Kind)) {
@@ -928,8 +943,8 @@ void BlockState::processRead(RLEContext &Ctx, SILInstruction *I, SILValue Mem,
   // Are we computing available values ?.
   bool CanForward = true;
   LSValueList Vals;
-  LSValue::expand(Val, &I->getModule(), TypeExpansionContext(*I->getFunction()),
-                  Vals, Ctx.getTE());
+  LSValue::expand(Val, &I->getModule(), TypeExpansionContext(*Fn), Vals,
+                  Ctx.getTE());
   if (isComputeAvailValue(Kind) || isPerformingRLE(Kind)) {
     for (unsigned i = 0; i < Locs.size(); ++i) {
       if (isTrackingLocation(ForwardSetIn, Ctx.getLocationBit(Locs[i])))
@@ -1194,10 +1209,13 @@ void BlockState::dump(RLEContext &Ctx) {
 
 RLEContext::RLEContext(SILFunction *F, SILPassManager *PM, AliasAnalysis *AA,
                        TypeExpansionAnalysis *TE, PostOrderFunctionInfo *PO,
-                       EpilogueARCFunctionInfo *EAFI, bool disableArrayLoads)
+                       EpilogueARCFunctionInfo *EAFI, bool disableArrayLoads,
+                       JointPostDominanceSetComputer &computer)
     : Fn(F), PM(PM), AA(AA), TE(TE), PO(PO), EAFI(EAFI),
-      ArrayType(disableArrayLoads ?
-                F->getModule().getASTContext().getArrayDecl() : nullptr)
+      ArrayType(disableArrayLoads
+                    ? F->getModule().getASTContext().getArrayDecl()
+                    : nullptr),
+      jointPostDomComputer(computer)
 #ifndef NDEBUG
       ,
       printCtx(llvm::dbgs(), /*Verbose=*/false, /*Sorted=*/true)
@@ -1322,8 +1340,8 @@ SILValue RLEContext::computePredecessorLocationValue(SILBasicBlock *BB,
 
     // Reduce the available values into a single SILValue we can use to forward
     SILInstruction *IPt = CurBB->getTerminator();
-    Values.push_back(
-        {CurBB, LSValue::reduce(L, &BB->getModule(), LSValues, IPt)});
+    Values.push_back({CurBB, LSValue::reduce(L, &BB->getModule(), LSValues, IPt,
+                                             &jointPostDomComputer)});
   }
 
   // Finally, collect all the values for the SILArgument, materialize it using
@@ -1335,7 +1353,8 @@ SILValue RLEContext::computePredecessorLocationValue(SILBasicBlock *BB,
     Updater.addAvailableValue(V.first, V.second);
   }
 
-  return Updater.getValueInMiddleOfBlock(BB);
+  auto Val = Updater.getValueInMiddleOfBlock(BB);
+  return makeNewValueAvailable(Val, BB, &jointPostDomComputer);
 }
 
 bool RLEContext::collectLocationValues(SILBasicBlock *BB, LSLocation &L,
@@ -1350,9 +1369,14 @@ bool RLEContext::collectLocationValues(SILBasicBlock *BB, LSLocation &L,
   // Find the locations that this basic block defines and the locations which
   // we do not have a concrete value in the current basic block.
   for (auto &X : Locs) {
-    Values[X] = getValue(VM[getLocationBit(X)]);
-    if (!Values[X].isCoveringValue())
+    auto Val = getValue(VM[getLocationBit(X)]);
+    if (!Val.isCoveringValue()) {
+      auto AvailValue =
+          makeCopiedValueAvailable(Val.getBase(), BB, &jointPostDomComputer);
+      Values[X] = LSValue(AvailValue, Val.getPath().getValue());
       continue;
+    }
+    Values[X] = Val;
     CSLocs.push_back(X);
   }
 
@@ -1367,7 +1391,6 @@ bool RLEContext::collectLocationValues(SILBasicBlock *BB, LSLocation &L,
     SILValue V = computePredecessorLocationValue(BB, X);
     if (!V)
       return false;
-
     // We've constructed a concrete value for the covering value. Expand and
     // collect the newly created forwardable values.
     LSLocationList Locs;
@@ -1610,9 +1633,15 @@ bool RLEContext::run() {
         continue;
       LLVM_DEBUG(llvm::dbgs() << "Replacing  " << SILValue(Iter->first)
                               << "With " << Iter->second);
+      auto *origLoad = cast<LoadInst>(Iter->first);
+      SILValue newValue = Iter->second;
+      if (origLoad->getOwnershipQualifier() == LoadOwnershipQualifier::Take) {
+        SILBuilderWithScope(origLoad).createDestroyAddr(origLoad->getLoc(),
+                                                        origLoad->getOperand());
+      }
       SILChanged = true;
-      Iter->first->replaceAllUsesWith(Iter->second);
-      InstsToDelete.push_back(Iter->first);
+      origLoad->replaceAllUsesWith(newValue);
+      InstsToDelete.push_back(origLoad);
       ++NumForwardedLoads;
     }
   }
@@ -1650,9 +1679,6 @@ class RedundantLoadElimination : public SILFunctionTransform {
   /// The entry point to the transformation.
   void run() override {
     SILFunction *F = getFunction();
-    // FIXME: Handle ownership.
-    if (F->hasOwnership())
-      return;
 
     LLVM_DEBUG(llvm::dbgs() << "*** RLE on function: " << F->getName()
                             << " ***\n");
@@ -1662,7 +1688,9 @@ class RedundantLoadElimination : public SILFunctionTransform {
     auto *PO = PM->getAnalysis<PostOrderAnalysis>()->get(F);
     auto *EAFI = PM->getAnalysis<EpilogueARCAnalysis>()->get(F);
 
-    RLEContext RLE(F, PM, AA, TE, PO, EAFI, disableArrayLoads);
+    DeadEndBlocks deadEndBlocks(F);
+    JointPostDominanceSetComputer computer(deadEndBlocks);
+    RLEContext RLE(F, PM, AA, TE, PO, EAFI, disableArrayLoads, computer);
     if (RLE.run()) {
       invalidateAnalysis(SILAnalysis::InvalidationKind::Instructions);
     }

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -1963,3 +1963,53 @@ SILBasicBlock::iterator swift::replaceSingleUse(Operand *use, SILValue newValue,
 
   return nextII;
 }
+
+SILValue swift::makeCopiedValueAvailable(
+    SILValue value, SILBasicBlock *inBlock,
+    JointPostDominanceSetComputer *jointPostDomComputer) {
+  if (!value->getFunction()->hasOwnership())
+    return value;
+
+  if (value->getType().isTrivial(*value->getFunction()))
+    return value;
+
+  auto insertPt = getInsertAfterPoint(value).getValue();
+  auto *copy =
+      SILBuilderWithScope(insertPt).createCopyValue(insertPt->getLoc(), value);
+
+  return makeNewValueAvailable(copy, inBlock, jointPostDomComputer);
+}
+
+SILValue swift::makeNewValueAvailable(
+    SILValue value, SILBasicBlock *inBlock,
+    JointPostDominanceSetComputer *jointPostDomComputer) {
+  if (!value->getFunction()->hasOwnership())
+    return value;
+
+  if (value->getType().isTrivial(*value->getFunction()))
+    return value;
+
+  assert(value->getUses().empty() &&
+         value.getOwnershipKind() == OwnershipKind::Owned);
+
+  // Use \p jointPostDomComputer to:
+  // 1. Create a control equivalent copy at \p inBlock if needed
+  // 2. Insert destroy_value at leaking blocks
+  SILValue controlEqCopy;
+  jointPostDomComputer->findJointPostDominatingSet(
+      value->getParentBlock(), inBlock,
+      [&](SILBasicBlock *loopBlock) {
+        assert(loopBlock == inBlock);
+        auto front = loopBlock->begin();
+        SILBuilderWithScope newBuilder(front);
+        controlEqCopy = newBuilder.createCopyValue(front->getLoc(), value);
+      },
+      [&](SILBasicBlock *postDomBlock) {
+        // Insert a destroy_value in the leaking block
+        auto front = postDomBlock->begin();
+        SILBuilderWithScope newBuilder(front);
+        newBuilder.createDestroyValue(front->getLoc(), value);
+      });
+
+  return controlEqCopy ? controlEqCopy : value;
+}