Merge pull request #40291 from eeckstein/cow-opt-improvements

Array related optimization improvements

Author: eeckstein

include/swift/SIL/InstructionUtils.h

@@ -26,8 +26,6 @@ namespace swift {
 /// nothing left to strip.
 SILValue getUnderlyingObject(SILValue V);
 
-SILValue getUnderlyingObjectStopAtMarkDependence(SILValue V);
-
 SILValue stripSinglePredecessorArgs(SILValue V);
 
 /// Return the underlying SILValue after stripping off all casts from the

include/swift/SIL/SILBridging.h

@@ -237,6 +237,8 @@ BridgedInstruction SILBuilder_createBuiltinBinaryFunction(
           BridgedType operandType, BridgedType resultType, BridgedValueArray arguments);
 BridgedInstruction SILBuilder_createCondFail(BridgedInstruction insertionPoint,
           BridgedLocation loc, BridgedValue condition, BridgedStringRef messge);
+BridgedInstruction SILBuilder_createIntegerLiteral(BridgedInstruction insertionPoint,
+          BridgedLocation loc, BridgedType type, SwiftInt value);
 
 SWIFT_END_NULLABILITY_ANNOTATIONS
 

include/swift/SIL/SILNodes.def

@@ -469,7 +469,7 @@ ABSTRACT_VALUE_AND_INST(SingleValueInstruction, ValueBase, SILInstruction)
                       LiteralInst, None, DoesNotRelease)
     BRIDGED_SINGLE_VALUE_INST(GlobalValueInst, global_value,
                       LiteralInst, None, DoesNotRelease)
-    SINGLE_VALUE_INST(IntegerLiteralInst, integer_literal,
+    BRIDGED_SINGLE_VALUE_INST(IntegerLiteralInst, integer_literal,
                       LiteralInst, None, DoesNotRelease)
     SINGLE_VALUE_INST(FloatLiteralInst, float_literal,
                       LiteralInst, None, DoesNotRelease)

include/swift/SILOptimizer/PassManager/Passes.def

@@ -431,6 +431,7 @@ PASS(PruneVTables, "prune-vtables",
      "Mark class methods that do not require vtable dispatch")
 PASS_RANGE(AllPasses, AADumper, PruneVTables)
 
+SWIFT_INSTRUCTION_PASS(BeginCOWMutationInst, "simplify-begin_cow_mutation")
 SWIFT_INSTRUCTION_PASS_WITH_LEGACY(GlobalValueInst, "simplify-global_value")
 SWIFT_INSTRUCTION_PASS_WITH_LEGACY(StrongRetainInst, "simplify-strong_retain")
 SWIFT_INSTRUCTION_PASS_WITH_LEGACY(StrongReleaseInst, "simplify-strong_release")

include/swift/SILOptimizer/Utils/LoadStoreOptUtils.h

@@ -379,21 +379,37 @@ class LSLocation : public LSBase {
   static void reduce(LSLocation Base, SILModule *Mod,
                      TypeExpansionContext context, LSLocationList &Locs);
 
+  /// Gets the base address for `v`.
+  /// If `stopAtImmutable` is true, the base address is only calculated up to
+  /// a `ref_element_addr [immutable]` or a `ref_tail_addr [immutable]`.
+  /// Return the base address and true if such an immutable class projection
+  /// is found.
+  static std::pair<SILValue, bool>
+  getBaseAddressOrObject(SILValue v, bool stopAtImmutable);
+
   /// Enumerate the given Mem LSLocation.
-  static void enumerateLSLocation(TypeExpansionContext context, SILModule *M,
+  /// If `stopAtImmutable` is true, the base address is only calculated up to
+  /// a `ref_element_addr [immutable]` or a `ref_tail_addr [immutable]`.
+  /// Returns true if it's an immutable location.
+  static bool enumerateLSLocation(TypeExpansionContext context, SILModule *M,
                                   SILValue Mem,
                                   std::vector<LSLocation> &LSLocationVault,
                                   LSLocationIndexMap &LocToBit,
                                   LSLocationBaseMap &BaseToLoc,
-                                  TypeExpansionAnalysis *TE);
+                                  TypeExpansionAnalysis *TE,
+                                  bool stopAtImmutable);
 
   /// Enumerate all the locations in the function.
+  /// If `stopAtImmutable` is true, the base addresses are only calculated up to
+  /// a `ref_element_addr [immutable]` or a `ref_tail_addr [immutable]`.
   static void enumerateLSLocations(SILFunction &F,
                                    std::vector<LSLocation> &LSLocationVault,
                                    LSLocationIndexMap &LocToBit,
                                    LSLocationBaseMap &BaseToLoc,
                                    TypeExpansionAnalysis *TE,
-                                   std::pair<int, int> &LSCount);
+                                   bool stopAtImmutable,
+                                   int &numLoads, int &numStores,
+                                   bool &immutableLoadsFound);
 };
 
 static inline llvm::hash_code hash_value(const LSLocation &L) {

lib/SIL/Utils/InstructionUtils.cpp

@@ -62,18 +62,6 @@ SILValue swift::getUnderlyingObject(SILValue v) {
   }
 }
 
-SILValue swift::getUnderlyingObjectStopAtMarkDependence(SILValue v) {
-  while (true) {
-    SILValue v2 = stripCastsWithoutMarkDependence(v);
-    v2 = stripAddressProjections(v2);
-    v2 = stripIndexingInsts(v2);
-    v2 = lookThroughOwnershipInsts(v2);
-    if (v2 == v)
-      return v2;
-    v = v2;
-  }
-}
-
 /// Return the underlying SILValue after stripping off identity SILArguments if
 /// we belong to a BB with one predecessor.
 SILValue swift::stripSinglePredecessorArgs(SILValue V) {

lib/SIL/Utils/SILBridging.cpp

@@ -513,3 +513,11 @@ BridgedInstruction SILBuilder_createCondFail(BridgedInstruction insertionPoint,
   return {builder.createCondFail(getRegularLocation(loc),
     castToSILValue(condition), getStringRef(messge))};
 }
+
+BridgedInstruction SILBuilder_createIntegerLiteral(BridgedInstruction insertionPoint,
+          BridgedLocation loc, BridgedType type, SwiftInt value) {
+  SILBuilder builder(castToInst(insertionPoint), getSILDebugScope(loc));
+  return {builder.createIntegerLiteral(getRegularLocation(loc),
+                                       getSILType(type), value)};
+}
+

lib/SILOptimizer/Analysis/MemoryBehavior.cpp

@@ -566,6 +566,17 @@ static SILValue getBeginScopeInst(SILValue V) {
   if (BorrowedValue borrowedObj = getSingleBorrowIntroducingValue(object)) {
     return borrowedObj.value;
   }
+  if (!object->getFunction()->hasOwnership()) {
+    // In non-OSSA, do a quick check if the object is a guaranteed function
+    // argument.
+    // Note that in OSSA, getSingleBorrowIntroducingValue will detect a
+    // guaranteed argument.
+    SILValue root = findOwnershipReferenceAggregate(object);
+    if (auto *funcArg = dyn_cast<SILFunctionArgument>(root)) {
+      if (funcArg->getArgumentConvention().isGuaranteedConvention())
+        return funcArg;
+    }
+  }
   return SILValue();
 }
 

lib/SILOptimizer/SILCombiner/SILCombinerMiscVisitors.cpp

@@ -804,6 +804,18 @@ SILInstruction *SILCombiner::optimizeLoadFromStringLiteral(LoadInst *LI) {
   return Builder.createIntegerLiteral(LI->getLoc(), LI->getType(), str[index]);
 }
 
+static bool isShiftRightByAtLeastOne(SILInstruction *inst) {
+  auto *bi = dyn_cast<BuiltinInst>(inst);
+  if (!bi)
+    return false;
+  if (bi->getBuiltinInfo().ID != BuiltinValueKind::LShr)
+    return false;
+  auto *shiftVal = dyn_cast<IntegerLiteralInst>(bi->getArguments()[1]);
+  if (!shiftVal)
+    return false;
+  return shiftVal->getValue().isStrictlyPositive();
+}
+
 /// Returns true if \p LI loads a zero integer from the empty Array, Dictionary
 /// or Set singleton.
 static bool isZeroLoadFromEmptyCollection(SingleValueInstruction *LI) {
@@ -826,15 +838,23 @@ static bool isZeroLoadFromEmptyCollection(SingleValueInstruction *LI) {
       }
       case ValueKind::StructElementAddrInst: {
         auto *SEA = cast<StructElementAddrInst>(addr);
-        // For Array, we only support "count". The value of "capacityAndFlags"
-        // is not defined in the ABI and could change in another version of the
-        // runtime (the capacity must be 0, but the flags may be not 0).
-        if (SEA->getStructDecl()->getName().is("_SwiftArrayBodyStorage") &&
-            !SEA->getField()->getName().is("count")) {
-          return false;
-        }
         addr = SEA->getOperand();
-        break;
+        if (!SEA->getStructDecl()->getName().is("_SwiftArrayBodyStorage"))
+          break;
+        if (SEA->getField()->getName().is("count"))
+          break;
+        // For Array, the value of `capacityAndFlags` has only a zero capacity
+        // but not necessarily a zero flag (in fact, the flag is 1).
+        // Therefore only replace `capacityAndFlags` with zero if the flag is
+        // masked out by a right-shift of 1.
+        if (SEA->getField()->getName().is("_capacityAndFlags")) {
+          for (Operand *loadUse : LI->getUses()) {
+            if (!isShiftRightByAtLeastOne(loadUse->getUser()))
+              return false;
+          }
+          break;
+        }
+        return false;
       }
       case ValueKind::RefElementAddrInst: {
         auto *REA = cast<RefElementAddrInst>(addr);
@@ -858,6 +878,13 @@ static bool isZeroLoadFromEmptyCollection(SingleValueInstruction *LI) {
       case ValueKind::EndCOWMutationInst:
         addr = cast<SingleValueInstruction>(addr)->getOperand(0);
         break;
+      case ValueKind::MultipleValueInstructionResult:
+        if (auto *bci = dyn_cast<BeginCOWMutationInst>(
+                                              addr->getDefiningInstruction())) {
+          addr = bci->getOperand();
+          break;
+        }
+        return false;
       default:
         return false;
     }

lib/SILOptimizer/Transforms/COWOpts.cpp

@@ -58,10 +58,6 @@ namespace {
 /// The optimization can also handle def-use chains between end_cow_mutation and
 /// begin_cow_mutation which involve phi-arguments.
 ///
-/// An additional peephole optimization is performed: if the begin_cow_mutation
-/// is the only use of the end_cow_mutation, the whole pair of instructions
-/// is eliminated.
-///
 class COWOptsPass : public SILFunctionTransform {
 public:
   COWOptsPass() {}
@@ -86,18 +82,19 @@ void COWOptsPass::run() {
   if (!F->shouldOptimize())
     return;
 
-  LLVM_DEBUG(llvm::dbgs() << "*** RedundantPhiElimination on function: "
+  LLVM_DEBUG(llvm::dbgs() << "*** COW optimization on function: "
                           << F->getName() << " ***\n");
 
   AA = PM->getAnalysis<AliasAnalysis>(F);
 
   bool changed = false;
   for (SILBasicBlock &block : *F) {
-    auto iter = block.begin();
-    while (iter != block.end()) {
-      SILInstruction *inst = &*iter++;
-      if (auto *beginCOW = dyn_cast<BeginCOWMutationInst>(inst))
-        changed |= optimizeBeginCOW(beginCOW);
+  
+    for (SILInstruction &inst : block) {
+      if (auto *beginCOW = dyn_cast<BeginCOWMutationInst>(&inst)) {
+        if (optimizeBeginCOW(beginCOW))
+          changed = true;
+      }
     }
   }
 
@@ -212,20 +209,6 @@ bool COWOptsPass::optimizeBeginCOW(BeginCOWMutationInst *BCM) {
                                     BCM->getUniquenessResult()->getType(), 1);
   BCM->getUniquenessResult()->replaceAllUsesWith(IL);
 
-  // Try the peephole optimization: remove an end_cow_mutation/begin_cow_mutation
-  // pair completely if the begin_cow_mutation is the only use of
-  // end_cow_mutation.
-  if (auto *singleEndCOW = dyn_cast<EndCOWMutationInst>(BCM->getOperand())) {
-    assert(endCOWMutationInsts.size() == 1 &&
-           *endCOWMutationInsts.begin() == singleEndCOW);
-    if (singleEndCOW->hasOneUse()) {
-      BCM->getBufferResult()->replaceAllUsesWith(singleEndCOW->getOperand());
-      BCM->eraseFromParent();
-      singleEndCOW->eraseFromParent();
-      return true;
-    }
-  }
-
   for (EndCOWMutationInst *ECM : endCOWMutationInsts) {
     // This is important for other optimizations: The code is now relying on
     // the buffer to be unique.

lib/SILOptimizer/Transforms/DeadStoreElimination.cpp

@@ -1182,14 +1182,17 @@ void DSEContext::runIterativeDSE() {
 }
 
 bool DSEContext::run() {
-  std::pair<int, int> LSCount = std::make_pair(0, 0);
+  int numLoads = 0, numStores = 0;
+  bool immutableLoadsFound = false;
   // Walk over the function and find all the locations accessed by
   // this function.
   LSLocation::enumerateLSLocations(*F, LocationVault, LocToBitIndex,
-                                   BaseToLocIndex, TE, LSCount);
+                                   BaseToLocIndex, TE,
+                                   /*stopAtImmutable*/ false,
+                                   numLoads, numStores, immutableLoadsFound);
 
   // Check how to optimize this function.
-  ProcessKind Kind = getProcessFunctionKind(LSCount.second);
+  ProcessKind Kind = getProcessFunctionKind(numStores);
 
   // We do not optimize this function at all.
   if (Kind == ProcessKind::ProcessNone)

lib/SILOptimizer/Transforms/RedundantLoadElimination.cpp

@@ -486,14 +486,21 @@ class RLEContext {
   /// If set, RLE ignores loads from that array type.
   NominalTypeDecl *ArrayType;
 
+  /// Se to true if loads with a `ref_element_addr [immutable]` or
+  /// `ref_tail_addr [immutable]` base address are found.
+  bool immutableLoadsFound = false;
+
+  /// Only optimize loads with a base address of  `ref_element_addr [immutable]`
+  /// `ref_tail_addr [immutable]`.
+  bool onlyImmutableLoads;
+
 #ifndef NDEBUG
   SILPrintContext printCtx;
 #endif
 
 public:
-  RLEContext(SILFunction *F, SILPassManager *PM, AliasAnalysis *AA,
-             TypeExpansionAnalysis *TE, PostOrderFunctionInfo *PO,
-             EpilogueARCFunctionInfo *EAFI, bool disableArrayLoads);
+  RLEContext(SILFunction *F, SILPassManager *PM,
+             bool disableArrayLoads, bool onlyImmutableLoads);
 
   RLEContext(const RLEContext &) = delete;
   RLEContext(RLEContext &&) = delete;
@@ -504,6 +511,8 @@ class RLEContext {
   /// Entry point to redundant load elimination.
   bool run();
 
+  bool shouldOptimizeImmutableLoads() const { return immutableLoadsFound; }
+
   SILFunction *getFunction() const { return Fn; }
 
   /// Use a set of ad hoc rules to tell whether we should run a pessimistic
@@ -570,6 +579,11 @@ class RLEContext {
           LI->getType().getNominalOrBoundGenericNominal() != ArrayType) {
         return LI;
       }
+      if (onlyImmutableLoads &&
+          !LSLocation::getBaseAddressOrObject(LI->getOperand(),
+                                             /*stopAtImmutable*/ true).second) {
+        return nullptr;
+      }
     }
     return nullptr;
   }
@@ -1200,14 +1214,17 @@ void BlockState::dump(RLEContext &Ctx) {
 //                          RLEContext Implementation
 //===----------------------------------------------------------------------===//
 
-RLEContext::RLEContext(SILFunction *F, SILPassManager *PM, AliasAnalysis *AA,
-                       TypeExpansionAnalysis *TE, PostOrderFunctionInfo *PO,
-                       EpilogueARCFunctionInfo *EAFI, bool disableArrayLoads)
-    : Fn(F), PM(PM), AA(AA), TE(TE), PO(PO), EAFI(EAFI), BBToLocState(F),
-      BBWithLoads(F),
+RLEContext::RLEContext(SILFunction *F, SILPassManager *PM,
+                       bool disableArrayLoads, bool onlyImmutableLoads)
+    : Fn(F), PM(PM), AA(PM->getAnalysis<AliasAnalysis>(F)),
+      TE(PM->getAnalysis<TypeExpansionAnalysis>()),
+      PO(PM->getAnalysis<PostOrderAnalysis>()->get(F)),
+      EAFI(PM->getAnalysis<EpilogueARCAnalysis>()->get(F)),
+      BBToLocState(F), BBWithLoads(F),
       ArrayType(disableArrayLoads
                     ? F->getModule().getASTContext().getArrayDecl()
-                    : nullptr)
+                    : nullptr),
+      onlyImmutableLoads(onlyImmutableLoads)
 #ifndef NDEBUG
       ,
       printCtx(llvm::dbgs(), /*Verbose=*/false, /*Sorted=*/true)
@@ -1567,14 +1584,15 @@ bool RLEContext::run() {
   // Phase 4. we perform the redundant load elimination.
   // Walk over the function and find all the locations accessed by
   // this function.
-  std::pair<int, int> LSCount = std::make_pair(0, 0); 
+  int numLoads = 0, numStores = 0;
   LSLocation::enumerateLSLocations(*Fn, LocationVault,
                                    LocToBitIndex,
                                    BaseToLocIndex, TE,
-                                   LSCount);
+                                   /*stopAtImmutable*/ onlyImmutableLoads,
+                                   numLoads, numStores, immutableLoadsFound);
 
   // Check how to optimize this function.
-  ProcessKind Kind = getProcessFunctionKind(LSCount.first, LSCount.second);
+  ProcessKind Kind = getProcessFunctionKind(numLoads, numStores);
 
   // We do not optimize this function at all.
   if (Kind == ProcessKind::ProcessNone)
@@ -1681,15 +1699,21 @@ class RedundantLoadElimination : public SILFunctionTransform {
     LLVM_DEBUG(llvm::dbgs() << "*** RLE on function: " << F->getName()
                             << " ***\n");
 
-    auto *AA = PM->getAnalysis<AliasAnalysis>(F);
-    auto *TE = PM->getAnalysis<TypeExpansionAnalysis>();
-    auto *PO = PM->getAnalysis<PostOrderAnalysis>()->get(F);
-    auto *EAFI = PM->getAnalysis<EpilogueARCAnalysis>()->get(F);
-
-    RLEContext RLE(F, PM, AA, TE, PO, EAFI, disableArrayLoads);
+    RLEContext RLE(F, PM, disableArrayLoads,
+                   /*onlyImmutableLoads*/ false);
     if (RLE.run()) {
       invalidateAnalysis(SILAnalysis::InvalidationKind::Instructions);
     }
+    if (RLE.shouldOptimizeImmutableLoads()) {
+      /// Re-running RLE with cutting base addresses off at
+      /// `ref_element_addr [immutable]` or `ref_tail_addr [immutable]` can
+      /// expose additional opportunities.
+      RLEContext RLE2(F, PM, disableArrayLoads,
+                      /*onlyImmutableLoads*/ true);
+      if (RLE2.run()) {
+        invalidateAnalysis(SILAnalysis::InvalidationKind::Instructions);
+      }
+    }
   }
 };