[opt] promote stack alloc_ref to object

lib/IRGen/GenClass.cpp

@@ -102,6 +102,10 @@ namespace {
     StructLayout *createLayoutWithTailElems(IRGenModule &IGM,
                                             SILType classType,
                                             ArrayRef<SILType> tailTypes) const;
+
+    using HeapTypeInfo<ClassTypeInfo>::initialize;
+    void initialize(IRGenFunction &IGF, Explosion &e, Address addr,
+                    bool isOutlined) const override;
   };
 } // end anonymous namespace
 
@@ -483,6 +487,29 @@ ClassTypeInfo::getClassLayout(IRGenModule &IGM, SILType classType,
   return *Layout;
 }
 
+void ClassTypeInfo::initialize(IRGenFunction &IGF, Explosion &src, Address addr,
+                               bool isOutlined) const {
+  // If the address is a poitner to the exploded type then we can simply emit a
+  // store and be done.
+  auto *exploded = src.getAll().front();
+  if (exploded->getType() == addr->getType()->getPointerElementType()) {
+    IGF.Builder.CreateStore(exploded, addr);
+    (void)src.claimNext();
+    return;
+  }
+
+  // If both are the same (address) type then just emit a memcpy.
+  if (exploded->getType() == addr->getType()) {
+    IGF.emitMemCpy(addr.getAddress(), exploded, getFixedSize(),
+                   addr.getAlignment());
+    (void)src.claimNext();
+    return;
+  }
+
+  // Otherwise, bail to the default implementation.
+  HeapTypeInfo<ClassTypeInfo>::initialize(IGF, src, addr, isOutlined);
+}
+
 /// Cast the base to i8*, apply the given inbounds offset (in bytes,
 /// as a size_t), and cast to a pointer to the given type.
 llvm::Value *IRGenFunction::emitByteOffsetGEP(llvm::Value *base,

lib/IRGen/IRGenSIL.cpp

@@ -877,9 +877,7 @@ class IRGenSILFunction :
   void visitDestroyValueInst(DestroyValueInst *i);
   void visitAutoreleaseValueInst(AutoreleaseValueInst *i);
   void visitSetDeallocatingInst(SetDeallocatingInst *i);
-  void visitObjectInst(ObjectInst *i)  {
-    llvm_unreachable("object instruction cannot appear in a function");
-  }
+  void visitObjectInst(ObjectInst *i);
   void visitStructInst(StructInst *i);
   void visitTupleInst(TupleInst *i);
   void visitEnumInst(EnumInst *i);
@@ -3560,6 +3558,49 @@ void IRGenSILFunction::visitDestroyValueInst(swift::DestroyValueInst *i) {
       .consume(*this, in, getDefaultAtomicity());
 }
 
+void IRGenSILFunction::visitObjectInst(swift::ObjectInst *i) {
+  SILType objType = i->getType().getObjectType();
+  // TODO: Currently generic classes aren't allowed but in the future we could
+  // support this.
+  assert(!objType.getClassOrBoundGenericClass()->getAsGenericContext() ||
+         !objType.getClassOrBoundGenericClass()
+                 ->getAsGenericContext()
+                 ->isGeneric() &&
+             "Generics are not yet supported");
+
+  const auto &typeInfo = cast<LoadableTypeInfo>(getTypeInfo(objType));
+  llvm::Value *metadata = emitClassHeapMetadataRef(
+      *this, objType.getASTType(), MetadataValueType::TypeMetadata,
+      MetadataState::Complete);
+  // TODO: this shouldn't be a stack alloc but, in order to maintain
+  // compatibility with alloc_ref we have to be a pointer.
+  Address alloca =
+      createAlloca(typeInfo.getStorageType()->getPointerElementType(),
+                   typeInfo.getFixedAlignment());
+  auto classAddr = Builder.CreateBitCast(alloca, IGM.RefCountedPtrTy);
+  auto classVal = emitInitStackObjectCall(metadata, classAddr.getAddress(),
+                                          "reference.new");
+  classVal = Builder.CreateBitCast(classVal, typeInfo.getStorageType());
+  // Match each property in the class decl to elements in the object
+  // instruction.
+  auto propsArr =
+      i->getType().getClassOrBoundGenericClass()->getStoredProperties();
+  SmallVector<VarDecl *, 8> props(propsArr.begin(), propsArr.end());
+  for (SILValue elt : i->getAllElements()) {
+    auto prop = props.pop_back_val();
+    auto elementExplosion = getLoweredExplosion(elt);
+    auto propType = IGM.getLoweredType(prop->getType());
+    const auto &propTypeInfo = cast<LoadableTypeInfo>(getTypeInfo(propType));
+    auto propAddr = projectPhysicalClassMemberAddress(*this, classVal, objType,
+                                                      propType, prop);
+    propTypeInfo.initialize(*this, elementExplosion, propAddr, false);
+  }
+
+  Explosion e;
+  e.add(classVal);
+  setLoweredExplosion(i, e);
+}
+
 void IRGenSILFunction::visitStructInst(swift::StructInst *i) {
   Explosion out;
   for (SILValue elt : i->getElements())

lib/SIL/Verifier/SILVerifier.cpp

@@ -1809,8 +1809,12 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     checkGlobalAccessInst(GVI);
   }
 
-  void checkObjectInst(ObjectInst *) {
-    require(false, "object instruction is only allowed in a static initializer");
+  void checkObjectInst(ObjectInst *objectInst) {
+    auto *classDecl = objectInst->getType().getClassOrBoundGenericClass();
+    require(classDecl, "Type of object instruction must be a class");
+    require(!classDecl->getAsGenericContext() ||
+                !classDecl->getAsGenericContext()->isGeneric(),
+            "Generics are not yet supported in object instructions");
   }
 
   void checkIntegerLiteralInst(IntegerLiteralInst *ILI) {

lib/SILOptimizer/Transforms/DeadStoreElimination.cpp

@@ -1038,6 +1038,9 @@ void DSEContext::processLoadInst(SILInstruction *I, DSEKind Kind) {
 
 void DSEContext::processStoreInst(SILInstruction *I, DSEKind Kind) {
   auto *SI = cast<StoreInst>(I);
+  // TODO: for some reason these stores are removed when they shouldn't be.
+  if (isa<ObjectInst>(SI->getSrc()))
+    return;
   processWrite(I, SI->getSrc(), SI->getDest(), Kind);
 }
 

lib/SILOptimizer/Transforms/StackPromotion.cpp

@@ -12,6 +12,7 @@
 
 #include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/CFG.h"
+#include "swift/SIL/DebugUtils.h"
 #include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SILOptimizer/Analysis/EscapeAnalysis.h"
@@ -51,6 +52,15 @@ class StackPromotion : public SILFunctionTransform {
   /// Tries to promote the allocation \p ARI.
   bool tryPromoteAlloc(AllocRefInst *ARI, EscapeAnalysis *EA,
                        DeadEndBlocks &DEBlocks);
+
+  /// Tries to promote the allocation \p ARI to an object. This optimization
+  /// will only happen if the class type has a compiler-generated constructor
+  /// and destructor. The promotion happens by scanning all uses in dominance
+  /// order. If all members are accounted for by ref_element_addr instruction
+  /// before we find any other use, then we can use those values to promote this
+  /// alloc_ref to an object.
+  bool tryPromoteToObject(AllocRefInst *allocRef,
+                          ValueLifetimeAnalysis::Frontier &frontier);
 };
 
 void StackPromotion::run() {
@@ -84,10 +94,11 @@ void StackPromotion::run() {
 bool StackPromotion::promoteInBlock(SILBasicBlock *BB, EscapeAnalysis *EA,
                                     DeadEndBlocks &DEBlocks) {
   bool Changed = false;
-  for (auto Iter = BB->begin(); Iter != BB->end();) {
-    // The allocation instruction may be moved, so increment Iter prior to
-    // doing the optimization.
-    SILInstruction *I = &*Iter++;
+  SmallVector<SILInstruction *, 64> allInstructions;
+  for (SILInstruction &inst : *BB) {
+    allInstructions.push_back(&inst);
+  }
+  for (auto *I : allInstructions) {
     if (auto *ARI = dyn_cast<AllocRefInst>(I)) {
       // Don't stack promote any allocation inside a code region which ends up
       // in a no-return block. Such allocations may missing their final release.
@@ -107,6 +118,8 @@ bool StackPromotion::tryPromoteAlloc(AllocRefInst *ARI, EscapeAnalysis *EA,
   if (ARI->isObjC() || ARI->canAllocOnStack())
     return false;
 
+  EA->invalidate(ARI->getFunction(),
+                 swift::AliasAnalysis::InvalidationKind::Everything);
   auto *ConGraph = EA->getConnectionGraph(ARI->getFunction());
   auto *contentNode = ConGraph->getValueContent(ARI);
   if (!contentNode)
@@ -142,6 +155,9 @@ bool StackPromotion::tryPromoteAlloc(AllocRefInst *ARI, EscapeAnalysis *EA,
   }
   NumStackPromoted++;
 
+  if (tryPromoteToObject(ARI, Frontier))
+    return true;
+
   // We set the [stack] attribute in the alloc_ref.
   ARI->setStackAllocatable();
 
@@ -153,6 +169,131 @@ bool StackPromotion::tryPromoteAlloc(AllocRefInst *ARI, EscapeAnalysis *EA,
   return true;
 }
 
+static void getOrderedNonDebugUses(SILValue v, DominanceInfo *domInfo,
+                                   SmallVectorImpl<Operand *> &uses) {
+  auto unsorted = getNonDebugUses(v);
+  uses.append(unsorted.begin(), unsorted.end());
+  llvm::sort(uses, [&domInfo](Operand *a, Operand *b) {
+    return domInfo->dominates(a->getUser(), b->getUser());
+  });
+}
+
+bool StackPromotion::tryPromoteToObject(
+    AllocRefInst *allocRef, ValueLifetimeAnalysis::Frontier &frontier) {
+  DominanceInfo *domInfo =
+      PM->getAnalysis<DominanceAnalysis>()->get(allocRef->getFunction());
+  auto *classDecl = allocRef->getType().getClassOrBoundGenericClass();
+  if (!classDecl || !classDecl->getDestructor()->isImplicit() ||
+      (classDecl->getAsGenericContext() &&
+       classDecl->getAsGenericContext()->isGeneric()))
+    return false;
+
+  SmallVector<VarDecl *, 8> props;
+  for (auto *prop : classDecl->getStoredProperties()) {
+    props.push_back(prop);
+  }
+
+  SmallVector<Operand *, 24> uses;
+  getOrderedNonDebugUses(allocRef, domInfo, uses);
+
+  llvm::reverse(props);
+  SmallVector<RefElementAddrInst *, 8> propertyInitializers;
+  for (auto *use : uses) {
+    auto propRef = dyn_cast<RefElementAddrInst>(use->getUser());
+    if (!propRef)
+      return false;
+
+    if (propRef->getField() != props.pop_back_val())
+      return false;
+
+    propertyInitializers.push_back(propRef);
+
+    if (props.empty())
+      break;
+  }
+
+  if (!props.empty())
+    return false;
+
+  SmallVector<StoreInst *, 8> deadStores;
+  SmallVector<SILValue, 8> elements;
+  for (auto *init : propertyInitializers) {
+    SmallVector<Operand *, 6> refElementUses;
+    getOrderedNonDebugUses(init, domInfo, refElementUses);
+    auto frontUser = refElementUses.front()->getUser();
+    if (auto *beginAccess = dyn_cast<BeginAccessInst>(frontUser)) {
+      SmallVector<Operand *, 4> beginAccessUses;
+      getOrderedNonDebugUses(beginAccess, domInfo, beginAccessUses);
+      frontUser = beginAccessUses.front()->getUser();
+    }
+    if (auto *store = dyn_cast<StoreInst>(frontUser)) {
+      elements.push_back(store->getSrc());
+      deadStores.push_back(store);
+    } else {
+      return false;
+    }
+  }
+
+  SILInstruction *lastElement = nullptr;
+  for (auto first = elements.rbegin(); first != elements.rend(); ++first) {
+    auto inst = first->getDefiningInstruction();
+    if (!inst)
+      continue;
+
+    if (!lastElement || domInfo->dominates(lastElement, inst))
+      lastElement = inst;
+  }
+
+  // If we didn't find anything, that means that all the elements are arguments,
+  // or there aren't any elements. Either way, we know that putting where the
+  // alloc_ref is will work.
+  if (!lastElement)
+    lastElement = allocRef;
+
+  for (auto *init : propertyInitializers) {
+    if (llvm::any_of(init->getUses(), [&domInfo, &lastElement](Operand *use) {
+          return domInfo->dominates(use->getUser(), lastElement);
+        }))
+      return false;
+  }
+
+  SILBuilder builder(std::next(lastElement->getIterator()));
+  auto object = builder.createObject(allocRef->getLoc(), allocRef->getType(),
+                                     elements, elements.size());
+
+  SmallVector<Operand *, 8> users(allocRef->use_begin(), allocRef->use_end());
+  for (auto *use : users) {
+    auto user = use->getUser();
+
+    // The only user that still may not dominate object is a debug_value.
+    if (isa<SetDeallocatingInst>(user) || isa<DeallocRefInst>(user) ||
+        isa<DebugValueInst>(user) || isa<StrongReleaseInst>(user))
+      user->eraseFromParent();
+
+    if (auto ref = dyn_cast<RefElementAddrInst>(user)) {
+      // We need to move the ref_element_addr up.
+      if (domInfo->dominates(ref, object)) {
+        auto newRef = builder.createRefElementAddr(ref->getLoc(), object,
+                                                   ref->getField());
+        ref->replaceAllUsesWith(newRef);
+        ref->eraseFromParent();
+      }
+    }
+  }
+
+  for (auto *store : deadStores) {
+    store->eraseFromParent();
+  }
+
+  allocRef->replaceAllUsesWith(object);
+  allocRef->eraseFromParent();
+
+  llvm::errs() << "Promoted to object in stack. Function: "
+               << object->getFunction()->getName() << "\n";
+
+  return true;
+}
+
 } // end anonymous namespace
 
 SILTransform *swift::createStackPromotion() {

lib/SILOptimizer/Utils/SILInliner.cpp

@@ -915,7 +915,7 @@ InlineCost swift::instructionInlineCost(SILInstruction &I) {
   case SILInstructionKind::MarkUninitializedInst:
     llvm_unreachable("not valid in canonical sil");
   case SILInstructionKind::ObjectInst:
-    llvm_unreachable("not valid in a function");
+    return InlineCost::Free;
   }
 
   llvm_unreachable("Unhandled ValueKind in switch.");