Merge pull request #32397 from gottesmm/pr-92e92ba00c4c145dbd650430947c43d7b3a743d5

[ownership] Loosen restrictions around what we specialize and add generic specialization tests behind a flag.

Author: gottesmm

include/swift/SIL/ApplySite.h

@@ -21,9 +21,11 @@
 #ifndef SWIFT_SIL_APPLYSITE_H
 #define SWIFT_SIL_APPLYSITE_H
 
+#include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBasicBlock.h"
-#include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILFunction.h"
+#include "swift/SIL/SILInstruction.h"
+#include "llvm/ADT/ArrayRef.h"
 
 namespace swift {
 
@@ -502,6 +504,34 @@ class FullApplySite : public ApplySite {
     return getSubstCalleeConv().hasIndirectSILResults();
   }
 
+  /// If our apply site has a single direct result SILValue, return that
+  /// SILValue. Return SILValue() otherwise.
+  ///
+  /// This means that:
+  ///
+  /// 1. If we have an ApplyInst, we just visit the apply.
+  /// 2. If we have a TryApplyInst, we visit the first argument of the normal
+  ///    block.
+  /// 3. If we have a BeginApplyInst, we return SILValue() since the begin_apply
+  ///    yields values instead of returning them. A returned value should only
+  ///    be valid after a full apply site has completely finished executing.
+  SILValue getSingleDirectResult() const {
+    switch (getKind()) {
+    case FullApplySiteKind::ApplyInst:
+      return SILValue(cast<ApplyInst>(getInstruction()));
+    case FullApplySiteKind::BeginApplyInst: {
+      return SILValue();
+    }
+    case FullApplySiteKind::TryApplyInst: {
+      auto *normalBlock = cast<TryApplyInst>(getInstruction())->getNormalBB();
+      assert(normalBlock->getNumArguments() == 1 &&
+             "Expected try apply to have a single result");
+      return normalBlock->getArgument(0);
+    }
+    }
+    llvm_unreachable("Covered switch isn't covered?!");
+  }
+
   unsigned getNumIndirectSILResults() const {
     return getSubstCalleeConv().getNumIndirectSILResults();
   }

include/swift/SIL/SILBuilder.h

@@ -726,7 +726,8 @@ class SILBuilder {
   }
 
   LoadBorrowInst *createLoadBorrow(SILLocation Loc, SILValue LV) {
-    assert(isLoadableOrOpaque(LV->getType()));
+    assert(isLoadableOrOpaque(LV->getType()) &&
+           !LV->getType().isTrivial(getFunction()));
     return insert(new (getModule())
                       LoadBorrowInst(getSILDebugLocation(Loc), LV));
   }
@@ -737,11 +738,19 @@ class SILBuilder {
                       BeginBorrowInst(getSILDebugLocation(Loc), LV));
   }
 
+  /// Convenience function for creating a load_borrow on non-trivial values and
+  /// load [trivial] on trivial values. Becomes load unqualified in non-ossa
+  /// functions.
   SILValue emitLoadBorrowOperation(SILLocation loc, SILValue v) {
     if (!hasOwnership()) {
       return emitLoadValueOperation(loc, v,
                                     LoadOwnershipQualifier::Unqualified);
     }
+
+    if (v->getType().isTrivial(getFunction())) {
+      return emitLoadValueOperation(loc, v, LoadOwnershipQualifier::Trivial);
+    }
+
     return createLoadBorrow(loc, v);
   }
 
@@ -877,6 +886,33 @@ class SILBuilder {
                       StoreBorrowInst(getSILDebugLocation(Loc), Src, DestAddr));
   }
 
+  /// A helper function for emitting store_borrow in operations where one must
+  /// handle both ossa and non-ossa code.
+  ///
+  /// In words:
+  ///
+  /// * If the function does not have ownership, this just emits an unqualified
+  ///   store.
+  ///
+  /// * If the function has ownership, but the type is trivial, use store
+  ///   [trivial].
+  ///
+  /// * Otherwise, emit an actual store_borrow.
+  void emitStoreBorrowOperation(SILLocation loc, SILValue src,
+                                SILValue destAddr) {
+    if (!hasOwnership()) {
+      return emitStoreValueOperation(loc, src, destAddr,
+                                     StoreOwnershipQualifier::Unqualified);
+    }
+
+    if (src->getType().isTrivial(getFunction())) {
+      return emitStoreValueOperation(loc, src, destAddr,
+                                     StoreOwnershipQualifier::Trivial);
+    }
+
+    createStoreBorrow(loc, src, destAddr);
+  }
+
   MarkUninitializedInst *
   createMarkUninitialized(SILLocation Loc, SILValue src,
                           MarkUninitializedInst::Kind k) {

lib/SILOptimizer/Transforms/GenericSpecializer.cpp

@@ -28,6 +28,10 @@
 
 using namespace swift;
 
+// For testing during bring up.
+static llvm::cl::opt<bool> EnableGenericSpecializerWithOwnership(
+    "sil-generic-specializer-enable-ownership", llvm::cl::init(false));
+
 namespace {
 
 class GenericSpecializer : public SILFunctionTransform {
@@ -39,7 +43,7 @@ class GenericSpecializer : public SILFunctionTransform {
     SILFunction &F = *getFunction();
 
     // TODO: We should be able to handle ownership.
-    if (F.hasOwnership())
+    if (F.hasOwnership() && !EnableGenericSpecializerWithOwnership)
       return;
 
     LLVM_DEBUG(llvm::dbgs() << "***** GenericSpecializer on function:"

lib/SILOptimizer/Utils/GenericCloner.cpp

@@ -79,8 +79,9 @@ void GenericCloner::populateCloned() {
 
     auto createAllocStack = [&]() {
       // We need an alloc_stack as a replacement for the indirect parameter.
-      assert(mappedType.isAddress());
-      mappedType = mappedType.getObjectType();
+      if (mappedType.isAddress()) {
+        mappedType = mappedType.getObjectType();
+      }
       auto AllocStackLoc = RegularLocation::getAutoGeneratedLocation();
       ASI = getBuilder().createAllocStack(AllocStackLoc, mappedType);
       AllocStacks.push_back(ASI);
@@ -106,24 +107,36 @@ void GenericCloner::populateCloned() {
         // Handle arguments for formal parameters.
         unsigned paramIdx = ArgIdx - origConv.getSILArgIndexOfFirstParam();
         if (ReInfo.isParamConverted(paramIdx)) {
-          // Store the new direct parameter to the alloc_stack.
-          createAllocStack();
+          assert(mappedType.isAddress());
+          mappedType = mappedType.getObjectType();
           auto *NewArg = ClonedEntryBB->createFunctionArgument(
               mappedType, OrigArg->getDecl());
-          getBuilder().createStore(Loc, NewArg, ASI,
-                                   StoreOwnershipQualifier::Unqualified);
 
-          // Try to create a new debug_value from an existing debug_value_addr.
+          // Try to create a new debug_value from an existing debug_value_addr
+          // for the argument. We do this before storing to ensure that when we
+          // are cloning code in ossa the argument has not been consumed by the
+          // store below.
           for (Operand *ArgUse : OrigArg->getUses()) {
             if (auto *DVAI = dyn_cast<DebugValueAddrInst>(ArgUse->getUser())) {
+              auto *oldScope = getBuilder().getCurrentDebugScope();
               getBuilder().setCurrentDebugScope(
                   remapScope(DVAI->getDebugScope()));
               getBuilder().createDebugValue(DVAI->getLoc(), NewArg,
                                             *DVAI->getVarInfo());
-              getBuilder().setCurrentDebugScope(nullptr);
+              getBuilder().setCurrentDebugScope(oldScope);
               break;
             }
           }
+
+          // Store the new direct parameter to an alloc_stack.
+          createAllocStack();
+          if (!NewArg->getArgumentConvention().isGuaranteedConvention()) {
+            getBuilder().emitStoreValueOperation(Loc, NewArg, ASI,
+                                                 StoreOwnershipQualifier::Init);
+          } else {
+            getBuilder().emitStoreBorrowOperation(Loc, NewArg, ASI);
+          }
+
           entryArgs.push_back(ASI);
           return true;
         }
@@ -150,9 +163,9 @@ void GenericCloner::visitTerminator(SILBasicBlock *BB) {
     if (ReturnValueAddr) {
       // The result is converted from indirect to direct. We have to load the
       // returned value from the alloc_stack.
-      ReturnValue =
-        getBuilder().createLoad(ReturnValueAddr->getLoc(), ReturnValueAddr,
-                                LoadOwnershipQualifier::Unqualified);
+      ReturnValue = getBuilder().emitLoadValueOperation(
+          ReturnValueAddr->getLoc(), ReturnValueAddr,
+          LoadOwnershipQualifier::Take);
     }
     for (AllocStackInst *ASI : reverse(AllocStacks)) {
       getBuilder().createDeallocStack(ASI->getLoc(), ASI);