Make autdiff more robust in presence of unreachable blocks (#71356)

Unreachable blocks possess some challenges to autodiff since in reverse pass (pullback generation) we need to execute the function backwards, pushing the values from return BB back to entry block. As a result, unreachable blocks might become reachable from the return BB and this might cause all kind of issues.

Author: asl

lib/SILOptimizer/Differentiation/LinearMapInfo.cpp

@@ -143,7 +143,9 @@ void LinearMapInfo::populateBranchingTraceDecl(SILBasicBlock *originalBB,
       decl->setInterfaceType(astCtx.TheRawPointerType);
     } else { // Otherwise the payload is the linear map tuple.
       auto *linearMapStructTy = getLinearMapTupleType(predBB);
-      assert(linearMapStructTy && "must have linear map struct type for predecessor BB");
+      // Do not create entries for unreachable predecessors
+      if (!linearMapStructTy)
+        continue;
       auto canLinearMapStructTy = linearMapStructTy->getCanonicalType();
       decl->setInterfaceType(
           canLinearMapStructTy->hasArchetype()

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -2087,7 +2087,10 @@ bool PullbackCloner::Implementation::run() {
       originalBlocks.push_back(BB);
 
       for (auto *nextBB : BB->getPredecessorBlocks()) {
-        workqueue.pushIfNotVisited(nextBB);
+        // If there is no linear map tuple for predecessor BB, then BB is
+        // unreachable from function entry. Do not run pullback cloner on it.
+        if (getPullbackInfo().getLinearMapTupleType(nextBB))
+          workqueue.pushIfNotVisited(nextBB);
       }
     }
   }
@@ -2803,6 +2806,11 @@ void PullbackCloner::Implementation::visitSILBasicBlock(SILBasicBlock *bb) {
   SmallDenseMap<SILValue, TrampolineBlockSet> pullbackTrampolineBlockMap;
   SmallDenseMap<SILBasicBlock *, SILBasicBlock *> origPredpullbackSuccBBMap;
   for (auto *predBB : bb->getPredecessorBlocks()) {
+    // If there is no linear map tuple for predecessor BB, then BB is
+    // unreachable from function entry. There is no branch tracing enum for it
+    // as well, so we should not create any branching to it in the pullback.
+    if (!getPullbackInfo().getLinearMapTupleType(predBB))
+      continue;
     auto *pullbackSuccBB =
         buildPullbackSuccessor(bb, predBB, pullbackTrampolineBlockMap);
     origPredpullbackSuccBBMap[predBB] = pullbackSuccBB;

test/AutoDiff/compiler_crashers_fixed/issue-71164-unreachable-blocks.swift

@@ -0,0 +1,44 @@
+// RUN: %target-swift-frontend -emit-sil -verify %s
+// REQUIRES: swift_in_compiler
+
+// https://github.com/apple/swift/issues/71164
+
+// There are few unreachble blocks created due to mandatory boolean constant
+// propagation
+// Ensure we do not ceeate linear map types for unreachable BB and that they
+// are skipped during VJP and pullback generation
+import _Differentiation
+
+struct A: Differentiable {}
+
+struct B: Differentiable {
+    @differentiable(reverse)
+    func c(b: B) -> A {
+        while true {
+            if true {
+                break
+            }
+        };
+        
+        return A()
+    }
+
+    @differentiable(reverse)
+    func d(b: B) -> A {
+        while true {
+            if true {
+                return c(b : b)
+            }
+            if true {
+                break
+            }
+            if false {
+                return c(b : b)
+            } else {
+                break
+            }
+        };
+        return A()
+    }
+
+}