Sink address projections in ArrayPropertyOpt and enable verification to ensure we have no address phis

include/swift/SILOptimizer/Utils/BasicBlockOptUtils.h

@@ -135,14 +135,20 @@ bool canCloneTerminator(TermInst *termInst);
 /// BasicBlockCloner handles this internally.
 class SinkAddressProjections {
   // Projections ordered from last to first in the chain.
-  SmallVector<SingleValueInstruction *, 4> projections;
+  SmallVector<SingleValueInstruction *, 4> oldProjections;
+  // Cloned projections to avoid address phis.
+  SmallVectorImpl<SingleValueInstruction *> *newProjections;
   SmallSetVector<SILValue, 4> inBlockDefs;
 
   // Transient per-projection data for use during cloning.
   SmallVector<Operand *, 4> usesToReplace;
   llvm::SmallDenseMap<SILBasicBlock *, Operand *, 4> firstBlockUse;
 
 public:
+  SinkAddressProjections(
+      SmallVectorImpl<SingleValueInstruction *> *newProjections = nullptr)
+      : newProjections(newProjections) {}
+
   /// Check for an address projection chain ending at \p inst. Return true if
   /// the given instruction is successfully analyzed.
   ///
@@ -163,6 +169,7 @@ class SinkAddressProjections {
   ArrayRef<SILValue> getInBlockDefs() const {
     return inBlockDefs.getArrayRef();
   }
+
   /// Clone the chain of projections at their use sites.
   ///
   /// Return true if anything was done.

lib/SIL/Verifier/SILVerifier.cpp

@@ -1065,24 +1065,6 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     return InstNumbers[a] < InstNumbers[b];
   }
 
-  // FIXME: For sanity, address-type phis should be prohibited at all SIL
-  // stages. However, the optimizer currently breaks the invariant in three
-  // places:
-  // 1. Normal Simplify CFG during conditional branch simplification
-  //    (sneaky jump threading).
-  // 2. Simplify CFG via Jump Threading.
-  // 3. Loop Rotation.
-  //
-  // BasicBlockCloner::canCloneInstruction and sinkAddressProjections is
-  // designed to avoid this issue, we just need to make sure all passes use it
-  // correctly.
-  //
-  // Minimally, we must prevent address-type phis as long as access markers are
-  // preserved. A goal is to preserve access markers in OSSA.
-  bool prohibitAddressPhis() {
-    return F.hasOwnership();
-  }
-
   void visitSILPhiArgument(SILPhiArgument *arg) {
     // Verify that the `isPhiArgument` property is sound:
     // - Phi arguments come from branches.
@@ -1106,7 +1088,7 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
                 "All phi argument inputs must be from branches.");
       }
     }
-    if (arg->isPhi() && prohibitAddressPhis()) {
+    if (arg->isPhi()) {
       // As a property of well-formed SIL, we disallow address-type
       // phis. Supporting them would prevent reliably reasoning about the
       // underlying storage of memory access. This reasoning is important for

lib/SILOptimizer/LoopTransforms/ArrayPropertyOpt.cpp

@@ -53,19 +53,19 @@
 #define DEBUG_TYPE "array-property-opt"
 
 #include "ArrayOpt.h"
+#include "swift/SIL/CFG.h"
+#include "swift/SIL/DebugUtils.h"
+#include "swift/SIL/InstructionUtils.h"
+#include "swift/SIL/LoopInfo.h"
+#include "swift/SIL/Projection.h"
+#include "swift/SIL/SILCloner.h"
 #include "swift/SILOptimizer/Analysis/ArraySemantic.h"
 #include "swift/SILOptimizer/Analysis/LoopAnalysis.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
+#include "swift/SILOptimizer/Utils/BasicBlockOptUtils.h"
 #include "swift/SILOptimizer/Utils/CFGOptUtils.h"
 #include "swift/SILOptimizer/Utils/LoopUtils.h"
 #include "swift/SILOptimizer/Utils/SILSSAUpdater.h"
-#include "swift/SIL/CFG.h"
-#include "swift/SIL/DebugUtils.h"
-#include "swift/SIL/InstructionUtils.h"
-#include "swift/SIL/Projection.h"
-#include "swift/SIL/LoopInfo.h"
-#include "swift/SIL/BasicBlockBits.h"
-#include "swift/SIL/SILCloner.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
@@ -86,6 +86,8 @@ class ArrayPropertiesAnalysis {
   SILBasicBlock *Preheader;
   DominanceInfo *DomTree;
 
+  SinkAddressProjections sinkProj;
+
   llvm::DenseMap<SILFunction *, uint32_t> InstCountCache;
   llvm::SmallSet<SILValue, 16> HoistableArray;
 
@@ -169,13 +171,18 @@ class ArrayPropertiesAnalysis {
 
     bool FoundHoistable = false;
     uint32_t LoopInstCount = 0;
+
     for (auto *BB : Loop->getBlocks()) {
       for (auto &Inst : *BB) {
         // Can't clone alloc_stack instructions whose dealloc_stack is outside
         // the loop.
         if (!canDuplicateLoopInstruction(Loop, &Inst))
           return false;
 
+        if (!sinkProj.analyzeAddressProjections(&Inst)) {
+          return false;
+        }
+
         ArraySemanticsCall ArrayPropsInst(&Inst, "array.props", true);
         if (!ArrayPropsInst)
           continue;
@@ -512,10 +519,11 @@ class RegionCloner : public SILCloner<RegionCloner> {
                          SILSSAUpdater &SSAUp) {
     // Collect outside uses.
     SmallVector<UseWrapper, 16> UseList;
-    for (auto Use : V->getUses())
+    for (auto Use : V->getUses()) {
       if (!isBlockCloned(Use->getUser()->getParent())) {
         UseList.push_back(UseWrapper(Use));
       }
+    }
     if (UseList.empty())
       return;
 
@@ -532,15 +540,40 @@ class RegionCloner : public SILCloner<RegionCloner> {
 
   void updateSSAForm() {
     SILSSAUpdater SSAUp;
+    SmallVector<SingleValueInstruction *, 4> newProjections;
+    SinkAddressProjections sinkProj(&newProjections);
+
     for (auto *origBB : originalPreorderBlocks()) {
       // Update outside used phi values.
-      for (auto *arg : origBB->getArguments())
+      for (auto *arg : origBB->getArguments()) {
         updateSSAForValue(origBB, arg, SSAUp);
+      }
 
       // Update outside used instruction values.
       for (auto &inst : *origBB) {
-        for (auto result : inst.getResults())
-          updateSSAForValue(origBB, result, SSAUp);
+        for (auto result : inst.getResults()) {
+          bool success = sinkProj.analyzeAddressProjections(&inst);
+          assert(success);
+          // Sink address projections by cloning to avoid address phis.
+          sinkProj.cloneProjections();
+
+          // If no new projections were created, update ssa for the result only.
+          if (newProjections.empty()) {
+            updateSSAForValue(origBB, result, SSAUp);
+            continue;
+          }
+
+          for (auto *newProj : newProjections) {
+            // Operand values of new projections may need ssa update.
+            for (auto opVal : newProj->getOperandValues()) {
+              if (!isBlockCloned(opVal->getParentBlock())) {
+                continue;
+              }
+              updateSSAForValue(origBB, opVal, SSAUp);
+            }
+          }
+          newProjections.clear();
+        }
       }
     }
   }

lib/SILOptimizer/Utils/BasicBlockOptUtils.cpp

@@ -223,7 +223,7 @@ void BasicBlockCloner::sinkAddressProjections() {
 //
 // Return true on success, even if projections is empty.
 bool SinkAddressProjections::analyzeAddressProjections(SILInstruction *inst) {
-  projections.clear();
+  oldProjections.clear();
   inBlockDefs.clear();
 
   SILBasicBlock *bb = inst->getParent();
@@ -237,7 +237,7 @@ bool SinkAddressProjections::analyzeAddressProjections(SILInstruction *inst) {
     }
     if (auto *addressProj = dyn_cast<SingleValueInstruction>(def)) {
       if (addressProj->isPure()) {
-        projections.push_back(addressProj);
+        oldProjections.push_back(addressProj);
         return true;
       }
     }
@@ -252,12 +252,12 @@ bool SinkAddressProjections::analyzeAddressProjections(SILInstruction *inst) {
       return false;
   }
   // Recurse upward through address projections.
-  for (unsigned idx = 0; idx < projections.size(); ++idx) {
+  for (unsigned idx = 0; idx < oldProjections.size(); ++idx) {
     // Only one address result/operand can be handled per instruction.
-    if (projections.size() != idx + 1)
+    if (oldProjections.size() != idx + 1)
       return false;
 
-    for (SILValue operandVal : projections[idx]->getOperandValues())
+    for (SILValue operandVal : oldProjections[idx]->getOperandValues())
       if (!pushOperandVal(operandVal))
         return false;
   }
@@ -267,13 +267,13 @@ bool SinkAddressProjections::analyzeAddressProjections(SILInstruction *inst) {
 // Clone the projections gathered by 'analyzeAddressProjections' at
 // their use site outside this block.
 bool SinkAddressProjections::cloneProjections() {
-  if (projections.empty())
+  if (oldProjections.empty())
     return false;
 
-  SILBasicBlock *bb = projections.front()->getParent();
+  SILBasicBlock *bb = oldProjections.front()->getParent();
   // Clone projections in last-to-first order.
-  for (unsigned idx = 0; idx < projections.size(); ++idx) {
-    auto *oldProj = projections[idx];
+  for (unsigned idx = 0; idx < oldProjections.size(); ++idx) {
+    auto *oldProj = oldProjections[idx];
     assert(oldProj->getParent() == bb);
     // Reset transient per-projection sets.
     usesToReplace.clear();
@@ -297,9 +297,12 @@ bool SinkAddressProjections::cloneProjections() {
       auto *useBB = use->getUser()->getParent();
       auto *firstUse = firstBlockUse.lookup(useBB);
       SingleValueInstruction *newProj;
-      if (use == firstUse)
+      if (use == firstUse) {
         newProj = cast<SingleValueInstruction>(oldProj->clone(use->getUser()));
-      else {
+        if (newProjections) {
+          newProjections->push_back(newProj);
+        }
+      } else {
         newProj = cast<SingleValueInstruction>(firstUse->get());
         assert(newProj->getParent() == useBB);
         newProj->moveFront(useBB);

lib/SILOptimizer/Utils/SILSSAUpdater.cpp

@@ -86,20 +86,30 @@ areIdentical(llvm::DenseMap<SILBasicBlock *, SILValue> &availableValues) {
     return true;
   }
 
-  auto *mvir =
-      dyn_cast<MultipleValueInstructionResult>(availableValues.begin()->second);
-  if (!mvir)
-    return false;
+  if (auto *mvir = dyn_cast<MultipleValueInstructionResult>(
+          availableValues.begin()->second)) {
+    for (auto value : availableValues) {
+      auto *result = dyn_cast<MultipleValueInstructionResult>(value.second);
+      if (!result)
+        return false;
+      if (!result->getParent()->isIdenticalTo(mvir->getParent()) ||
+          result->getIndex() != mvir->getIndex()) {
+        return false;
+      }
+    }
+    return true;
+  }
 
+  auto *firstArg = cast<SILArgument>(availableValues.begin()->second);
   for (auto value : availableValues) {
-    auto *result = dyn_cast<MultipleValueInstructionResult>(value.second);
-    if (!result)
+    auto *arg = dyn_cast<SILArgument>(value.second);
+    if (!arg)
       return false;
-    if (!result->getParent()->isIdenticalTo(mvir->getParent()) ||
-        result->getIndex() != mvir->getIndex()) {
+    if (arg != firstArg) {
       return false;
     }
   }
+
   return true;
 }