[loop-rotate] In OSSA, instead of creating address phis, sneak the address through the phi using a RawPointer.

In OSSA, we do not allow for address phis, but in certain cases the logic of
LoopRotate really wants them. To work around this issue, I added some code in
this PR to loop rotate that as a post-pass fixes up any address phis by
inserting address <-> raw pointer adapters and changing the address phi to
instead be of raw pointer type.

Author: gottesmm

include/swift/SIL/SILArgument.h

@@ -215,6 +215,12 @@ class SILPhiArgument : public SILArgument {
   /// this will be guaranteed to return a valid SILValue.
   SILValue getIncomingPhiValue(SILBasicBlock *predBlock) const;
 
+  /// If this argument is a true phi, return the operand in the \p predBLock
+  /// associated with an incoming value.
+  ///
+  /// \returns the operand or nullptr if this is not a true phi.
+  Operand *getIncomingPhiOperand(SILBasicBlock *predBlock) const;
+
   /// If this argument is a phi, populate `OutArray` with the incoming phi
   /// values for each predecessor BB. If this argument is not a phi, return
   /// false.

lib/SIL/IR/SILArgument.cpp

@@ -161,6 +161,12 @@ bool SILPhiArgument::getIncomingPhiValues(
   return true;
 }
 
+Operand *SILPhiArgument::getIncomingPhiOperand(SILBasicBlock *predBlock) const {
+  if (!isPhiArgument())
+    return nullptr;
+  return getIncomingPhiOperandForPred(getParent(), predBlock, getIndex());
+}
+
 bool SILPhiArgument::getIncomingPhiOperands(
     SmallVectorImpl<Operand *> &returnedPhiOperands) const {
   if (!isPhiArgument())

lib/SILOptimizer/LoopTransforms/LoopRotate.cpp

@@ -125,7 +125,9 @@ static void mapOperands(SILInstruction *inst,
 static void updateSSAForUseOfValue(
     SILSSAUpdater &updater, SmallVectorImpl<SILPhiArgument *> &insertedPhis,
     const llvm::DenseMap<ValueBase *, SILValue> &valueMap,
-    SILBasicBlock *Header, SILBasicBlock *EntryCheckBlock, SILValue Res) {
+    SILBasicBlock *Header, SILBasicBlock *EntryCheckBlock, SILValue Res,
+    SmallVectorImpl<std::pair<SILBasicBlock *, unsigned>>
+        &accumulatedAddressPhis) {
   // Find the mapped instruction.
   assert(valueMap.count(Res) && "Expected to find value in map!");
   SILValue MappedValue = valueMap.find(Res)->second;
@@ -159,49 +161,96 @@ static void updateSSAForUseOfValue(
            && "The entry check block should dominate the header");
     updater.rewriteUse(*use);
   }
-  // Canonicalize inserted phis to avoid extra BB Args.
+
+  // Canonicalize inserted phis to avoid extra BB Args and if we find an address
+  // phi, stash it so we can handle it after we are done rewriting.
+  bool hasOwnership = Header->getParent()->hasOwnership();
   for (SILPhiArgument *arg : insertedPhis) {
     if (SILValue inst = replaceBBArgWithCast(arg)) {
       arg->replaceAllUsesWith(inst);
       // DCE+SimplifyCFG runs as a post-pass cleanup.
       // DCE replaces dead arg values with undef.
       // SimplifyCFG deletes the dead BB arg.
+      continue;
     }
+
+    // If we didn't simplify and have an address phi, stash the value so we can
+    // fix it up.
+    if (hasOwnership && arg->getType().isAddress())
+      accumulatedAddressPhis.emplace_back(arg->getParent(), arg->getIndex());
   }
 }
 
-static void
-updateSSAForUseOfInst(SILSSAUpdater &updater,
-                      SmallVectorImpl<SILPhiArgument *> &insertedPhis,
-                      const llvm::DenseMap<ValueBase *, SILValue> &valueMap,
-                      SILBasicBlock *header, SILBasicBlock *entryCheckBlock,
-                      SILInstruction *inst) {
+static void updateSSAForUseOfInst(
+    SILSSAUpdater &updater, SmallVectorImpl<SILPhiArgument *> &insertedPhis,
+    const llvm::DenseMap<ValueBase *, SILValue> &valueMap,
+    SILBasicBlock *header, SILBasicBlock *entryCheckBlock, SILInstruction *inst,
+    SmallVectorImpl<std::pair<SILBasicBlock *, unsigned>>
+        &accumulatedAddressPhis) {
   for (auto result : inst->getResults())
     updateSSAForUseOfValue(updater, insertedPhis, valueMap, header,
-                           entryCheckBlock, result);
+                           entryCheckBlock, result, accumulatedAddressPhis);
 }
 
 /// Rewrite the code we just created in the preheader and update SSA form.
 static void rewriteNewLoopEntryCheckBlock(
     SILBasicBlock *header, SILBasicBlock *entryCheckBlock,
     const llvm::DenseMap<ValueBase *, SILValue> &valueMap) {
-  SmallVector<SILPhiArgument *, 4> insertedPhis;
+  SmallVector<std::pair<SILBasicBlock *, unsigned>, 8> accumulatedAddressPhis;
+  SmallVector<SILPhiArgument *, 8> insertedPhis;
   SILSSAUpdater updater(&insertedPhis);
 
-  // Fix PHIs (incoming arguments).
-  for (auto *arg : header->getArguments())
+  // Fix PHIs (incoming arguments). We iterate by index in case we replace the
+  // phi argument so we do not invalidate iterators.
+  for (unsigned i : range(header->getNumArguments())) {
+    auto *arg = header->getArguments()[i];
     updateSSAForUseOfValue(updater, insertedPhis, valueMap, header,
-                           entryCheckBlock, arg);
+                           entryCheckBlock, arg, accumulatedAddressPhis);
+  }
 
   auto instIter = header->begin();
 
   // The terminator might change from under us.
   while (instIter != header->end()) {
     auto &inst = *instIter;
     updateSSAForUseOfInst(updater, insertedPhis, valueMap, header,
-                          entryCheckBlock, &inst);
+                          entryCheckBlock, &inst, accumulatedAddressPhis);
     ++instIter;
   }
+
+  // Then see if any of our phis were address phis. In such a case, rewrite the
+  // address to be a smuggled through raw pointer. We do this late to
+  // conservatively not interfere with the previous code's invariants.
+  //
+  // We also translate the phis into a BasicBlock, Index form so we are careful
+  // with invalidation issues around branches/args.
+  auto rawPointerTy =
+      SILType::getRawPointerType(header->getParent()->getASTContext());
+  auto rawPointerUndef = SILUndef::get(rawPointerTy, header->getModule());
+  auto loc = RegularLocation::getAutoGeneratedLocation();
+  while (!accumulatedAddressPhis.empty()) {
+    SILBasicBlock *block;
+    unsigned argIndex;
+    std::tie(block, argIndex) = accumulatedAddressPhis.pop_back_val();
+    auto *arg = cast<SILPhiArgument>(block->getArgument(argIndex));
+    assert(arg->getType().isAddress() && "Not an address phi?!");
+    for (auto *predBlock : block->getPredecessorBlocks()) {
+      Operand *predUse = arg->getIncomingPhiOperand(predBlock);
+      SILBuilderWithScope builder(predUse->getUser());
+      auto *newIncomingValue =
+          builder.createAddressToPointer(loc, predUse->get(), rawPointerTy);
+      predUse->set(newIncomingValue);
+    }
+    SILBuilderWithScope builder(arg->getNextInstruction());
+    SILType oldArgType = arg->getType();
+    auto *phiShim = builder.createPointerToAddress(
+        loc, rawPointerUndef, oldArgType, true /*isStrict*/,
+        false /*is invariant*/);
+    arg->replaceAllUsesWith(phiShim);
+    SILArgument *newArg = block->replacePhiArgument(
+        argIndex, rawPointerTy, OwnershipKind::None, nullptr);
+    phiShim->setOperand(newArg);
+  }
 }
 
 /// Update the dominator tree after rotating the loop.

test/SILOptimizer/looprotate_ossa.sil

@@ -472,3 +472,58 @@ bb2:
  return %1 : $()
 }
 
+// Make sure that we do not create address phis.
+//
+// CHECK-LABEL: sil [ossa] @addressPhiFixUp : $@convention(thin) (Builtin.RawPointer) -> Builtin.RawPointer {
+// CHECK: bb1:
+// CHECK:   [[NEXT:%.*]] = address_to_pointer {{%.*}} : $*UInt8 to $Builtin.RawPointer
+// CHECK:   cond_br {{%.*}}, bb3, bb2
+//
+// CHECK: bb2:
+// CHECK-NEXT: br bb7(%0 : $Builtin.RawPointer)
+//
+// CHECK: bb3:
+// CHECK-NEXT: br bb4([[NEXT]] : $Builtin.RawPointer)
+//
+// CHECK: bb4([[ARG:%.*]] :
+// CHECK:  [[CAST_BACK:%.*]] = pointer_to_address [[ARG]] : $Builtin.RawPointer to [strict] $*UInt8
+// CHECK:  [[GEP:%.*]] = index_addr [[CAST_BACK]] :
+// CHECK:  [[CAST_ROUND_TRIP_START:%.*]] = address_to_pointer [[GEP]]
+// CHECK:  [[CAST_ROUND_TRIP_END:%.*]] = pointer_to_address [[CAST_ROUND_TRIP_START]] : $Builtin.RawPointer to [strict] $*UInt8
+// CHECK:  [[BACK_TO_RAWPOINTER:%.*]] = address_to_pointer [[CAST_ROUND_TRIP_END]]
+// CHECK:  cond_br {{%.*}}, bb6, bb5
+//
+// CHECK: bb5:
+// CHECK-NEXT: br bb7([[CAST_ROUND_TRIP_START]] :
+//
+// CHECK: bb6:
+// CHECK-NEXT: br bb4([[BACK_TO_RAWPOINTER]] :
+//
+// CHECK: bb7([[RESULT:%.*]] :
+// CHECK-NEXT: return [[RESULT]]
+// CHECK: } // end sil function 'addressPhiFixUp'
+sil [ossa] @addressPhiFixUp : $@convention(thin) (Builtin.RawPointer) -> Builtin.RawPointer {
+bb0(%0 : $Builtin.RawPointer):
+  br bb1
+
+bb1:
+  br bb2(%0 : $Builtin.RawPointer)
+
+bb2(%1 : $Builtin.RawPointer):
+  %2 = pointer_to_address %1 : $Builtin.RawPointer to [strict] $*UInt8
+  %3 = load [trivial] %2 : $*UInt8
+  %4 = destructure_struct %3 : $UInt8
+  %5 = integer_literal $Builtin.Int64, 0
+  %6 = builtin "zextOrBitCast_Int8_Int64"(%4 : $Builtin.Int8) : $Builtin.Int64
+  %7 = builtin "cmp_eq_Int64"(%6 : $Builtin.Int64, %5 : $Builtin.Int64) : $Builtin.Int1
+  cond_br %7, bb3, bb4
+
+bb3:
+  %8 = integer_literal $Builtin.Word, 1
+  %9 = index_addr %2 : $*UInt8, %8 : $Builtin.Word
+  %10 = address_to_pointer %9 : $*UInt8 to $Builtin.RawPointer
+  br bb2(%10 : $Builtin.RawPointer)
+
+bb4:
+  return %1 : $Builtin.RawPointer
+}