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

semantics attribute that is used by the top-level array initializer (in ArrayShared.swift),
which is the entry point used by the compiler to initialize array from array literals.
This initializer is early-inlined so that other optimizations can work on its body.

Fix DeadObjectElimination and ArrayCOWOpts optimization passes to work with this
semantics attribute in addition to "array.uninitialized", which they already use.

Refactor mapInitializationStores function from ArrayElementValuePropagation.cpp to
ArraySemantic.cpp so that the array-initialization pattern matching functionality
implemented by the function can be reused by other optimizations.

Author: ravikandhadai

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>)