Enable LoopRotate on OSSA

include/swift/SILOptimizer/Utils/SILSSAUpdater.h

@@ -49,6 +49,8 @@ class SILSSAUpdater {
 
   SILType type;
 
+  ValueOwnershipKind ownershipKind;
+
   // The SSAUpdaterTraits specialization uses this sentinel to mark 'new' phi
   // nodes (all the incoming edge arguments have this sentinel set).
   std::unique_ptr<SILUndef, void (*)(SILUndef *)> phiSentinel;
@@ -69,8 +71,8 @@ class SILSSAUpdater {
     insertedPhis = inputInsertedPhis;
   }
 
-  /// Initialize for a use of a value of type.
-  void initialize(SILType type);
+  /// Initialize for a use of a value of type and ownershipKind
+  void initialize(SILType type, ValueOwnershipKind ownershipKind);
 
   bool hasValueForBlock(SILBasicBlock *block) const;
   void addAvailableValue(SILBasicBlock *block, SILValue value);

lib/SILOptimizer/LoopTransforms/ArrayPropertyOpt.cpp

@@ -516,7 +516,7 @@ class RegionCloner : public SILCloner<RegionCloner> {
       return;
 
     // Update SSA form.
-    SSAUp.initialize(V->getType());
+    SSAUp.initialize(V->getType(), V.getOwnershipKind());
     SSAUp.addAvailableValue(OrigBB, V);
     SILValue NewVal = getMappedValue(V);
     SSAUp.addAvailableValue(getOpBasicBlock(OrigBB), NewVal);

lib/SILOptimizer/LoopTransforms/LICM.cpp

@@ -1326,7 +1326,8 @@ hoistLoadsAndStores(AccessPath accessPath, SILLoop *loop) {
 
       if (!storeAddr) {
         storeAddr = SI->getDest();
-        ssaUpdater.initialize(storeAddr->getType().getObjectType());
+        ssaUpdater.initialize(storeAddr->getType().getObjectType(),
+                              storeAddr.getOwnershipKind());
       } else if (SI->getDest()->getType() != storeAddr->getType()) {
         // This transformation assumes that the values of all stores in the loop
         // must be interchangeable. It won't work if stores different types

lib/SILOptimizer/LoopTransforms/LoopRotate.cpp

@@ -39,6 +39,8 @@ using namespace swift;
 /// loops in the swift benchmarks).
 static llvm::cl::opt<int> LoopRotateSizeLimit("looprotate-size-limit",
                                               llvm::cl::init(20));
+static llvm::cl::opt<bool> RotateSingleBlockLoop("looprotate-single-block-loop",
+                                                 llvm::cl::init(false));
 
 /// Check whether all operands are loop invariant.
 static bool
@@ -131,7 +133,7 @@ static void updateSSAForUseOfValue(
   assert(Res->getType() == MappedValue->getType() && "The types must match");
 
   insertedPhis.clear();
-  updater.initialize(Res->getType());
+  updater.initialize(Res->getType(), Res.getOwnershipKind());
   updater.addAvailableValue(Header, Res);
   updater.addAvailableValue(EntryCheckBlock, MappedValue);
 
@@ -228,9 +230,9 @@ static bool rotateLoopAtMostUpToLatch(SILLoop *loop, DominanceInfo *domInfo,
     return false;
   }
 
-  bool didRotate =
-      rotateLoop(loop, domInfo, loopInfo, false /* rotateSingleBlockLoops */,
-                 latch, ShouldVerify);
+  bool didRotate = rotateLoop(
+      loop, domInfo, loopInfo,
+      RotateSingleBlockLoop /* rotateSingleBlockLoops */, latch, ShouldVerify);
 
   // Keep rotating at most until we hit the original latch.
   if (didRotate)
@@ -440,9 +442,6 @@ class LoopRotation : public SILFunctionTransform {
 
     SILFunction *f = getFunction();
     assert(f);
-    // FIXME: Add ownership support.
-    if (f->hasOwnership())
-      return;
 
     SILLoopInfo *loopInfo = loopAnalysis->get(f);
     assert(loopInfo);

lib/SILOptimizer/LoopTransforms/LoopUnroll.cpp

@@ -335,7 +335,7 @@ updateSSA(SILModule &M, SILLoop *Loop,
       if (!Loop->contains(Use->getUser()->getParent()))
         UseList.push_back(UseWrapper(Use));
     // Update SSA of use with the available values.
-    SSAUp.initialize(OrigValue->getType());
+    SSAUp.initialize(OrigValue->getType(), OrigValue.getOwnershipKind());
     SSAUp.addAvailableValue(OrigValue->getParentBlock(), OrigValue);
     for (auto NewValue : MapEntry.second)
       SSAUp.addAvailableValue(NewValue->getParentBlock(), NewValue);

lib/SILOptimizer/Mandatory/ClosureLifetimeFixup.cpp

@@ -245,7 +245,9 @@ static void extendLifetimeToEndOfFunction(SILFunction &fn,
   // lifetime respecting loops.
   SmallVector<SILPhiArgument *, 8> insertedPhis;
   SILSSAUpdater updater(&insertedPhis);
-  updater.initialize(optionalEscapingClosureTy);
+  updater.initialize(optionalEscapingClosureTy, fn.hasOwnership()
+                                                    ? OwnershipKind::Owned
+                                                    : OwnershipKind::None);
 
   // Create an Optional<() -> ()>.none in the entry block of the function and
   // add it as an available value to the SSAUpdater.
@@ -850,7 +852,9 @@ static bool fixupCopyBlockWithoutEscaping(CopyBlockWithoutEscapingInst *cb,
 
   SmallVector<SILPhiArgument *, 8> insertedPhis;
   SILSSAUpdater updater(&insertedPhis);
-  updater.initialize(optionalEscapingClosureTy);
+  updater.initialize(optionalEscapingClosureTy, fn.hasOwnership()
+                                                    ? OwnershipKind::Owned
+                                                    : OwnershipKind::None);
 
   // Create the Optional.none as the beginning available value.
   SILValue entryBlockOptionalNone;

lib/SILOptimizer/Mandatory/PredictableMemOpt.cpp

@@ -684,7 +684,8 @@ AvailableValueAggregator::aggregateFullyAvailableValue(SILType loadTy,
   // have multiple insertion points if we are storing exactly the same value
   // implying that we can just copy firstVal at each insertion point.
   SILSSAUpdater updater(&insertedPhiNodes);
-  updater.initialize(loadTy);
+  updater.initialize(loadTy, B.hasOwnership() ? OwnershipKind::Owned
+                                              : OwnershipKind::None);
 
   Optional<SILValue> singularValue;
   for (auto *insertPt : insertPts) {
@@ -860,7 +861,8 @@ SILValue AvailableValueAggregator::handlePrimitiveValue(SILType loadTy,
   // never have the same value along all paths unless we have a trivial value
   // meaning the SSA updater given a non-trivial value must /always/ be used.
   SILSSAUpdater updater(&insertedPhiNodes);
-  updater.initialize(loadTy);
+  updater.initialize(loadTy, B.hasOwnership() ? OwnershipKind::Owned
+                                              : OwnershipKind::None);
 
   Optional<SILValue> singularValue;
   for (auto *i : insertPts) {

lib/SILOptimizer/Transforms/DeadCodeElimination.cpp

@@ -524,7 +524,7 @@ bool DCE::removeDead(SILFunction &F) {
 
       arg->replaceAllUsesWithUndef();
 
-      if (!F.hasOwnership() || arg->getType().isTrivial(F)) {
+      if (!F.hasOwnership() || arg->getOwnershipKind() == OwnershipKind::None) {
         i++;
         Changed = true;
         BranchesChanged = true;

lib/SILOptimizer/Transforms/DeadObjectElimination.cpp

@@ -559,7 +559,7 @@ static void insertReleases(ArrayRef<StoreInst*> Stores,
   assert(!Stores.empty());
   SILValue StVal = Stores.front()->getSrc();
 
-  SSAUp.initialize(StVal->getType());
+  SSAUp.initialize(StVal->getType(), StVal.getOwnershipKind());
 
   for (auto *Store : Stores)
     SSAUp.addAvailableValue(Store->getParent(), Store->getSrc());

lib/SILOptimizer/Transforms/RedundantLoadElimination.cpp

@@ -1349,7 +1349,9 @@ SILValue RLEContext::computePredecessorLocationValue(SILBasicBlock *BB,
   // the SSAUpdater.
   Updater.initialize(
       L.getType(&BB->getModule(), TypeExpansionContext(*BB->getParent()))
-          .getObjectType());
+          .getObjectType(),
+      Values[0].second.getOwnershipKind());
+
   for (auto V : Values) {
     Updater.addAvailableValue(V.first, V.second);
   }

lib/SILOptimizer/Utils/BasicBlockOptUtils.cpp

@@ -101,7 +101,7 @@ void BasicBlockCloner::updateSSAAfterCloning() {
 
   SILSSAUpdater ssaUpdater;
   for (auto availValPair : availVals) {
-    ValueBase *inst = availValPair.first;
+    auto inst = availValPair.first;
     if (inst->use_empty())
       continue;
 
@@ -112,7 +112,7 @@ void BasicBlockCloner::updateSSAAfterCloning() {
     for (auto *use : inst->getUses())
       useList.push_back(UseWrapper(use));
 
-    ssaUpdater.initialize(inst->getType());
+    ssaUpdater.initialize(inst->getType(), inst.getOwnershipKind());
     ssaUpdater.addAvailableValue(origBB, inst);
     ssaUpdater.addAvailableValue(getNewBB(), newResult);
 

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -1997,7 +1997,7 @@ SILValue swift::makeCopiedValueAvailable(
   if (!value->getFunction()->hasOwnership())
     return value;
 
-  if (value->getType().isTrivial(*value->getFunction()))
+  if (value.getOwnershipKind() == OwnershipKind::None)
     return value;
 
   auto insertPt = getInsertAfterPoint(value).getValue();
@@ -2013,7 +2013,7 @@ SILValue swift::makeNewValueAvailable(
   if (!value->getFunction()->hasOwnership())
     return value;
 
-  if (value->getType().isTrivial(*value->getFunction()))
+  if (value.getOwnershipKind() == OwnershipKind::None)
     return value;
 
   assert(value->getUses().empty() &&

lib/SILOptimizer/Utils/SILSSAUpdater.cpp

@@ -34,13 +34,14 @@ void SILSSAUpdater::deallocateSentinel(SILUndef *undef) {
 }
 
 SILSSAUpdater::SILSSAUpdater(SmallVectorImpl<SILPhiArgument *> *phis)
-    : blockToAvailableValueMap(nullptr),
+    : blockToAvailableValueMap(nullptr), ownershipKind(OwnershipKind::None),
       phiSentinel(nullptr, deallocateSentinel), insertedPhis(phis) {}
 
 SILSSAUpdater::~SILSSAUpdater() = default;
 
-void SILSSAUpdater::initialize(SILType inputType) {
+void SILSSAUpdater::initialize(SILType inputType, ValueOwnershipKind kind) {
   type = inputType;
+  ownershipKind = kind;
 
   phiSentinel = std::unique_ptr<SILUndef, void (*)(SILUndef *)>(
       SILUndef::getSentinelValue(inputType, this),
@@ -59,6 +60,7 @@ bool SILSSAUpdater::hasValueForBlock(SILBasicBlock *block) const {
 /// Indicate that a rewritten value is available in the specified block with the
 /// specified value.
 void SILSSAUpdater::addAvailableValue(SILBasicBlock *block, SILValue value) {
+  assert(value.getOwnershipKind().isCompatibleWith(ownershipKind));
   (*blockToAvailableValueMap)[block] = value;
 }
 
@@ -316,7 +318,7 @@ class SSAUpdaterTraits<SILSSAUpdater> {
                                  SILSSAUpdater *ssaUpdater) {
     // Add the argument to the block.
     SILValue phi(
-        block->createPhiArgument(ssaUpdater->type, OwnershipKind::Owned));
+        block->createPhiArgument(ssaUpdater->type, ssaUpdater->ownershipKind));
 
     // Mark all predecessor blocks with the sentinel undef value.
     SmallVector<SILBasicBlock *, 4> predBlockList(