Merge pull request #69450 from meg-gupta/uninarrayfix

Add a mark_dependence while emitting SIL for uninitialized array allocation

Author: meg-gupta

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ObjectOutliner.swift

@@ -41,7 +41,6 @@ import SIL
 ///
 let objectOutliner = FunctionPass(name: "object-outliner") {
   (function: Function, context: FunctionPassContext) in
-
   for inst in function.instructions {
     if let ari = inst as? AllocRefInstBase {
       if let globalValue = optimizeObjectAllocation(allocRef: ari, context) {
@@ -85,6 +84,10 @@ private func findEndCOWMutation(of object: Value) -> EndCOWMutationInst? {
       if let ecm = findEndCOWMutation(of: uci) {
         return ecm
       }
+    case let urci as UncheckedRefCastInst:
+      if let ecm = findEndCOWMutation(of: urci) {
+        return ecm
+      }
     case let mv as MoveValueInst:
       if let ecm = findEndCOWMutation(of: mv) {
         return ecm
@@ -147,7 +150,7 @@ private func findInitStores(of object: Value,
         return false
       }
     default:
-      if !isValidUseOfObject(use.instruction) {
+      if !isValidUseOfObject(use) {
         return false
       }
     }
@@ -174,6 +177,18 @@ private func findStores(toTailAddress tailAddr: Value, tailElementIndex: Int, st
       if !findStores(inUsesOf: tea, index: tailElementIndex * numTupleElements + tupleIdx, stores: &stores) {
         return false
       }
+    case let atp as AddressToPointerInst:
+      if !findStores(toTailAddress: atp, tailElementIndex: tailElementIndex, stores: &stores) {
+        return false
+      }
+    case let mdi as MarkDependenceInst:
+      if !findStores(toTailAddress: mdi, tailElementIndex: tailElementIndex, stores: &stores) {
+        return false
+      }
+    case let pta as PointerToAddressInst:
+      if !findStores(toTailAddress: pta, tailElementIndex: tailElementIndex, stores: &stores) {
+        return false
+      }
     case let store as StoreInst:
       if store.source.type.isTuple {
         // This kind of SIL is never generated because tuples are stored with separated stores to tuple_element_addr.
@@ -184,7 +199,7 @@ private func findStores(toTailAddress tailAddr: Value, tailElementIndex: Int, st
         return false
       }
     default:
-      if !isValidUseOfObject(use.instruction) {
+      if !isValidUseOfObject(use) {
         return false
       }
     }
@@ -198,7 +213,7 @@ private func findStores(inUsesOf address: Value, index: Int, stores: inout [Stor
       if !handleStore(store, index: index, stores: &stores) {
         return false
       }
-    } else if !isValidUseOfObject(use.instruction) {
+    } else if !isValidUseOfObject(use) {
       return false
     }
   }
@@ -215,7 +230,8 @@ private func handleStore(_ store: StoreInst, index: Int, stores: inout [StoreIns
   return false
 }
 
-private func isValidUseOfObject(_ inst: Instruction) -> Bool {
+private func isValidUseOfObject(_ use: Operand) -> Bool {
+  let inst = use.instruction
   switch inst {
   case is DebugValueInst,
        is LoadInst,
@@ -227,6 +243,17 @@ private func isValidUseOfObject(_ inst: Instruction) -> Bool {
        is EndCOWMutationInst:
     return true
 
+  case let mdi as MarkDependenceInst:
+    if (use == mdi.baseOperand) {
+      return true;
+    }
+    for mdiUse in mdi.uses {
+      if !isValidUseOfObject(mdiUse) {
+        return false
+      }
+    }
+    return true
+
   case is StructElementAddrInst,
        is AddressToPointerInst,
        is StructInst,
@@ -238,9 +265,12 @@ private func isValidUseOfObject(_ inst: Instruction) -> Bool {
        is UpcastInst,
        is BeginDeallocRefInst,
        is RefTailAddrInst,
-       is RefElementAddrInst:
-    for use in (inst as! SingleValueInstruction).uses {
-      if !isValidUseOfObject(use.instruction) {
+       is RefElementAddrInst,
+       is StructInst,
+       is PointerToAddressInst,
+       is IndexAddrInst:
+    for instUse in (inst as! SingleValueInstruction).uses {
+      if !isValidUseOfObject(instUse) {
         return false
       }
     }
@@ -342,6 +372,8 @@ private func rewriteUses(of startValue: Value, _ context: FunctionPassContext) {
       context.erase(instruction: endMutation)
     case let upCast as UpcastInst:
       worklist.pushIfNotVisited(usersOf: upCast)
+    case let urci as UncheckedRefCastInst:
+      worklist.pushIfNotVisited(usersOf: urci)
     case let moveValue as MoveValueInst:
       worklist.pushIfNotVisited(usersOf: moveValue)
     case is DeallocRefInst, is DeallocStackRefInst:

lib/SILGen/SILGenApply.cpp

@@ -6252,7 +6252,10 @@ SILGenFunction::emitUninitializedArrayAllocation(Type ArrayTy,
   SmallVector<ManagedValue, 2> resultElts;
   std::move(result).getAll(resultElts);
 
-  return {resultElts[0], resultElts[1].getUnmanagedValue()};
+  // Add a mark_dependence between the interior pointer and the array value
+  auto dependentValue = B.createMarkDependence(Loc, resultElts[1].getValue(),
+                                               resultElts[0].getValue());
+  return {resultElts[0], dependentValue};
 }
 
 /// Deallocate an uninitialized array.

lib/SILOptimizer/Analysis/ArraySemantic.cpp

@@ -841,7 +841,8 @@ bool swift::ArraySemanticsCall::mapInitializationStores(
     return false;
 
   // Match initialization stores into ElementBuffer. E.g.
-  // %83 = struct_extract %element_buffer : $UnsafeMutablePointer<Int>
+  // %82 = struct_extract %element_buffer : $UnsafeMutablePointer<Int>
+  // %83 = mark_dependence %82 : $Builtin.RawPointer on ArrayVal
   // %84 = pointer_to_address %83 : $Builtin.RawPointer to strict $*Int
   // store %85 to %84 : $*Int
   // %87 = integer_literal $Builtin.Word, 1
@@ -850,12 +851,27 @@ bool swift::ArraySemanticsCall::mapInitializationStores(
 
   // If this an ArrayUninitializedIntrinsic 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;
+  // be struct-extracted to obtain a builtin.RawPointer. In this case
+  // mark_dependence can be an operand of the struct_extract or its user.
+
+  SILValue UnsafeMutablePointerExtract;
+  if (getKind() == ArrayCallKind::kArrayUninitializedIntrinsic) {
+    UnsafeMutablePointerExtract = dyn_cast_or_null<MarkDependenceInst>(
+        getSingleNonDebugUser(ElementBuffer));
+  } else {
+    auto user = getSingleNonDebugUser(ElementBuffer);
+    // Match mark_dependence (struct_extract or
+    // struct_extract (mark_dependence
+    if (auto *MDI = dyn_cast_or_null<MarkDependenceInst>(user)) {
+      UnsafeMutablePointerExtract =
+          dyn_cast_or_null<StructExtractInst>(getSingleNonDebugUser(MDI));
+    } else {
+      if (auto *SEI = dyn_cast_or_null<StructExtractInst>(user)) {
+        UnsafeMutablePointerExtract =
+            dyn_cast_or_null<MarkDependenceInst>(getSingleNonDebugUser(SEI));
+      }
+    }
+  }
   if (!UnsafeMutablePointerExtract)
     return false;
 

lib/SILOptimizer/Analysis/DifferentiableActivityAnalysis.cpp

@@ -429,12 +429,15 @@ void DifferentiableActivityInfo::setUsefulThroughArrayInitialization(
       continue;
     // The second tuple field of the return value is the `RawPointer`.
     for (auto use : dti->getResult(1)->getUses()) {
-      // The `RawPointer` passes through a `pointer_to_address`. That
-      // instruction's first use is a `store` whose source is useful; its
+      // The `RawPointer` passes through a `mark_dependence(pointer_to_address`.
+      // That instruction's first use is a `store` whose source is useful; its
       // subsequent uses are `index_addr`s whose only use is a useful `store`.
-      auto *ptai = dyn_cast<PointerToAddressInst>(use->getUser());
-      assert(ptai && "Expected `pointer_to_address` user for uninitialized "
-                     "array intrinsic");
+      auto *mdi = dyn_cast<MarkDependenceInst>(use->getUser());
+      assert(
+          mdi &&
+          "Expected a mark_dependence user for uninitialized array intrinsic.");
+      auto *ptai = dyn_cast<PointerToAddressInst>(getSingleNonDebugUser(mdi));
+      assert(ptai && "Expected a pointer_to_address.");
       setUseful(ptai, dependentVariableIndex);
       // Propagate usefulness through array element addresses:
       // `pointer_to_address` and `index_addr` instructions.

lib/SILOptimizer/Differentiation/Common.cpp

@@ -37,30 +37,18 @@ ApplyInst *getAllocateUninitializedArrayIntrinsicElementAddress(SILValue v) {
     ptai = dyn_cast<PointerToAddressInst>(iai->getOperand(0));
   if (!ptai)
     return nullptr;
+  auto *mdi = dyn_cast<MarkDependenceInst>(
+      ptai->getOperand()->getDefiningInstruction());
+  if (!mdi)
+    return nullptr;
   // Return the `array.uninitialized_intrinsic` application, if it exists.
   if (auto *dti = dyn_cast<DestructureTupleInst>(
-          ptai->getOperand()->getDefiningInstruction()))
+          mdi->getValue()->getDefiningInstruction()))
     return ArraySemanticsCall(dti->getOperand(),
                               semantics::ARRAY_UNINITIALIZED_INTRINSIC);
   return nullptr;
 }
 
-DestructureTupleInst *getSingleDestructureTupleUser(SILValue value) {
-  bool foundDestructureTupleUser = false;
-  if (!value->getType().is<TupleType>())
-    return nullptr;
-  DestructureTupleInst *result = nullptr;
-  for (auto *use : value->getUses()) {
-    if (auto *dti = dyn_cast<DestructureTupleInst>(use->getUser())) {
-      assert(!foundDestructureTupleUser &&
-             "There should only be one `destructure_tuple` user of a tuple");
-      foundDestructureTupleUser = true;
-      result = dti;
-    }
-  }
-  return result;
-}
-
 bool isSemanticMemberAccessor(SILFunction *original) {
   auto *dc = original->getDeclContext();
   if (!dc)
@@ -109,7 +97,7 @@ void forEachApplyDirectResult(
       resultCallback(ai);
       return;
     }
-    if (auto *dti = getSingleDestructureTupleUser(ai))
+    if (auto *dti = ai->getSingleUserOfType<DestructureTupleInst>())
       for (auto directResult : dti->getResults())
         resultCallback(directResult);
     break;

lib/SILOptimizer/Differentiation/JVPCloner.cpp

@@ -1312,7 +1312,8 @@ class JVPCloner::Implementation final
         if (!origResult->getType().is<TupleType>()) {
           setTangentValue(bb, origResult,
                           makeConcreteTangentValue(differentialResult));
-        } else if (auto *dti = getSingleDestructureTupleUser(ai)) {
+        } else if (auto *dti =
+                       ai->getSingleUserOfType<DestructureTupleInst>()) {
           bool notSetValue = true;
           for (auto result : dti->getResults()) {
             if (activityInfo.isActive(result, getConfig())) {

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -3331,7 +3331,11 @@ void PullbackCloner::Implementation::
   builder.setCurrentDebugScope(remapScope(dti->getDebugScope()));
   builder.setInsertionPoint(arrayAdjoint->getParentBlock());
   for (auto use : dti->getResult(1)->getUses()) {
-    auto *ptai = dyn_cast<PointerToAddressInst>(use->getUser());
+    auto *mdi = dyn_cast<MarkDependenceInst>(use->getUser());
+    assert(mdi && "Expected mark_dependence user");
+    auto *ptai =
+        dyn_cast_or_null<PointerToAddressInst>(getSingleNonDebugUser(mdi));
+    assert(ptai && "Expected pointer_to_address user");
     auto adjBuf = getAdjointBuffer(origBB, ptai);
     auto *eltAdjBuf = getArrayAdjointElementBuffer(arrayAdjoint, 0, loc);
     builder.emitInPlaceAdd(loc, adjBuf, eltAdjBuf);

lib/SILOptimizer/LoopTransforms/ForEachLoopUnroll.cpp

@@ -302,6 +302,11 @@ void ArrayInfo::classifyUsesOfArray(SILValue arrayValue) {
     // above as the array would be passed indirectly.
     if (isFixLifetimeUseOfArray(user, arrayValue))
       continue;
+    if (auto *MDI = dyn_cast<MarkDependenceInst>(user)) {
+      if (MDI->getBase() == arrayValue) {
+        continue;
+      }
+    }
     // Check if this is a forEach call on the array.
     if (TryApplyInst *forEachCall = isForEachUseOfArray(user, arrayValue)) {
       forEachCalls.insert(forEachCall);

lib/SILOptimizer/Transforms/ArrayCountPropagation.cpp

@@ -127,6 +127,12 @@ bool ArrayAllocation::recursivelyCollectUses(ValueBase *Def) {
         isa<DebugValueInst>(User))
       continue;
 
+    if (auto *MDI = dyn_cast<MarkDependenceInst>(User)) {
+      if (Def == MDI->getBase()) {
+        continue;
+      }
+    }
+
     // Array value projection.
     if (auto *SEI = dyn_cast<StructExtractInst>(User)) {
       if (!recursivelyCollectUses(SEI))

lib/SILOptimizer/Transforms/ArrayElementValuePropagation.cpp

@@ -122,6 +122,9 @@ bool ArrayAllocation::replacementsAreValid() {
 /// Recursively look at all uses of this definition. Abort if the array value
 /// could escape or be changed. Collect all uses that are calls to array.count.
 bool ArrayAllocation::recursivelyCollectUses(ValueBase *Def) {
+  LLVM_DEBUG(llvm::dbgs() << "Collecting uses of:");
+  LLVM_DEBUG(Def->dump());
+
   for (auto *Opd : Def->getUses()) {
     auto *User = Opd->getUser();
     // Ignore reference counting and debug instructions.
@@ -148,6 +151,12 @@ bool ArrayAllocation::recursivelyCollectUses(ValueBase *Def) {
       continue;
     }
 
+    if (auto *MDI = dyn_cast<MarkDependenceInst>(User)) {
+      if (Def != MDI->getBase())
+        return false;
+      continue;
+    }
+
     // Check array semantic calls.
     ArraySemanticsCall ArrayOp(User);
     switch (ArrayOp.getKind()) {
@@ -181,17 +190,28 @@ bool ArrayAllocation::analyze(ApplyInst *Alloc) {
   if (!Uninitialized)
     return false;
 
+  LLVM_DEBUG(llvm::dbgs() << "Found array allocation: ");
+  LLVM_DEBUG(Alloc->dump());
+
   ArrayValue = Uninitialized.getArrayValue();
-  if (!ArrayValue)
+  if (!ArrayValue) {
+    LLVM_DEBUG(llvm::dbgs() << "Did not find array value\n");
     return false;
+  }
 
+  LLVM_DEBUG(llvm::dbgs() << "ArrayValue: ");
+  LLVM_DEBUG(ArrayValue->dump());
   // Figure out all stores to the array.
-  if (!mapInitializationStores(Uninitialized))
+  if (!mapInitializationStores(Uninitialized)) {
+    LLVM_DEBUG(llvm::dbgs() << "Could not map initializing stores\n");
     return false;
+  }
 
   // Check if the array value was stored or has escaped.
-  if (!recursivelyCollectUses(ArrayValue))
+  if (!recursivelyCollectUses(ArrayValue)) {
+    LLVM_DEBUG(llvm::dbgs() << "Array value stored or escaped\n");
     return false;
+  }
 
   return true;
 }
@@ -328,7 +348,9 @@ class ArrayElementPropagation : public SILFunctionTransform {
     auto &Fn = *getFunction();
     bool Changed = false;
 
-    for (auto &BB :Fn) {
+    LLVM_DEBUG(llvm::dbgs() << "ArrayElementPropagation looking at function: "
+                            << Fn.getName() << "\n");
+    for (auto &BB : Fn) {
       for (auto &Inst : BB) {
         if (auto *Apply = dyn_cast<ApplyInst>(&Inst)) {
           ArrayAllocation ALit;