[SIL Optimization] Make ArraySemantics.cpp aware of "array.uninitialized_intrinsic"

include/swift/SILOptimizer/Analysis/ArraySemantic.h

@@ -41,7 +41,8 @@ enum class ArrayCallKind {
   // a function, and it has a self parameter, make sure that it is defined
   // before this comment.
   kArrayInit,
-  kArrayUninitialized
+  kArrayUninitialized,
+  kArrayUninitializedIntrinsic
 };
 
 /// Wrapper around array semantic calls.
@@ -183,6 +184,21 @@ class ArraySemanticsCall {
   /// Can this function be inlined by the early inliner.
   bool canInlineEarly() const;
 
+  /// If this is a call to  ArrayUninitialized (or
+  /// ArrayUninitializedInstrinsic), identify the instructions that store
+  /// elements into the array indices. For every index, add the store
+  /// instruction that stores to that index to \p ElementStoreMap.
+  ///
+  /// \returns true iff this is an "array.uninitialized" semantic call, and the
+  /// stores into the array indices are identified and the \p ElementStoreMap is
+  /// populated.
+  ///
+  /// Note that this function does not support array initializations that use
+  /// copy_addr, which implies that arrays with address-only types would not
+  /// be recognized by this function as yet.
+  bool mapInitializationStores(
+      llvm::DenseMap<uint64_t, StoreInst *> &ElementStoreMap);
+
 protected:
   /// Validate the signature of this call.
   bool isValidSignature();

lib/SILOptimizer/Analysis/ArraySemantic.cpp

@@ -170,6 +170,7 @@ ArrayCallKind swift::ArraySemanticsCall::getKind() const {
                   ArrayCallKind::kArrayPropsIsNativeTypeChecked)
             .StartsWith("array.init", ArrayCallKind::kArrayInit)
             .Case("array.uninitialized", ArrayCallKind::kArrayUninitialized)
+            .Case("array.uninitialized_intrinsic", ArrayCallKind::kArrayUninitializedIntrinsic)
             .Case("array.check_subscript", ArrayCallKind::kCheckSubscript)
             .Case("array.check_index", ArrayCallKind::kCheckIndex)
             .Case("array.get_count", ArrayCallKind::kGetCount)
@@ -591,11 +592,15 @@ bool swift::ArraySemanticsCall::canInlineEarly() const {
     case ArrayCallKind::kAppendContentsOf:
     case ArrayCallKind::kReserveCapacityForAppend:
     case ArrayCallKind::kAppendElement:
+    case ArrayCallKind::kArrayUninitializedIntrinsic:
       // append(Element) calls other semantics functions. Therefore it's
       // important that it's inlined by the early inliner (which is before all
       // the array optimizations). Also, this semantics is only used to lookup
       // Array.append(Element), so inlining it does not prevent any other
       // optimization.
+      //
+      // Early inlining array.uninitialized_intrinsic semantic call helps in
+      // stack promotion.
       return true;
   }
 }
@@ -622,61 +627,57 @@ SILValue swift::ArraySemanticsCall::getInitializationCount() const {
   return SILValue();
 }
 
-SILValue swift::ArraySemanticsCall::getArrayValue() const {
-  if (getKind() == ArrayCallKind::kArrayUninitialized) {
-    TupleExtractInst *ArrayDef = nullptr;
-    for (auto *Op : SemanticsCall->getUses()) {
-      auto *TupleElt = dyn_cast<TupleExtractInst>(Op->getUser());
-      if (!TupleElt)
-        return SILValue();
-      switch (TupleElt->getFieldNo()) {
-      default:
-        return SILValue();
-      case 0: {
-          // Should only have one tuple extract after CSE.
-        if (ArrayDef)
-          return SILValue();
-        ArrayDef = TupleElt;
-        break;
-      }
-      case 1: /*Ignore the storage address */ break;
-      }
+/// Given an array semantic call \c arrayCall, if it is an "array.uninitialized"
+/// initializer, which returns a two-element tuple, return the element of the
+/// tuple at \c tupleElementIndex. Return a null SILValue if the
+/// array call is not an "array.uninitialized" initializer or if the extraction
+/// of the result tuple fails.
+static SILValue getArrayUninitializedInitResult(ArraySemanticsCall arrayCall,
+                                                unsigned tupleElementIndex) {
+  assert(tupleElementIndex <= 1 && "tupleElementIndex must be 0 or 1");
+  ArrayCallKind arrayCallKind = arrayCall.getKind();
+  if (arrayCallKind != ArrayCallKind::kArrayUninitialized &&
+      arrayCallKind != ArrayCallKind::kArrayUninitializedIntrinsic)
+    return SILValue();
+
+  // In OSSA, the call result will be extracted through a destructure_tuple
+  // instruction.
+  ApplyInst *callInst = arrayCall;
+  if (callInst->getFunction()->hasOwnership()) {
+    Operand *singleUse = callInst->getSingleUse();
+    if (!singleUse)
+      return SILValue();
+    if (DestructureTupleInst *destructTuple =
+            dyn_cast<DestructureTupleInst>(singleUse->getUser())) {
+      return destructTuple->getResult(tupleElementIndex);
     }
-    return SILValue(ArrayDef);
+    return SILValue();
+  }
+
+  // In non-OSSA, look for a tuple_extract instruction of the call result with
+  // the requested tupleElementIndex.
+  TupleExtractInst *tupleExtractInst = nullptr;
+  for (auto *op : callInst->getUses()) {
+    auto *tupleElt = dyn_cast<TupleExtractInst>(op->getUser());
+    if (!tupleElt)
+      return SILValue();
+    if (tupleElt->getFieldNo() != tupleElementIndex)
+      continue;
+    tupleExtractInst = tupleElt;
+    break;
   }
+  return SILValue(tupleExtractInst);
+}
 
-  if (getKind() == ArrayCallKind::kArrayInit)
+SILValue swift::ArraySemanticsCall::getArrayValue() const {
+  ArrayCallKind arrayCallKind = getKind();
+  if (arrayCallKind == ArrayCallKind::kArrayInit)
     return SILValue(SemanticsCall);
-
-  return SILValue();
+  return getArrayUninitializedInitResult(*this, 0);
 }
 
 SILValue swift::ArraySemanticsCall::getArrayElementStoragePointer() const {
-  if (getKind() == ArrayCallKind::kArrayUninitialized) {
-    TupleExtractInst *ArrayElementStorage = nullptr;
-    for (auto *Op : SemanticsCall->getUses()) {
-      auto *TupleElt = dyn_cast<TupleExtractInst>(Op->getUser());
-      if (!TupleElt)
-        return SILValue();
-      switch (TupleElt->getFieldNo()) {
-      default:
-        return SILValue();
-      case 0: {
-        // Ignore the array value.
-        break;
-      }
-      case 1:
-        // Should only have one tuple extract after CSE.
-        if (ArrayElementStorage)
-          return SILValue();
-        ArrayElementStorage = TupleElt;
-        break;
-      }
-    }
-    return SILValue(ArrayElementStorage);
-  }
-
-  return SILValue();
+  return getArrayUninitializedInitResult(*this, 1);
 }
 
 bool swift::ArraySemanticsCall::replaceByValue(SILValue V) {
@@ -786,3 +787,75 @@ bool swift::ArraySemanticsCall::replaceByAppendingValues(
 
   return true;
 }
+
+bool swift::ArraySemanticsCall::mapInitializationStores(
+    llvm::DenseMap<uint64_t, StoreInst *> &ElementValueMap) {
+  if (getKind() != ArrayCallKind::kArrayUninitialized &&
+      getKind() != ArrayCallKind::kArrayUninitializedIntrinsic)
+    return false;
+  SILValue ElementBuffer = getArrayElementStoragePointer();
+  if (!ElementBuffer)
+    return false;
+
+  // Match initialization stores into ElementBuffer. E.g.
+  // %83 = struct_extract %element_buffer : $UnsafeMutablePointer<Int>
+  // %84 = pointer_to_address %83 : $Builtin.RawPointer to strict $*Int
+  // store %85 to %84 : $*Int
+  // %87 = integer_literal $Builtin.Word, 1
+  // %88 = index_addr %84 : $*Int, %87 : $Builtin.Word
+  // store %some_value to %88 : $*Int
+
+  // If this an ArrayUinitializedIntrinsic then the ElementBuffer is a
+  // builtin.RawPointer. Otherwise, it is an UnsafeMutablePointer, which would
+  // be struct-extracted to obtain a builtin.RawPointer.
+  SILValue UnsafeMutablePointerExtract =
+      (getKind() == ArrayCallKind::kArrayUninitialized)
+          ? dyn_cast_or_null<StructExtractInst>(
+                getSingleNonDebugUser(ElementBuffer))
+          : ElementBuffer;
+  if (!UnsafeMutablePointerExtract)
+    return false;
+
+  auto *PointerToAddress = dyn_cast_or_null<PointerToAddressInst>(
+      getSingleNonDebugUser(UnsafeMutablePointerExtract));
+  if (!PointerToAddress)
+    return false;
+
+  // Match the stores. We can have either a store directly to the address or
+  // to an index_addr projection.
+  for (auto *Op : PointerToAddress->getUses()) {
+    auto *Inst = Op->getUser();
+
+    // Store to the base.
+    auto *SI = dyn_cast<StoreInst>(Inst);
+    if (SI && SI->getDest() == PointerToAddress) {
+      // We have already seen an entry for this index bail.
+      if (ElementValueMap.count(0))
+        return false;
+      ElementValueMap[0] = SI;
+      continue;
+    } else if (SI)
+      return false;
+
+    // Store to an index_addr projection.
+    auto *IndexAddr = dyn_cast<IndexAddrInst>(Inst);
+    if (!IndexAddr)
+      return false;
+    SI = dyn_cast_or_null<StoreInst>(getSingleNonDebugUser(IndexAddr));
+    if (!SI || SI->getDest() != IndexAddr)
+      return false;
+    auto *Index = dyn_cast<IntegerLiteralInst>(IndexAddr->getIndex());
+    if (!Index)
+      return false;
+    auto IndexVal = Index->getValue();
+    // Let's not blow up our map.
+    if (IndexVal.getActiveBits() > 16)
+      return false;
+    // Already saw an entry.
+    if (ElementValueMap.count(IndexVal.getZExtValue()))
+      return false;
+
+    ElementValueMap[IndexVal.getZExtValue()] = SI;
+  }
+  return !ElementValueMap.empty();
+}

lib/SILOptimizer/LoopTransforms/COWArrayOpt.cpp

@@ -314,6 +314,7 @@ static bool isNonMutatingArraySemanticCall(SILInstruction *Inst) {
   case ArrayCallKind::kWithUnsafeMutableBufferPointer:
   case ArrayCallKind::kArrayInit:
   case ArrayCallKind::kArrayUninitialized:
+  case ArrayCallKind::kArrayUninitializedIntrinsic:
   case ArrayCallKind::kAppendContentsOf:
   case ArrayCallKind::kAppendElement:
     return false;
@@ -662,7 +663,8 @@ bool COWArrayOpt::hasLoopOnlyDestructorSafeArrayOperations() {
         auto Kind = Sem.getKind();
         // Safe because they create new arrays.
         if (Kind == ArrayCallKind::kArrayInit ||
-            Kind == ArrayCallKind::kArrayUninitialized)
+            Kind == ArrayCallKind::kArrayUninitialized ||
+            Kind == ArrayCallKind::kArrayUninitializedIntrinsic)
           continue;
         // All array types must be the same. This is a stronger guaranteed than
         // we actually need. The requirement is that we can't create another

lib/SILOptimizer/Transforms/ArrayElementValuePropagation.cpp

@@ -67,7 +67,7 @@ class ArrayAllocation {
   /// A map of Array indices to element values
   llvm::DenseMap<uint64_t, SILValue> ElementValueMap;
 
-  bool mapInitializationStores(SILValue ElementBuffer);
+  bool mapInitializationStores(ArraySemanticsCall arrayUninitCall);
   bool recursivelyCollectUses(ValueBase *Def);
   bool replacementsAreValid();
 
@@ -93,64 +93,16 @@ class ArrayAllocation {
 };
 
 /// Map the indices of array element initialization stores to their values.
-bool ArrayAllocation::mapInitializationStores(SILValue ElementBuffer) {
-  assert(ElementBuffer &&
-         "Must have identified an array element storage pointer");
-
-  // Match initialization stores.
-  // %83 = struct_extract %element_buffer : $UnsafeMutablePointer<Int>
-  // %84 = pointer_to_address %83 : $Builtin.RawPointer to strict $*Int
-  // store %85 to %84 : $*Int
-  // %87 = integer_literal $Builtin.Word, 1
-  // %88 = index_addr %84 : $*Int, %87 : $Builtin.Word
-  // store %some_value to %88 : $*Int
-
-  auto *UnsafeMutablePointerExtract =
-      dyn_cast_or_null<StructExtractInst>(getSingleNonDebugUser(ElementBuffer));
-  if (!UnsafeMutablePointerExtract)
+bool ArrayAllocation::mapInitializationStores(
+    ArraySemanticsCall arrayUninitCall) {
+  llvm::DenseMap<uint64_t, StoreInst *> elementStoreMap;
+  if (!arrayUninitCall.mapInitializationStores(elementStoreMap))
     return false;
-  auto *PointerToAddress = dyn_cast_or_null<PointerToAddressInst>(
-      getSingleNonDebugUser(UnsafeMutablePointerExtract));
-  if (!PointerToAddress)
-    return false;
-
-  // Match the stores. We can have either a store directly to the address or
-  // to an index_addr projection.
-  for (auto *Op : PointerToAddress->getUses()) {
-    auto *Inst = Op->getUser();
-
-    // Store to the base.
-    auto *SI = dyn_cast<StoreInst>(Inst);
-    if (SI && SI->getDest() == PointerToAddress) {
-      // We have already seen an entry for this index bail.
-      if (ElementValueMap.count(0))
-        return false;
-      ElementValueMap[0] = SI->getSrc();
-      continue;
-    } else if (SI)
-      return false;
-
-    // Store an index_addr projection.
-    auto *IndexAddr = dyn_cast<IndexAddrInst>(Inst);
-    if (!IndexAddr)
-      return false;
-    SI = dyn_cast_or_null<StoreInst>(getSingleNonDebugUser(IndexAddr));
-    if (!SI || SI->getDest() != IndexAddr)
-      return false;
-    auto *Index = dyn_cast<IntegerLiteralInst>(IndexAddr->getIndex());
-    if (!Index)
-      return false;
-    auto IndexVal = Index->getValue();
-    // Let's not blow up our map.
-    if (IndexVal.getActiveBits() > 16)
-      return false;
-    // Already saw an entry.
-    if (ElementValueMap.count(IndexVal.getZExtValue()))
-      return false;
-
-    ElementValueMap[IndexVal.getZExtValue()] = SI->getSrc();
-  }
-  return !ElementValueMap.empty();
+  // Extract the SIL values of the array elements from the stores.
+  ElementValueMap.grow(elementStoreMap.size());
+  for (auto keyValue : elementStoreMap)
+    ElementValueMap[keyValue.getFirst()] = keyValue.getSecond()->getSrc();
+  return true;
 }
 
 bool ArrayAllocation::replacementsAreValid() {
@@ -217,12 +169,8 @@ bool ArrayAllocation::analyze(ApplyInst *Alloc) {
   if (!ArrayValue)
     return false;
 
-  SILValue ElementBuffer = Uninitialized.getArrayElementStoragePointer();
-  if (!ElementBuffer)
-    return false;
-
   // Figure out all stores to the array.
-  if (!mapInitializationStores(ElementBuffer))
+  if (!mapInitializationStores(Uninitialized))
     return false;
 
   // Check if the array value was stored or has escaped.

lib/SILOptimizer/Transforms/DeadObjectElimination.cpp

@@ -641,7 +641,8 @@ static bool removeAndReleaseArray(SingleValueInstruction *NewArrayValue,
 /// side effect?
 static bool isAllocatingApply(SILInstruction *Inst) {
   ArraySemanticsCall ArrayAlloc(Inst);
-  return ArrayAlloc.getKind() == ArrayCallKind::kArrayUninitialized;
+  return ArrayAlloc.getKind() == ArrayCallKind::kArrayUninitialized ||
+         ArrayAlloc.getKind() == ArrayCallKind::kArrayUninitializedIntrinsic;
 }
 
 namespace {
@@ -832,7 +833,9 @@ static bool getDeadInstsAfterInitializerRemoved(
 bool DeadObjectElimination::processAllocApply(ApplyInst *AI,
                                               DeadEndBlocks &DEBlocks) {
   // Currently only handle array.uninitialized
-  if (ArraySemanticsCall(AI).getKind() != ArrayCallKind::kArrayUninitialized)
+  if (ArraySemanticsCall(AI).getKind() != ArrayCallKind::kArrayUninitialized &&
+      ArraySemanticsCall(AI).getKind() !=
+          ArrayCallKind::kArrayUninitializedIntrinsic)
     return false;
 
   llvm::SmallVector<SILInstruction *, 8> instsDeadAfterInitializerRemoved;

test/SILOptimizer/dead_array_elim.sil

@@ -21,7 +21,7 @@ class TrivialDestructor {
 // Remove a dead array.
 // rdar://20980377 Add dead array elimination to DeadObjectElimination
 // Swift._allocateUninitializedArray <A> (Builtin.Word) -> (Swift.Array<A>, Builtin.RawPointer)
-sil [_semantics "array.uninitialized"] @allocArray : $@convention(thin) <τ_0_0> (Builtin.Word) -> @owned (Array<τ_0_0>, Builtin.RawPointer)
+sil [_semantics "array.uninitialized_intrinsic"] @allocArray : $@convention(thin) <τ_0_0> (Builtin.Word) -> @owned (Array<τ_0_0>, Builtin.RawPointer)
 
 sil [_semantics "array.uninitialized"] @adoptStorageSpecialiedForInt : $@convention(method) (@guaranteed _ContiguousArrayStorage<Int>, Builtin.Word, @thin Array<Int>.Type) -> (@owned Array<Int>, UnsafeMutablePointer<Int>)