stdlib, SIL optimizer: use the SIL copy-on-write representation in the Array types.

Use the new builtins for COW representation in Array, ContiguousArray and ArraySlice.
The basic idea is to strictly separate code which mutates an array buffer from code which reads from an array.
The concept is explained in more detail in docs/SIL.rst, section "Copy-on-Write Representation".

The main change is to use beginCOWMutation() instead of isUniquelyReferenced() and insert endCOWMutation() at the end of all mutating functions. Also, reading from the array buffer must be done differently, depending on if the buffer is in a mutable or immutable state.

All the required invariants are enforced by runtime checks - but only in an assert-build of the library: a bit in the buffer flags indicates if the buffer is mutable or not.

Along with the library changes, also two optimizations needed to be updated: COWArrayOpt and ObjectOutliner.

Author: eeckstein

lib/SILOptimizer/LoopTransforms/ArrayOpt.h

@@ -118,15 +118,17 @@ class StructUseCollector {
     }
   }
 
-  /// Returns true if there is a single address user of the value.
-  bool hasSingleAddressUse(SILInstruction *SingleAddressUser) {
+  /// Returns true if there are only address users of the value.
+  bool hasOnlyAddressUses(ApplyInst *use1, ApplyInst *use2) {
     if (!AggregateAddressUsers.empty())
       return false;
     if (!ElementAddressUsers.empty())
       return false;
-    if (StructAddressUsers.size() != 1)
-      return false;
-    return StructAddressUsers[0] == SingleAddressUser;
+    for (SILInstruction *user : StructAddressUsers) {
+      if (user != use1 && user != use2)
+        return false;
+    }
+    return true;
   }
 
 protected:

lib/SILOptimizer/LoopTransforms/COWArrayOpt.cpp

@@ -436,6 +436,18 @@ bool COWArrayOpt::checkSafeArrayAddressUses(UserList &AddressUsers) {
   return true;
 }
 
+template <typename UserRange>
+ArraySemanticsCall getEndMutationCall(const UserRange &AddressUsers) {
+  for (auto *UseInst : AddressUsers) {
+    if (auto *AI = dyn_cast<ApplyInst>(UseInst)) {
+      ArraySemanticsCall ASC(AI);
+      if (ASC.getKind() == ArrayCallKind::kEndMutation)
+        return ASC;
+    }
+  }
+  return ArraySemanticsCall();
+}
+
 /// Returns true if this instruction is a safe array use if all of its users are
 /// also safe array users.
 static SILValue isTransitiveSafeUser(SILInstruction *I) {
@@ -811,8 +823,14 @@ void COWArrayOpt::hoistAddressProjections(Operand &ArrayOp) {
   }
 }
 
-/// Check if this call to "make_mutable" is hoistable, and move it, or delete it
-/// if it's already hoisted.
+/// Check if this call to "make_mutable" is hoistable, and copy it, along with
+/// the corresponding end_mutation call, to the loop pre-header.
+///
+/// The origial make_mutable/end_mutation calls remain in the loop, because
+/// removing them would violate the COW representation rules.
+/// Having those calls in the pre-header will then enable COWOpts (after
+/// inlining) to constant fold the uniqueness check of the begin_cow_mutation
+/// in the loop.
 bool COWArrayOpt::hoistMakeMutable(ArraySemanticsCall MakeMutable,
                                    bool dominatesExits) {
   LLVM_DEBUG(llvm::dbgs() << "    Checking mutable array: " <<CurrentArrayAddr);
@@ -872,6 +890,18 @@ bool COWArrayOpt::hoistMakeMutable(ArraySemanticsCall MakeMutable,
       return false;
   }
 
+  auto ArrayUsers = llvm::map_range(MakeMutable.getSelf()->getUses(),
+                                    ValueBase::UseToUser());
+
+  // There should be a call to end_mutation. Find it so that we can copy it to
+  // the pre-header.
+  ArraySemanticsCall EndMutation = getEndMutationCall(ArrayUsers);
+  if (!EndMutation) {
+    EndMutation = getEndMutationCall(StructUses.StructAddressUsers);
+    if (!EndMutation)
+      return false;
+  }
+
   // Hoist the make_mutable.
   LLVM_DEBUG(llvm::dbgs() << "    Hoisting make_mutable: " << *MakeMutable);
 
@@ -880,12 +910,18 @@ bool COWArrayOpt::hoistMakeMutable(ArraySemanticsCall MakeMutable,
   assert(MakeMutable.canHoist(Preheader->getTerminator(), DomTree) &&
          "Should be able to hoist make_mutable");
 
-  MakeMutable.hoist(Preheader->getTerminator(), DomTree);
+  // Copy the make_mutable and end_mutation calls to the pre-header.
+  TermInst *insertionPoint = Preheader->getTerminator();
+  ApplyInst *hoistedMM = MakeMutable.copyTo(insertionPoint, DomTree);
+  ApplyInst *EMInst = EndMutation;
+  ApplyInst *hoistedEM = cast<ApplyInst>(EMInst->clone(insertionPoint));
+  hoistedEM->setArgument(0, hoistedMM->getArgument(0));
+  placeFuncRef(hoistedEM, DomTree);
 
   // Register array loads. This is needed for hoisting make_mutable calls of
   // inner arrays in the two-dimensional case.
   if (arrayContainerIsUnique &&
-      StructUses.hasSingleAddressUse((ApplyInst *)MakeMutable)) {
+      StructUses.hasOnlyAddressUses((ApplyInst *)MakeMutable, EMInst)) {
     for (auto use : MakeMutable.getSelf()->getUses()) {
       if (auto *LI = dyn_cast<LoadInst>(use->getUser()))
         HoistableLoads.insert(LI);
@@ -917,39 +953,33 @@ bool COWArrayOpt::run() {
   // is only mapped to a call once the analysis has determined that no
   // make_mutable calls are required within the loop body for that array.
   llvm::SmallDenseMap<SILValue, ApplyInst*> ArrayMakeMutableMap;
-
+  
+  llvm::SmallVector<ArraySemanticsCall, 8> makeMutableCalls;
+  
   for (auto *BB : Loop->getBlocks()) {
     if (ColdBlocks.isCold(BB))
       continue;
-    bool dominatesExits = dominatesExitingBlocks(BB);
-    for (auto II = BB->begin(), IE = BB->end(); II != IE;) {
-      // Inst may be moved by hoistMakeMutable.
-      SILInstruction *Inst = &*II;
-      ++II;
-      ArraySemanticsCall MakeMutableCall(Inst, "array.make_mutable");
-      if (!MakeMutableCall)
-        continue;
+      
+    // Instructions are getting moved around. To not mess with iterator
+    // invalidation, first collect all calls, and then do the transformation.
+    for (SILInstruction &I : *BB) {
+      ArraySemanticsCall MakeMutableCall(&I, "array.make_mutable");
+      if (MakeMutableCall)
+        makeMutableCalls.push_back(MakeMutableCall);
+    }
 
+    bool dominatesExits = dominatesExitingBlocks(BB);
+    for (ArraySemanticsCall MakeMutableCall : makeMutableCalls) {
       CurrentArrayAddr = MakeMutableCall.getSelf();
       auto HoistedCallEntry = ArrayMakeMutableMap.find(CurrentArrayAddr);
       if (HoistedCallEntry == ArrayMakeMutableMap.end()) {
-        if (!hoistMakeMutable(MakeMutableCall, dominatesExits)) {
+        if (hoistMakeMutable(MakeMutableCall, dominatesExits)) {
+          ArrayMakeMutableMap[CurrentArrayAddr] = MakeMutableCall;
+          HasChanged = true;
+        } else {
           ArrayMakeMutableMap[CurrentArrayAddr] = nullptr;
-          continue;
         }
-
-        ArrayMakeMutableMap[CurrentArrayAddr] = MakeMutableCall;
-        HasChanged = true;
-        continue;
       }
-
-      if (!HoistedCallEntry->second)
-        continue;
-
-      LLVM_DEBUG(llvm::dbgs() << "    Removing make_mutable call: "
-                              << *MakeMutableCall);
-      MakeMutableCall.removeCall();
-      HasChanged = true;
     }
   }
   return HasChanged;

lib/SILOptimizer/Transforms/ObjectOutliner.cpp

@@ -35,15 +35,14 @@ class ObjectOutliner {
     return type.getNominalOrBoundGenericNominal() == ArrayDecl;
   }
 
-  bool isValidUseOfObject(SILInstruction *Val,
-                          bool isCOWObject,
-                          ApplyInst **FindStringCall = nullptr);
+  bool isValidUseOfObject(SILInstruction *Val);
+
+  ApplyInst *findFindStringCall(SILValue V);
 
   bool getObjectInitVals(SILValue Val,
                          llvm::DenseMap<VarDecl *, StoreInst *> &MemberStores,
                          llvm::SmallVectorImpl<StoreInst *> &TailStores,
-                         unsigned NumTailTupleElements,
-                         ApplyInst **FindStringCall);
+                         unsigned NumTailTupleElements);
   bool handleTailAddr(int TailIdx, SILInstruction *I, unsigned NumTailTupleElements,
                       llvm::SmallVectorImpl<StoreInst *> &TailStores);
 
@@ -116,13 +115,7 @@ static bool isValidInitVal(SILValue V) {
 }
 
 /// Check if a use of an object may prevent outlining the object.
-///
-/// If \p isCOWObject is true, then the object reference is wrapped into a
-/// COW container. Currently this is just Array<T>.
-/// If a use is a call to the findStringSwitchCase semantic call, the apply
-/// is returned in \p FindStringCall.
-bool ObjectOutliner::isValidUseOfObject(SILInstruction *I, bool isCOWObject,
-                                      ApplyInst **FindStringCall) {
+bool ObjectOutliner::isValidUseOfObject(SILInstruction *I) {
   switch (I->getKind()) {
   case SILInstructionKind::DebugValueAddrInst:
   case SILInstructionKind::DebugValueInst:
@@ -132,51 +125,25 @@ bool ObjectOutliner::isValidUseOfObject(SILInstruction *I, bool isCOWObject,
   case SILInstructionKind::StrongReleaseInst:
   case SILInstructionKind::FixLifetimeInst:
   case SILInstructionKind::SetDeallocatingInst:
+  case SILInstructionKind::EndCOWMutationInst:
     return true;
 
-  case SILInstructionKind::ReturnInst:
-  case SILInstructionKind::TryApplyInst:
-  case SILInstructionKind::PartialApplyInst:
-  case SILInstructionKind::StoreInst:
-    /// We don't have a representation for COW objects in SIL, so we do some
-    /// ad-hoc testing: We can ignore uses of a COW object if any use after
-    /// this will do a uniqueness checking before the object is modified.
-    return isCOWObject;
-
-  case SILInstructionKind::ApplyInst:
-    if (!isCOWObject)
-      return false;
-    // There should only be a single call to findStringSwitchCase. But even
-    // if there are multiple calls, it's not problem - we'll just optimize the
-    // last one we find.
-    if (cast<ApplyInst>(I)->hasSemantics(semantics::FIND_STRING_SWITCH_CASE))
-      *FindStringCall = cast<ApplyInst>(I);
-    return true;
-
-  case SILInstructionKind::StructInst:
-    if (isCOWType(cast<StructInst>(I)->getType())) {
-      // The object is wrapped into a COW container.
-      isCOWObject = true;
-    }
-    break;
-
-  case SILInstructionKind::UncheckedRefCastInst:
   case SILInstructionKind::StructElementAddrInst:
   case SILInstructionKind::AddressToPointerInst:
-    assert(!isCOWObject && "instruction cannot have a COW object as operand");
-    break;
-
+  case SILInstructionKind::StructInst:
   case SILInstructionKind::TupleInst:
   case SILInstructionKind::TupleExtractInst:
   case SILInstructionKind::EnumInst:
-    break;
-
   case SILInstructionKind::StructExtractInst:
-    // To be on the safe side we don't consider the object as COW if it is
-    // extracted again from the COW container: the uniqueness check may be
-    // optimized away in this case.
-    isCOWObject = false;
-    break;
+  case SILInstructionKind::UncheckedRefCastInst:
+  case SILInstructionKind::UpcastInst: {
+    auto SVI = cast<SingleValueInstruction>(I);
+    for (Operand *Use : getNonDebugUses(SVI)) {
+      if (!isValidUseOfObject(Use->getUser()))
+        return false;
+    }
+    return true;
+  }
 
   case SILInstructionKind::BuiltinInst: {
     // Handle the case for comparing addresses. This occurs when the Array
@@ -198,13 +165,37 @@ bool ObjectOutliner::isValidUseOfObject(SILInstruction *I, bool isCOWObject,
   default:
     return false;
   }
+}
 
-  auto SVI = cast<SingleValueInstruction>(I);
-  for (Operand *Use : getNonDebugUses(SVI)) {
-    if (!isValidUseOfObject(Use->getUser(), isCOWObject, FindStringCall))
-      return false;
+/// Finds a call to findStringSwitchCase in the uses of \p V.
+ApplyInst *ObjectOutliner::findFindStringCall(SILValue V) {
+  for (Operand *use : V->getUses()) {
+    SILInstruction *user = use->getUser();
+    switch (user->getKind()) {
+    case SILInstructionKind::ApplyInst:
+      // There should only be a single call to findStringSwitchCase. But even
+      // if there are multiple calls, it's not problem - we'll just optimize the
+      // last one we find.
+      if (cast<ApplyInst>(user)->hasSemantics(semantics::FIND_STRING_SWITCH_CASE))
+        return cast<ApplyInst>(user);
+      break;
+
+    case SILInstructionKind::StructInst:
+    case SILInstructionKind::TupleInst:
+    case SILInstructionKind::UncheckedRefCastInst:
+    case SILInstructionKind::UpcastInst: {
+      if (ApplyInst *foundCall =
+           findFindStringCall(cast<SingleValueInstruction>(user))) {
+        return foundCall;
+      }
+      break;
+    }
+
+    default:
+      break;
+    }
   }
-  return true;
+  return nullptr;
 }
 
 /// Handle the address of a tail element.
@@ -232,20 +223,19 @@ bool ObjectOutliner::handleTailAddr(int TailIdx, SILInstruction *TailAddr,
       }
     }
   }
-  return isValidUseOfObject(TailAddr, /*isCOWObject*/false);
+  return isValidUseOfObject(TailAddr);
 }
 
 /// Get the init values for an object's stored properties and its tail elements.
 bool ObjectOutliner::getObjectInitVals(SILValue Val,
                         llvm::DenseMap<VarDecl *, StoreInst *> &MemberStores,
                         llvm::SmallVectorImpl<StoreInst *> &TailStores,
-                        unsigned NumTailTupleElements,
-                        ApplyInst **FindStringCall) {
+                        unsigned NumTailTupleElements) {
   for (Operand *Use : Val->getUses()) {
     SILInstruction *User = Use->getUser();
     if (auto *UC = dyn_cast<UpcastInst>(User)) {
       // Upcast is transparent.
-      if (!getObjectInitVals(UC, MemberStores, TailStores, NumTailTupleElements, FindStringCall))
+      if (!getObjectInitVals(UC, MemberStores, TailStores, NumTailTupleElements))
         return false;
     } else if (auto *REA = dyn_cast<RefElementAddrInst>(User)) {
       // The address of a stored property.
@@ -255,7 +245,7 @@ bool ObjectOutliner::getObjectInitVals(SILValue Val,
           if (!isValidInitVal(SI->getSrc()) || MemberStores[REA->getField()])
             return false;
           MemberStores[REA->getField()] = SI;
-        } else if (!isValidUseOfObject(ElemAddrUser, /*isCOWObject*/false)) {
+        } else if (!isValidUseOfObject(ElemAddrUser)) {
           return false;
         }
       }
@@ -280,7 +270,7 @@ bool ObjectOutliner::getObjectInitVals(SILValue Val,
           return false;
         }
       }
-    } else if (!isValidUseOfObject(User, /*isCOWObject*/false, FindStringCall)) {
+    } else if (!isValidUseOfObject(User)) {
       return false;
     }
   }
@@ -305,7 +295,8 @@ class GlobalVariableMangler : public Mangle::ASTMangler {
 /// Try to convert an object allocation into a statically initialized object.
 ///
 /// In general this works for any class, but in practice it will only kick in
-/// for array buffer objects. The use cases are array literals in a function.
+/// for copy-on-write buffers, like array buffer objects.
+/// The use cases are array literals in a function.
 /// For example:
 ///     func getarray() -> [Int] {
 ///       return [1, 2, 3]
@@ -314,6 +305,19 @@ bool ObjectOutliner::optimizeObjectAllocation(AllocRefInst *ARI) {
   if (ARI->isObjC())
     return false;
 
+  // Find the end_cow_mutation. Only for such COW buffer objects we do the
+  // transformation.
+  EndCOWMutationInst *endCOW = nullptr;
+  for (Operand *use : ARI->getUses()) {
+    if (auto *ecm = dyn_cast<EndCOWMutationInst>(use->getUser())) {
+      if (endCOW)
+        return false;
+      endCOW = ecm;
+    }
+  }
+  if (!endCOW || endCOW->doKeepUnique())
+    return false;
+
   // Check how many tail allocated elements are on the object.
   ArrayRef<Operand> TailCounts = ARI->getTailAllocatedCounts();
   SILType TailType;
@@ -363,11 +367,10 @@ bool ObjectOutliner::optimizeObjectAllocation(AllocRefInst *ARI) {
   }
 
   TailStores.resize(NumStores);
-  ApplyInst *FindStringCall = nullptr;
 
   // Get the initialization stores of the object's properties and tail
   // allocated elements. Also check if there are any "bad" uses of the object.
-  if (!getObjectInitVals(ARI, MemberStores, TailStores, NumTailTupleElems, &FindStringCall))
+  if (!getObjectInitVals(ARI, MemberStores, TailStores, NumTailTupleElems))
     return false;
 
   // Is there a store for all the class properties?
@@ -452,6 +455,12 @@ bool ObjectOutliner::optimizeObjectAllocation(AllocRefInst *ARI) {
   SILBuilder B(ARI);
   GlobalValueInst *GVI = B.createGlobalValue(ARI->getLoc(), Glob);
   B.createStrongRetain(ARI->getLoc(), GVI, B.getDefaultAtomicity());
+
+  ApplyInst *FindStringCall = findFindStringCall(endCOW);
+  assert(endCOW->getOperand() == ARI);
+  endCOW->replaceAllUsesWith(ARI);
+  ToRemove.push_back(endCOW);
+
   llvm::SmallVector<Operand *, 8> Worklist(ARI->use_begin(), ARI->use_end());
   while (!Worklist.empty()) {
     auto *Use = Worklist.pop_back_val();