Handle mark_dependence in Differentiation

Author: meg-gupta

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/Transforms/ArrayElementValuePropagation.cpp

@@ -151,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()) {