Merge pull request #32134 from eeckstein/cow-support2

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

Author: eeckstein

include/swift/AST/KnownDecls.def

@@ -44,6 +44,8 @@ FUNC_DECL(AllocateUninitializedArray,
           "_allocateUninitializedArray")
 FUNC_DECL(DeallocateUninitializedArray,
           "_deallocateUninitializedArray")
+FUNC_DECL(FinalizeUninitializedArray,
+          "_finalizeUninitializedArray")
 
 FUNC_DECL(ForceBridgeFromObjectiveC,
           "_forceBridgeFromObjectiveC")

include/swift/AST/SemanticAttrs.def

@@ -60,6 +60,7 @@ SEMANTICS_ATTR(ARRAY_WITH_UNSAFE_MUTABLE_BUFFER_POINTER, "array.withUnsafeMutabl
 SEMANTICS_ATTR(ARRAY_COUNT, "array.count")
 SEMANTICS_ATTR(ARRAY_DEALLOC_UNINITIALIZED, "array.dealloc_uninitialized")
 SEMANTICS_ATTR(ARRAY_UNINITIALIZED_INTRINSIC, "array.uninitialized_intrinsic")
+SEMANTICS_ATTR(ARRAY_FINALIZE_INTRINSIC, "array.finalize_intrinsic")
 
 SEMANTICS_ATTR(SEQUENCE_FOR_EACH, "sequence.forEach")
 

include/swift/AST/SemanticAttrs.h

@@ -19,6 +19,7 @@
 #ifndef SWIFT_SEMANTICS_H
 #define SWIFT_SEMANTICS_H
 
+#include "swift/Basic/LLVM.h"
 #include "llvm/ADT/StringRef.h"
 
 namespace swift {

include/swift/SIL/SILNodes.def

@@ -880,8 +880,13 @@ NODE_RANGE(NonValueInstruction, UnreachableInst, CondFailInst)
 ABSTRACT_INST(MultipleValueInstruction, SILInstruction)
 FULLAPPLYSITE_MULTIPLE_VALUE_INST(BeginApplyInst, begin_apply,
                                   MultipleValueInstruction, MayHaveSideEffects, MayRelease)
+                                  
+// begin_cow_mutation is defined to have side effects, because it has
+// dependencies with instructions which retain the buffer operand. This prevents
+// optimizations from moving begin_cow_mutation instructions across such retain
+// instructions.
 MULTIPLE_VALUE_INST(BeginCOWMutationInst, begin_cow_mutation,
-                    MultipleValueInstruction, None, DoesNotRelease)
+                    MultipleValueInstruction, MayHaveSideEffects, DoesNotRelease)
 MULTIPLE_VALUE_INST(DestructureStructInst, destructure_struct,
                     MultipleValueInstruction, None, DoesNotRelease)
 MULTIPLE_VALUE_INST(DestructureTupleInst, destructure_tuple,

include/swift/SILOptimizer/Differentiation/Common.h

@@ -17,10 +17,12 @@
 #ifndef SWIFT_SILOPTIMIZER_UTILS_DIFFERENTIATION_COMMON_H
 #define SWIFT_SILOPTIMIZER_UTILS_DIFFERENTIATION_COMMON_H
 
+#include "swift/AST/SemanticAttrs.h"
 #include "swift/SIL/SILDifferentiabilityWitness.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILModule.h"
 #include "swift/SIL/TypeSubstCloner.h"
+#include "swift/SILOptimizer/Analysis/ArraySemantic.h"
 #include "swift/SILOptimizer/Analysis/DifferentiableActivityAnalysis.h"
 
 namespace swift {

lib/SIL/IR/SILModule.cpp

@@ -355,6 +355,7 @@ bool SILModule::linkFunction(SILFunction *F, SILModule::LinkingMode Mode) {
 
 SILFunction *SILModule::findFunction(StringRef Name, SILLinkage Linkage) {
   assert((Linkage == SILLinkage::Public ||
+          Linkage == SILLinkage::SharedExternal ||
           Linkage == SILLinkage::PublicExternal) &&
          "Only a lookup of public functions is supported currently");
 
@@ -405,6 +406,9 @@ SILFunction *SILModule::findFunction(StringRef Name, SILLinkage Linkage) {
   // compilation, simply convert it into an external declaration,
   // so that a compiled version from the shared library is used.
   if (F->isDefinition() &&
+      // Don't eliminate bodies of _alwaysEmitIntoClient functions
+      // (PublicNonABI linkage is de-serialized as SharedExternal)
+      F->getLinkage() != SILLinkage::SharedExternal &&
       !F->getModule().getOptions().shouldOptimize()) {
     F->convertToDeclaration();
   }

lib/SILGen/SILGenApply.cpp

@@ -4962,6 +4962,22 @@ void SILGenFunction::emitUninitializedArrayDeallocation(SILLocation loc,
                               SGFContext());
 }
 
+ManagedValue SILGenFunction::emitUninitializedArrayFinalization(SILLocation loc,
+                                                      SILValue array) {
+  auto &Ctx = getASTContext();
+  auto finalize = Ctx.getFinalizeUninitializedArray();
+
+  CanType arrayTy = array->getType().getASTType();
+
+  // Invoke the intrinsic.
+  auto subMap = arrayTy->getContextSubstitutionMap(SGM.M.getSwiftModule(),
+                                                   Ctx.getArrayDecl());
+  RValue result = emitApplyOfLibraryIntrinsic(loc, finalize, subMap,
+                              ManagedValue::forUnmanaged(array),
+                              SGFContext());
+  return std::move(result).getScalarValue();
+}
+
 namespace {
   /// A cleanup that deallocates an uninitialized array.
   class DeallocateUninitializedArray: public Cleanup {

lib/SILGen/SILGenExpr.cpp

@@ -2113,10 +2113,11 @@ ManagedValue Lowering::emitEndVarargs(SILGenFunction &SGF, SILLocation loc,
   SGF.Cleanups.setCleanupState(varargs.getAbortCleanup(), CleanupState::Dead);
 
   // Reactivate the result cleanup.
-  auto result = varargs.getArray();
-  if (result.hasCleanup())
-    SGF.Cleanups.setCleanupState(result.getCleanup(), CleanupState::Active);
-  return result;
+  auto array = varargs.getArray();
+  if (array.hasCleanup())
+    SGF.Cleanups.setCleanupState(array.getCleanup(), CleanupState::Active);
+
+  return SGF.emitUninitializedArrayFinalization(loc, array.forward(SGF));
 }
 
 RValue RValueEmitter::visitTupleExpr(TupleExpr *E, SGFContext C) {

lib/SILGen/SILGenFunction.h

@@ -1187,6 +1187,7 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
 
   CleanupHandle enterDeallocateUninitializedArrayCleanup(SILValue array);
   void emitUninitializedArrayDeallocation(SILLocation loc, SILValue array);
+  ManagedValue emitUninitializedArrayFinalization(SILLocation loc, SILValue array);
 
   /// Emit a cleanup for an owned value that should be written back at end of
   /// scope if the value is not forwarded.

lib/SILOptimizer/Analysis/MemoryBehavior.cpp

@@ -180,6 +180,7 @@ class MemoryBehaviorVisitor
   MemBehavior visitStrongReleaseInst(StrongReleaseInst *BI);
   MemBehavior visitReleaseValueInst(ReleaseValueInst *BI);
   MemBehavior visitSetDeallocatingInst(SetDeallocatingInst *BI);
+  MemBehavior visitBeginCOWMutationInst(BeginCOWMutationInst *BCMI);
 #define ALWAYS_OR_SOMETIMES_LOADABLE_CHECKED_REF_STORAGE(Name, ...) \
   MemBehavior visit##Name##ReleaseInst(Name##ReleaseInst *BI);
 #include "swift/AST/ReferenceStorage.def"
@@ -395,6 +396,14 @@ MemBehavior MemoryBehaviorVisitor::visitSetDeallocatingInst(SetDeallocatingInst
   return MemBehavior::None;
 }
 
+MemBehavior MemoryBehaviorVisitor::
+visitBeginCOWMutationInst(BeginCOWMutationInst *BCMI) {
+  // begin_cow_mutation is defined to have side effects, because it has
+  // dependencies with instructions which retain the buffer operand.
+  // But it never interferes with any memory address.
+  return MemBehavior::None;
+}
+
 //===----------------------------------------------------------------------===//
 //                            Top Level Entrypoint
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/Differentiation/LinearMapInfo.cpp

@@ -414,6 +414,8 @@ void LinearMapInfo::generateDifferentiationDataStructures(
         // initialization is linear and handled separately.
         if (!shouldDifferentiateApplySite(ai) || isArrayLiteralIntrinsic(ai))
           continue;
+        if (ArraySemanticsCall(ai, semantics::ARRAY_FINALIZE_INTRINSIC))
+          continue;
         LLVM_DEBUG(getADDebugStream()
                    << "Adding linear map struct field for " << *ai);
         addLinearMapToStruct(context, ai);

lib/SILOptimizer/Differentiation/PullbackEmitter.cpp

@@ -1424,6 +1424,19 @@ void PullbackEmitter::visitApplyInst(ApplyInst *ai) {
   // special `store` and `copy_addr` support.
   if (isArrayLiteralIntrinsic(ai))
     return;
+  auto loc = ai->getLoc();
+  auto *bb = ai->getParent();
+  // Handle `array.finalize_intrinsic` applications. `array.finalize_intrinsic`
+  // semantically behaves like an identity function.
+  if (ArraySemanticsCall(ai, semantics::ARRAY_FINALIZE_INTRINSIC)) {
+    assert(ai->getNumArguments() == 1 &&
+           "Expected intrinsic to have one operand");
+    // Accumulate result's adjoint into argument's adjoint.
+    auto adjResult = getAdjointValue(bb, ai);
+    auto origArg = ai->getArgumentsWithoutIndirectResults().front();
+    addAdjointValue(bb, origArg, adjResult, loc);
+    return;
+  }
   // Replace a call to a function with a call to its pullback.
   auto &nestedApplyInfo = getContext().getNestedApplyInfo();
   auto applyInfoLookup = nestedApplyInfo.find(ai);
@@ -1439,7 +1452,6 @@ void PullbackEmitter::visitApplyInst(ApplyInst *ai) {
   // Get the pullback.
   auto *field = getPullbackInfo().lookUpLinearMapDecl(ai);
   assert(field);
-  auto loc = ai->getLoc();
   auto pullback = getPullbackStructElement(ai->getParent(), field);
 
   // Get the original result of the `apply` instruction.
@@ -1478,7 +1490,6 @@ void PullbackEmitter::visitApplyInst(ApplyInst *ai) {
   }
 
   // Get formal callee pullback arguments.
-  auto *bb = ai->getParent();
   assert(applyInfo.indices.results->getNumIndices() == 1);
   for (auto resultIndex : applyInfo.indices.results->getIndices()) {
     assert(resultIndex < origAllResults.size());

lib/SILOptimizer/Differentiation/VJPEmitter.cpp

@@ -541,6 +541,15 @@ void VJPEmitter::visitApplyInst(ApplyInst *ai) {
     TypeSubstCloner::visitApplyInst(ai);
     return;
   }
+  // If callee is `array.finalize_intrinsic`, do standard cloning.
+  // `array.finalize_intrinsic` has special-case pullback generation.
+  if (ArraySemanticsCall(ai, semantics::ARRAY_FINALIZE_INTRINSIC)) {
+    LLVM_DEBUG(getADDebugStream()
+               << "Cloning `array.finalize_intrinsic` `apply`:\n"
+               << *ai << '\n');
+    TypeSubstCloner::visitApplyInst(ai);
+    return;
+  }
   // If the original function is a semantic member accessor, do standard
   // cloning. Semantic member accessors have special pullback generation logic,
   // so all `apply` instructions can be directly cloned to the VJP.

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) {
@@ -642,7 +654,7 @@ bool COWArrayOpt::hasLoopOnlyDestructorSafeArrayOperations() {
 
       // Semantic calls are safe.
       ArraySemanticsCall Sem(Inst);
-      if (Sem) {
+      if (Sem && Sem.hasSelf()) {
         auto Kind = Sem.getKind();
         // Safe because they create new arrays.
         if (Kind == ArrayCallKind::kArrayInit ||
@@ -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/LoopTransforms/ForEachLoopUnroll.cpp

@@ -323,8 +323,15 @@ void ArrayInfo::classifyUsesOfArray(SILValue arrayValue) {
     // as the array itself is not modified (which is possible with reference
     // types).
     ArraySemanticsCall arrayOp(user);
-    if (!arrayOp.doesNotChangeArray())
-      mayBeWritten = true;
+    if (arrayOp.doesNotChangeArray())
+      continue;
+    
+    if (arrayOp.getKind() == swift::ArrayCallKind::kArrayFinalizeIntrinsic) {
+      classifyUsesOfArray((ApplyInst *)arrayOp);
+      continue;
+    }
+    
+    mayBeWritten = true;
   }
 }