Merge pull request #80652 from atrick/62-address-specialize

[6.2] Fix GenericSpecializer for addressable parameters.

Author: atrick

include/swift/AST/LifetimeDependence.h

@@ -316,11 +316,6 @@ class LifetimeDependenceInfo {
       && scopeLifetimeParamIndices->contains(index);
   }
 
-  bool checkAddressable(int index) const {
-    return hasAddressableParamIndices()
-      && getAddressableIndices()->contains(index);
-  }
-
   std::string getString() const;
   void Profile(llvm::FoldingSetNodeID &ID) const;
   void getConcatenatedData(SmallVectorImpl<bool> &concatenatedData) const;

include/swift/AST/Types.h

@@ -5649,6 +5649,13 @@ class SILFunctionType final
     return getLifetimeDependenceFor(getNumParameters());
   }
 
+  /// Return true of the specified parameter is addressable based on its type
+  /// lowering in 'caller's context. This includes @_addressableForDependencies
+  /// parameter types.
+  ///
+  /// Defined in SILType.cpp.
+  bool isAddressable(unsigned paramIdx, SILFunction *caller);
+
   /// Returns true if the function type stores a Clang type that cannot
   /// be derived from its Swift type. Returns false otherwise, including if
   /// the function type is not @convention(c) or @convention(block).

include/swift/SIL/ApplySite.h

@@ -641,6 +641,10 @@ class ApplySite {
     return getArgumentParameterInfo(oper).hasOption(SILParameterInfo::Sending);
   }
 
+  /// Return true if 'operand' is addressable after type substitution in the
+  /// caller's context.
+  bool isAddressable(const Operand &operand) const;
+
   static ApplySite getFromOpaqueValue(void *p) { return ApplySite(p); }
 
   friend bool operator==(ApplySite lhs, ApplySite rhs) {

include/swift/SILOptimizer/Utils/Generics.h

@@ -84,7 +84,8 @@ class ReabstractionInfo {
   /// specializer.
   bool ConvertIndirectToDirect = true;
 
-  /// If true, drop unused arguments.
+  /// If true, drop unused arguments. Dropping unused arguments is a
+  /// prerequisite before promoting an indirect argument to a direct argument.
   /// See `droppedArguments`.
   bool dropUnusedArguments = false;
 
@@ -204,7 +205,7 @@ class ReabstractionInfo {
   ReabstractionInfo(ModuleDecl *targetModule, bool isModuleWholeModule,
                     ApplySite Apply, SILFunction *Callee,
                     SubstitutionMap ParamSubs, SerializedKind_t Serialized,
-                    bool ConvertIndirectToDirect, bool dropMetatypeArgs,
+                    bool ConvertIndirectToDirect, bool dropUnusedArguments,
                     OptRemark::Emitter *ORE = nullptr);
 
   /// Constructs the ReabstractionInfo for generic function \p Callee with

lib/SIL/IR/ApplySite.cpp

@@ -45,3 +45,12 @@ void ApplySite::insertAfterApplication(
   llvm_unreachable("covered switch isn't covered");
 }
 
+bool ApplySite::isAddressable(const Operand &operand) const {
+  unsigned calleeArgIndex = getCalleeArgIndex(operand);
+  assert(calleeArgIndex >= getSubstCalleeConv().getSILArgIndexOfFirstParam());
+  unsigned paramIdx =
+    calleeArgIndex - getSubstCalleeConv().getSILArgIndexOfFirstParam();
+
+  CanSILFunctionType calleeType = getSubstCalleeType();
+  return calleeType->isAddressable(paramIdx, getFunction());
+}

lib/SIL/IR/SILType.cpp

@@ -706,6 +706,29 @@ bool SILFunctionType::isNoReturnFunction(SILModule &M,
   return false;
 }
 
+bool SILFunctionType::isAddressable(unsigned paramIdx, SILFunction *caller) {
+  SILParameterInfo paramInfo = getParameters()[paramIdx];
+  for (auto &depInfo : getLifetimeDependencies()) {
+    auto *addressableIndices = depInfo.getAddressableIndices();
+    if (addressableIndices && addressableIndices->contains(paramIdx)) {
+      return true;
+    }
+    auto *condAddressableIndices = depInfo.getConditionallyAddressableIndices();
+    if (condAddressableIndices && condAddressableIndices->contains(paramIdx)) {
+      CanType argType = paramInfo.getArgumentType(
+        caller->getModule(), this, caller->getTypeExpansionContext());
+      CanType contextType =
+        argType->hasTypeParameter()
+        ? caller->mapTypeIntoContext(argType)->getCanonicalType()
+        : argType;
+      auto &tl = caller->getTypeLowering(contextType);
+      if (tl.getRecursiveProperties().isAddressableForDependencies())
+        return true;
+    }
+  }
+  return false;
+}
+
 #ifndef NDEBUG
 static bool areOnlyAbstractionDifferent(CanType type1, CanType type2) {
   assert(type1->isLegalSILType());

lib/SILOptimizer/IPO/CapturePropagation.cpp

@@ -496,7 +496,7 @@ static SILFunction *getSpecializedWithDeadParams(
         FuncBuilder.getModule().getSwiftModule(),
         FuncBuilder.getModule().isWholeModule(), ApplySite(), Specialized,
         PAI->getSubstitutionMap(), Specialized->getSerializedKind(),
-        /* ConvertIndirectToDirect */ false, /*dropMetatypeArgs=*/false);
+        /* ConvertIndirectToDirect */ false, /*dropUnusedArguments=*/false);
     GenericFuncSpecializer FuncSpecializer(FuncBuilder,
                                            Specialized,
                                            ReInfo.getClonerParamSubstitutionMap(),

lib/SILOptimizer/IPO/UsePrespecialized.cpp

@@ -96,7 +96,7 @@ bool UsePrespecialized::replaceByPrespecialized(SILFunction &F) {
     ReabstractionInfo ReInfo(M.getSwiftModule(), M.isWholeModule(), AI,
                              ReferencedF, Subs, IsNotSerialized,
                              /*ConvertIndirectToDirect=*/ true,
-                             /*dropMetatypeArgs=*/ false);
+                             /*dropUnusedArguments=*/ false);
 
     if (!ReInfo.canBeSpecialized())
       continue;

lib/SILOptimizer/Utils/Generics.cpp

@@ -455,6 +455,82 @@ static bool shouldNotSpecialize(SILFunction *Callee, SILFunction *Caller,
   return false;
 }
 
+// Addressable parameters cannot be dropped because the address may
+// escape. They also can't be promoted to direct convention, so there
+// is no danger in preserving them.
+static bool canConvertArg(CanSILFunctionType substType, unsigned paramIdx,
+                          SILFunction *caller) {
+  return !substType->isAddressable(paramIdx, caller);
+}
+
+// If there is no read from an indirect argument, this argument has to be
+// dropped. At the call site the store to the argument's memory location could
+// have been removed (based on the callee's memory effects).  Therefore,
+// converting such an unused indirect argument to a direct argument, would load
+// an uninitialized value at the call site.  This would lead to verifier errors
+// and in worst case to a miscompile because IRGen can implicitly use dead
+// arguments, e.g. for getting the type of a class reference.
+static bool canDropUnusedArg(ApplySite apply, SILFunction *callee,
+                             CanSILFunctionType substType,
+                             unsigned paramIdx) {
+  FullApplySite fas = apply.asFullApplySite();
+  if (!fas) {
+    return false;
+  }
+  Operand &op = fas.getOperandsWithoutIndirectResults()[paramIdx];
+  return !callee->argumentMayRead(&op, op.get());
+}
+
+static bool isUsedAsDynamicSelf(SILArgument *arg) {
+  for (Operand *use : arg->getUses()) {
+    if (use->isTypeDependent())
+      return true;
+  }
+  return false;
+}
+
+static bool canDropMetatypeArg(ApplySite apply, SILFunction *callee,
+                               unsigned paramIdx) {
+  if (!callee->isDefinition())
+    return false;
+
+  unsigned calleeArgIdx =
+    apply.getSubstCalleeConv().getSILArgIndexOfFirstParam() + paramIdx;
+  SILArgument *calleeArg = callee->getArguments()[calleeArgIdx];
+
+  if (isUsedAsDynamicSelf(calleeArg))
+    return false;
+
+  if (calleeArg->getType().getASTType()->hasDynamicSelfType())
+    return false;
+
+  // We don't drop metatype arguments of not applied arguments (in case of
+  // `partial_apply`).
+  unsigned firstAppliedArgIdx = apply.getCalleeArgIndexOfFirstAppliedArg();
+  if (firstAppliedArgIdx > calleeArgIdx)
+      return false;
+
+  auto mt = calleeArg->getType().castTo<MetatypeType>();
+  if (mt->hasRepresentation()
+      && mt->getRepresentation() == MetatypeRepresentation::Thin) {
+    return true;
+  }
+  // If the passed thick metatype value is not a `metatype` instruction
+  // we don't know the real metatype at runtime. It's not necessarily the
+  // same as the declared metatype. It could e.g. be an upcast of a class
+  // metatype.
+  SILValue callerArg = apply.getArguments()[calleeArgIdx - firstAppliedArgIdx];
+  if (isa<MetatypeInst>(callerArg))
+    return true;
+
+  // But: if the metatype is not used in the callee we don't have to care
+  // what metatype value is passed. We can just remove it.
+  if (onlyHaveDebugUses(calleeArg))
+    return true;
+
+  return false;
+}
+
 /// Prepares the ReabstractionInfo object for further processing and checks
 /// if the current function can be specialized at all.
 /// Returns false, if the current function cannot be specialized.
@@ -771,7 +847,7 @@ void ReabstractionInfo::createSubstitutedAndSpecializedTypes() {
   for (SILParameterInfo PI : SubstitutedType->getParameters()) {
     auto IdxToInsert = IdxForParam;
     ++IdxForParam;
-    unsigned argIdx = i++;
+    unsigned paramIdx = i++;
 
     SILFunctionConventions substConv(SubstitutedType, getModule());
     TypeCategory tc = getParamTypeCategory(PI, substConv, getResilienceExpansion());
@@ -782,22 +858,14 @@ void ReabstractionInfo::createSubstitutedAndSpecializedTypes() {
     case ParameterConvention::Indirect_In_CXX:
     case ParameterConvention::Indirect_In:
     case ParameterConvention::Indirect_In_Guaranteed: {
-      if (Callee && Apply && dropUnusedArguments) {
-        // If there is no read from an indirect argument, this argument has to
-        // be dropped. At the call site the store to the argument's memory location
-        // could have been removed (based on the callee's memory effects).
-        // Therefore, converting such an unused indirect argument to a direct
-        // argument, would load an uninitialized value at the call site.
-        // This would lead to verifier errors and in worst case to a miscompile
-        // because IRGen can implicitly use dead arguments, e.g. for getting the
-        // type of a class reference.
-        if (FullApplySite fas = Apply.asFullApplySite()) {
-          Operand &op = fas.getOperandsWithoutIndirectResults()[argIdx];
-          if (!Callee->argumentMayRead(&op, op.get())) {
-            droppedArguments.set(IdxToInsert);
-            break;
-          }
-        }
+      if (Apply && !canConvertArg(SubstitutedType, paramIdx,
+                                  Apply.getFunction())) {
+        continue;
+      }
+      if (Callee && Apply && dropUnusedArguments
+          && canDropUnusedArg(Apply, Callee, SubstitutedType, paramIdx)) {
+        droppedArguments.set(IdxToInsert);
+        break;
       }
       Conversions.set(IdxToInsert);
       if (tc == LoadableAndTrivial)
@@ -822,8 +890,10 @@ void ReabstractionInfo::createSubstitutedAndSpecializedTypes() {
     case ParameterConvention::Direct_Unowned:
     case ParameterConvention::Direct_Guaranteed: {
       CanType ty = PI.getInterfaceType();
-      if (dropUnusedArguments && isa<MetatypeType>(ty) && !ty->hasArchetype())
+      if (dropUnusedArguments && isa<MetatypeType>(ty) && !ty->hasArchetype()
+          && Apply && Callee && canDropMetatypeArg(Apply, Callee, paramIdx)) {
         droppedArguments.set(IdxToInsert);
+      }
       break;
     }
     }
@@ -2916,7 +2986,8 @@ static bool createPrespecialized(StringRef UnspecializedName,
   ReabstractionInfo ReInfo(M.getSwiftModule(), M.isWholeModule(), ApplySite(),
                            UnspecFunc, Apply.getSubstitutionMap(),
                            IsNotSerialized,
-                           /*ConvertIndirectToDirect= */true, /*dropMetatypeArgs=*/ false);
+                           /*ConvertIndirectToDirect= */true,
+                           /*dropUnusedArguments=*/ false);
 
   if (!ReInfo.canBeSpecialized())
     return false;
@@ -3005,7 +3076,7 @@ static bool usePrespecialized(
                                       funcBuilder.getModule().isWholeModule(), apply, refF,
                                       apply.getSubstitutionMap(), IsNotSerialized,
                                       /*ConvertIndirectToDirect=*/ true,
-                                      /*dropMetatypeArgs=*/ false);
+                                      /*dropUnusedArguments=*/ false);
 
   for (auto *SA : refF->getSpecializeAttrs()) {
     if (!SA->isExported())
@@ -3149,7 +3220,8 @@ static bool usePrespecialized(
           funcBuilder.getModule().getSwiftModule(),
           funcBuilder.getModule().isWholeModule(), apply, refF, newSubstMap,
           apply.getFunction()->getSerializedKind(),
-          /*ConvertIndirectToDirect=*/ true, /*dropMetatypeArgs=*/ false, nullptr);
+          /*ConvertIndirectToDirect=*/ true,
+          /*dropUnusedArguments=*/ false, nullptr);
 
       if (layoutReInfo.getSpecializedType() == reInfo.getSpecializedType()) {
         layoutMatches.push_back(
@@ -3209,57 +3281,6 @@ static bool usePrespecialized(
   return false;
 }
 
-static bool isUsedAsDynamicSelf(SILArgument *arg) {
-  for (Operand *use : arg->getUses()) {
-    if (use->isTypeDependent())
-      return true;
-  }
-  return false;
-}
-
-static bool canDropMetatypeArgs(ApplySite apply, SILFunction *callee) {
-  if (!callee->isDefinition())
-    return false;
-
-  auto calleeArgs = callee->getArguments();
-  unsigned firstAppliedArgIdx = apply.getCalleeArgIndexOfFirstAppliedArg();
-  for (unsigned calleeArgIdx = 0; calleeArgIdx < calleeArgs.size(); ++calleeArgIdx) {
-    SILArgument *calleeArg = calleeArgs[calleeArgIdx];
-    auto mt = calleeArg->getType().getAs<MetatypeType>();
-    if (!mt)
-      continue;
-
-    if (isUsedAsDynamicSelf(calleeArg))
-      return false;
-
-    if (calleeArg->getType().getASTType()->hasDynamicSelfType())
-      return false;
-
-    // We don't drop metatype arguments of not applied arguments (in case of `partial_apply`).
-    if (firstAppliedArgIdx > calleeArgIdx)
-      return false;
-
-    if (mt->hasRepresentation() && mt->getRepresentation() == MetatypeRepresentation::Thin)
-      continue;
-
-    // If the passed thick metatype value is not a `metatype` instruction
-    // we don't know the real metatype at runtime. It's not necessarily the
-    // same as the declared metatype. It could e.g. be an upcast of a class
-    // metatype.
-    SILValue callerArg = apply.getArguments()[calleeArgIdx - firstAppliedArgIdx];
-    if (isa<MetatypeInst>(callerArg))
-      continue;
-
-    // But: if the metatype is not used in the callee we don't have to care
-    // what metatype value is passed. We can just remove it.
-    if (callee->isDefinition() && onlyHaveDebugUses(calleeArg))
-      continue;
-
-    return false;
-  }
-  return true;
-}
-
 void swift::trySpecializeApplyOfGeneric(
     SILOptFunctionBuilder &FuncBuilder,
     ApplySite Apply, DeadInstructionSet &DeadApplies,
@@ -3312,7 +3333,7 @@ void swift::trySpecializeApplyOfGeneric(
                            FuncBuilder.getModule().isWholeModule(), Apply, RefF,
                            Apply.getSubstitutionMap(), serializedKind,
                            /*ConvertIndirectToDirect=*/ true,
-                           /*dropMetatypeArgs=*/ canDropMetatypeArgs(Apply, RefF),
+                           /*dropUnusedArguments=*/ true,
                            &ORE);
   if (!ReInfo.canBeSpecialized())
     return;

lib/SILOptimizer/Utils/OptimizerBridging.cpp

@@ -206,7 +206,7 @@ OptionalBridgedFunction BridgedPassContext::specializeFunction(BridgedFunction f
   ReabstractionInfo ReInfo(mod->getSwiftModule(), mod->isWholeModule(),
                            ApplySite(), origFunc, subs, IsNotSerialized,
                            /*ConvertIndirectToDirect=*/true,
-                           /*dropMetatypeArgs=*/false);
+                           /*dropUnusedArguments=*/false);
 
   if (!ReInfo.canBeSpecialized()) {
     return {nullptr};