Merge pull request #35461 from gottesmm/ossa-interior-ptr-fixup

[ownership] Implement Interior Pointer handling API for RAUWing addresses

Author: gottesmm

include/swift/SIL/InstructionUtils.h

@@ -27,6 +27,10 @@ SILValue getUnderlyingObject(SILValue V);
 
 SILValue getUnderlyingObjectStopAtMarkDependence(SILValue V);
 
+/// Given an address look through address to address projections and indexing
+/// insts.
+SILValue getUnderlyingObjectStoppingAtObjectToAddrProjections(SILValue v);
+
 SILValue stripSinglePredecessorArgs(SILValue V);
 
 /// Return the underlying SILValue after stripping off all casts from the
@@ -55,8 +59,14 @@ SILValue stripClassCasts(SILValue V);
 
 /// Return the underlying SILValue after stripping off all address projection
 /// instructions.
+///
+/// FIXME: Today address projections are referring to the result of the
+/// projection and doesn't consider the operand. Should we change this?
 SILValue stripAddressProjections(SILValue V);
 
+/// Look through any projections that transform an address -> an address.
+SILValue lookThroughAddressToAddressProjections(SILValue v);
+
 /// Return the underlying SILValue after stripping off all aggregate projection
 /// instructions.
 ///

include/swift/SIL/OwnershipUtils.h

@@ -672,7 +672,17 @@ struct InteriorPointerOperand {
   /// requirements to ensure that the underlying class is alive at all use
   /// points.
   bool getImplicitUses(SmallVectorImpl<Operand *> &foundUses,
-                       std::function<void(Operand *)> *onError = nullptr);
+                       std::function<void(Operand *)> *onError = nullptr) {
+    return getImplicitUsesForAddress(getProjectedAddress(), foundUses, onError);
+  }
+
+  /// The algorithm that is used to determine what the verifier will consider to
+  /// be implicit uses of the given address. Used to implement \see
+  /// getImplicitUses.
+  static bool
+  getImplicitUsesForAddress(SILValue address,
+                            SmallVectorImpl<Operand *> &foundUses,
+                            std::function<void(Operand *)> *onError = nullptr);
 
   Operand *operator->() { return operand; }
   const Operand *operator->() const { return operand; }

include/swift/SIL/Projection.h

@@ -391,6 +391,10 @@ class Projection {
     }
   }
 
+  static bool isAddressToAddressProjection(SILValue v) {
+    return isAddressProjection(v) && !isObjectToAddressProjection(v);
+  }
+
   /// Returns true if this instruction projects from an object type into an
   /// address subtype.
   static bool isObjectToAddressProjection(SILValue V) {

include/swift/SILOptimizer/Analysis/SimplifyInstruction.h

@@ -27,31 +27,30 @@ namespace swift {
 class SILInstruction;
 class InstModCallbacks;
 
-/// Try to simplify the specified instruction, performing local
-/// analysis of the operands of the instruction, without looking at its uses
-/// (e.g. constant folding).  If a simpler result can be found, it is
-/// returned, otherwise a null SILValue is returned.
-///
-/// This is assumed to implement read-none transformations.
-SILValue simplifyInstruction(SILInstruction *I);
-
 /// Replace an instruction with a simplified result and erase it. If the
 /// instruction initiates a scope, do not replace the end of its scope; it will
 /// be deleted along with its parent.
 ///
-/// If it is nonnull, eraseNotify will be called before each instruction is
-/// deleted.
-///
-/// If it is nonnull and inst is in OSSA, newInstNotify will be called with each
-/// new instruction inserted to compensate for ownership.
-///
 /// NOTE: When OSSA is enabled this API assumes OSSA is properly formed and will
 /// insert compensating instructions.
 SILBasicBlock::iterator
 replaceAllSimplifiedUsesAndErase(SILInstruction *I, SILValue result,
                                  InstModCallbacks &callbacks,
                                  DeadEndBlocks *deadEndBlocks = nullptr);
 
+/// Attempt to map \p inst to a simplified result. Upon success, replace \p inst
+/// with this simplified result and erase \p inst. If the instruction initiates
+/// a scope, do not replace the end of its scope; it will be deleted along with
+/// its parent.
+///
+/// NOTE: When OSSA is enabled this API assumes OSSA is properly formed and will
+/// insert compensating instructions.
+/// NOTE: When \p I is in an OSSA function, this fails to optimize if \p
+/// deadEndBlocks is null.
+SILBasicBlock::iterator simplifyAndReplaceAllSimplifiedUsesAndErase(
+    SILInstruction *I, InstModCallbacks &callbacks,
+    DeadEndBlocks *deadEndBlocks = nullptr);
+
 // Simplify invocations of builtin operations that may overflow.
 /// All such operations return a tuple (result, overflow_flag).
 /// This function try to simplify such operations, but returns only a

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -400,6 +400,10 @@ class InstModCallbacks {
   }
 
   bool hadCallbackInvocation() const { return wereAnyCallbacksInvoked; }
+
+  /// Set \p wereAnyCallbacksInvoked to false. Useful if one wants to reuse an
+  /// InstModCallback in between iterations.
+  void resetHadCallbackInvocation() { wereAnyCallbacksInvoked = false; }
 };
 
 /// Get all consumed arguments of a partial_apply.

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -19,6 +19,7 @@
 #ifndef SWIFT_SILOPTIMIZER_UTILS_OWNERSHIPOPTUTILS_H
 #define SWIFT_SILOPTIMIZER_UTILS_OWNERSHIPOPTUTILS_H
 
+#include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/SILModule.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
@@ -28,31 +29,98 @@ namespace swift {
 // Defined in BasicBlockUtils.h
 struct JointPostDominanceSetComputer;
 
+/// A struct that contains context shared in between different operation +
+/// "ownership fixup" utilities. Please do not put actual methods on this, it is
+/// meant to be composed with.
 struct OwnershipFixupContext {
   Optional<InstModCallbacks> inlineCallbacks;
   InstModCallbacks &callbacks;
   DeadEndBlocks &deBlocks;
   JointPostDominanceSetComputer &jointPostDomSetComputer;
 
+  /// Extra state initialized by OwnershipRAUWFixupHelper::get() that we use
+  /// when RAUWing addresses. This ensures we do not need to recompute this
+  /// state when we perform the actual RAUW.
+  struct AddressFixupContext {
+    /// When determining if we need to perform an address pointer fixup, we
+    /// compute all transitive address uses of oldValue. If we find that we do
+    /// need this fixed up, then we will copy our interior pointer base value
+    /// and use this to seed that new lifetime.
+    SmallVector<Operand *, 8> allAddressUsesFromOldValue;
+
+    /// This is the interior pointer operand that the new value we want to RAUW
+    /// is transitively derived from and enables us to know the underlying
+    /// borrowed base value that we need to lifetime extend.
+    InteriorPointerOperand intPtrOp;
+
+    void clear() {
+      allAddressUsesFromOldValue.clear();
+      intPtrOp = InteriorPointerOperand();
+    }
+  };
+  AddressFixupContext extraAddressFixupInfo;
+
   OwnershipFixupContext(InstModCallbacks &callbacks, DeadEndBlocks &deBlocks,
-                        JointPostDominanceSetComputer &inputJPDComputer)
-      : inlineCallbacks(), callbacks(callbacks), deBlocks(deBlocks),
-        jointPostDomSetComputer(inputJPDComputer) {}
+                        JointPostDominanceSetComputer &jointPostDomSetComputer)
+      : callbacks(callbacks), deBlocks(deBlocks),
+        jointPostDomSetComputer(jointPostDomSetComputer) {}
+
+  void clear() {
+    jointPostDomSetComputer.clear();
+    extraAddressFixupInfo.allAddressUsesFromOldValue.clear();
+    extraAddressFixupInfo.intPtrOp = InteriorPointerOperand();
+  }
+
+private:
+  /// Helper method called to determine if we discovered we needed interior
+  /// pointer fixups while simplifying.
+  bool needsInteriorPointerFixups() const {
+    return bool(extraAddressFixupInfo.intPtrOp);
+  }
+};
 
-  OwnershipFixupContext(DeadEndBlocks &deBlocks,
-                        JointPostDominanceSetComputer &inputJPDComputer)
-      : inlineCallbacks(InstModCallbacks()), callbacks(*inlineCallbacks),
-        deBlocks(deBlocks), jointPostDomSetComputer(inputJPDComputer) {}
+/// A utility composed ontop of OwnershipFixupContext that knows how to RAUW a
+/// value or a single value instruction with a new value and then fixup
+/// ownership invariants afterwards.
+class OwnershipRAUWHelper {
+  OwnershipFixupContext *ctx;
+  SingleValueInstruction *oldValue;
+  SILValue newValue;
 
-  SILBasicBlock::iterator
-  replaceAllUsesAndErase(SingleValueInstruction *oldValue, SILValue newValue);
+public:
+  OwnershipRAUWHelper() : ctx(nullptr), oldValue(nullptr), newValue(nullptr) {}
 
-  /// We can not RAUW all old values with new values.
+  /// Return an instance of this class if we can perform the specific RAUW
+  /// operation ignoring if the types line up. Returns None otherwise.
   ///
-  /// Namely, we do not support RAUWing values with ValueOwnershipKind::None
-  /// that have uses that do not require ValueOwnershipKind::None or
-  /// ValueOwnershipKind::Any.
-  static bool canFixUpOwnershipForRAUW(SILValue oldValue, SILValue newValue);
+  /// DISCUSSION: We do not check that the types line up here so that we can
+  /// allow for our users to transform our new value in ways that preserve
+  /// ownership at \p oldValue before we perform the actual RAUW. If \p newValue
+  /// is an object, any instructions in the chain of transforming instructions
+  /// from \p newValue at \p oldValue's must be forwarding. If \p newValue is an
+  /// address, then these transforms can only transform the address into a
+  /// derived address.
+  OwnershipRAUWHelper(OwnershipFixupContext &ctx,
+                      SingleValueInstruction *oldValue, SILValue newValue);
+
+  /// Returns true if this helper was initialized into a valid state.
+  operator bool() const { return isValid(); }
+  bool isValid() const { return bool(ctx) && bool(oldValue) && bool(newValue); }
+
+  /// Perform the actual RAUW. We require that \p newValue and \p oldValue have
+  /// the same type at this point (in contrast to when calling
+  /// OwnershipRAUWFixupHelper::get()).
+  ///
+  /// This is so that we can avoid creating "forwarding" transformation
+  /// instructions before we know if we can perform the RAUW. Any such
+  /// "forwarding" transformation must be performed upon \p newValue at \p
+  /// oldValue's insertion point so that we can then here RAUW the transformed
+  /// \p newValue.
+  SILBasicBlock::iterator perform();
+
+private:
+  SILBasicBlock::iterator replaceAddressUses(SingleValueInstruction *oldValue,
+                                             SILValue newValue);
 };
 
 } // namespace swift

lib/SIL/Utils/InstructionUtils.cpp

@@ -62,6 +62,20 @@ SILValue swift::getUnderlyingObject(SILValue v) {
   }
 }
 
+SILValue
+swift::getUnderlyingObjectStoppingAtObjectToAddrProjections(SILValue v) {
+  if (!v->getType().isAddress())
+    return SILValue();
+
+  while (true) {
+    auto v2 = lookThroughAddressToAddressProjections(v);
+    v2 = stripIndexingInsts(v2);
+    if (v2 == v)
+      return v2;
+    v = v2;
+  }
+}
+
 SILValue swift::getUnderlyingObjectStopAtMarkDependence(SILValue v) {
   while (true) {
     SILValue v2 = stripCastsWithoutMarkDependence(v);
@@ -207,6 +221,15 @@ SILValue swift::stripAddressProjections(SILValue V) {
   }
 }
 
+SILValue swift::lookThroughAddressToAddressProjections(SILValue v) {
+  while (true) {
+    v = stripSinglePredecessorArgs(v);
+    if (!Projection::isAddressToAddressProjection(v))
+      return v;
+    v = cast<SingleValueInstruction>(v)->getOperand(0);
+  }
+}
+
 SILValue swift::stripValueProjections(SILValue V) {
   while (true) {
     V = stripSinglePredecessorArgs(V);

lib/SIL/Utils/OwnershipUtils.cpp

@@ -520,10 +520,9 @@ bool BorrowedValue::visitInteriorPointerOperands(
 //                           InteriorPointerOperand
 //===----------------------------------------------------------------------===//
 
-bool InteriorPointerOperand::getImplicitUses(
-    SmallVectorImpl<Operand *> &foundUses,
+bool InteriorPointerOperand::getImplicitUsesForAddress(
+    SILValue projectedAddress, SmallVectorImpl<Operand *> &foundUses,
     std::function<void(Operand *)> *onError) {
-  SILValue projectedAddress = getProjectedAddress();
   SmallVector<Operand *, 8> worklist(projectedAddress->getUses());
 
   bool foundError = false;

lib/SILOptimizer/Analysis/SimplifyInstruction.cpp

@@ -300,14 +300,6 @@ SILValue InstSimplifier::visitAddressToPointerInst(AddressToPointerInst *ATPI) {
 }
 
 SILValue InstSimplifier::visitPointerToAddressInst(PointerToAddressInst *PTAI) {
-  // (pointer_to_address strict (address_to_pointer x)) -> x
-  //
-  // NOTE: We can not perform this optimization in OSSA without dealing with
-  // interior pointers since we may be escaping an interior pointer address from
-  // a borrow scope.
-  if (PTAI->getFunction()->hasOwnership())
-    return SILValue();
-
   // If this address is not strict, then it cannot be replaced by an address
   // that may be strict.
   if (auto *ATPI = dyn_cast<AddressToPointerInst>(PTAI->getOperand()))
@@ -755,7 +747,8 @@ swift::replaceAllSimplifiedUsesAndErase(SILInstruction *i, SILValue result,
   if (svi->getFunction()->hasOwnership()) {
     JointPostDominanceSetComputer computer(*deadEndBlocks);
     OwnershipFixupContext ctx{callbacks, *deadEndBlocks, computer};
-    return ctx.replaceAllUsesAndErase(svi, result);
+    OwnershipRAUWHelper helper(ctx, svi, result);
+    return helper.perform();
   }
   return replaceAllUsesAndErase(svi, result, callbacks);
 }
@@ -779,19 +772,39 @@ SILValue swift::simplifyOverflowBuiltinInstruction(BuiltinInst *BI) {
 ///
 /// NOTE: We assume that the insertion point associated with the SILValue must
 /// dominate \p i.
-SILValue swift::simplifyInstruction(SILInstruction *i) {
-  SILValue result = InstSimplifier().visit(i);
-  if (!result)
-    return SILValue();
-
-  // If we have a result, we know that we must have a single value instruction
-  // by assumption since we have not implemented support in the rest of inst
-  // simplify for non-single value instructions. We put the cast here so that
-  // this code is not updated at this point in time.
-  auto *svi = cast<SingleValueInstruction>(i);
-  if (svi->getFunction()->hasOwnership())
-    if (!OwnershipFixupContext::canFixUpOwnershipForRAUW(svi, result))
-      return SILValue();
+static SILValue simplifyInstruction(SILInstruction *i) {
+  return InstSimplifier().visit(i);
+}
 
-  return result;
+SILBasicBlock::iterator swift::simplifyAndReplaceAllSimplifiedUsesAndErase(
+    SILInstruction *i, InstModCallbacks &callbacks,
+    DeadEndBlocks *deadEndBlocks) {
+  auto next = std::next(i->getIterator());
+  auto *svi = dyn_cast<SingleValueInstruction>(i);
+  if (!svi)
+    return next;
+  SILValue result = simplifyInstruction(i);
+
+  // If we fail to simplify or the simplified value returned is our passed in
+  // value, just return std::next since we can't simplify.
+  if (!result || svi == result)
+    return next;
+
+  if (!svi->getFunction()->hasOwnership())
+    return replaceAllUsesAndErase(svi, result, callbacks);
+
+  // If we weren't passed a dead end blocks, we can't optimize without ownership
+  // enabled.
+  if (!deadEndBlocks)
+    return next;
+
+  JointPostDominanceSetComputer computer(*deadEndBlocks);
+  OwnershipFixupContext ctx{callbacks, *deadEndBlocks, computer};
+  OwnershipRAUWHelper helper(ctx, svi, result);
+
+  // If our RAUW helper is invalid, we do not support RAUWing this case, so
+  // just return next.
+  if (!helper.isValid())
+    return next;
+  return helper.perform();
 }

lib/SILOptimizer/Mandatory/OwnershipModelEliminator.cpp

@@ -479,13 +479,12 @@ static bool stripOwnership(SILFunction &func) {
     auto value = visitor.instructionsToSimplify.pop_back_val();
     if (!value.hasValue())
       continue;
-    if (SILValue newValue = simplifyInstruction(*value)) {
-      InstModCallbacks callbacks([&](SILInstruction *instToErase) {
-        visitor.eraseInstruction(instToErase);
-      });
-      replaceAllSimplifiedUsesAndErase(*value, newValue, callbacks);
-      madeChange = true;
-    }
+    InstModCallbacks callbacks([&](SILInstruction *instToErase) {
+      visitor.eraseInstruction(instToErase);
+    });
+    // We are no longer in OSSA, so we don't need to pass in a deBlocks.
+    simplifyAndReplaceAllSimplifiedUsesAndErase(*value, callbacks);
+    madeChange |= callbacks.hadCallbackInvocation();
   }
 
   return madeChange;