[AutoDiff] Make force-unwrapping differentiable. (#26826)

`Optional` force-unwrapping is a mathematically transposable operation. This patch makes force-unwrapping differentiable by applying its transpose on adjoint buffers.

```swift
func bla<T: Differentiable & FloatingPoint>(_ t: T) -> (T, Float) where T == T.TangentVector {
    gradient(at: t, Float(1)) { (x, y) in (x as! Float) * y }
}
print(bla(Float(2))) // (1, 2)
```

Resolves [TF-455](https://bugs.swift.org/browse/TF-455).

Author: rxwei

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -1820,6 +1820,12 @@ void DifferentiableActivityInfo::analyze(DominanceInfo *di,
           if (isVaried(cai->getSrc(), i))
             recursivelySetVaried(cai->getDest(), i);
         }
+        // Handle `unconditional_checked_cast_addr`.
+        else if (auto *uccai =
+                     dyn_cast<UnconditionalCheckedCastAddrInst>(&inst)) {
+          if (isVaried(uccai->getSrc(), i))
+            recursivelySetVaried(uccai->getDest(), i);
+        }
         // Handle `tuple_element_addr`.
         else if (auto *teai = dyn_cast<TupleElementAddrInst>(&inst)) {
           if (isVaried(teai->getOperand(), i)) {
@@ -1941,6 +1947,12 @@ void DifferentiableActivityInfo::analyze(DominanceInfo *di,
           if (isUseful(cai->getDest(), i))
             propagateUsefulThroughBuffer(cai->getSrc(), i);
         }
+        // Handle `unconditional_checked_cast_addr`.
+        else if (auto *uccai =
+                     dyn_cast<UnconditionalCheckedCastAddrInst>(&inst)) {
+          if (isUseful(uccai->getDest(), i))
+            propagateUsefulThroughBuffer(uccai->getSrc(), i);
+        }
         // Handle everything else.
         else if (llvm::any_of(inst.getResults(),
           [&](SILValue res) { return isUseful(res, i); })) {
@@ -6601,6 +6613,27 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
     }
   }
 
+  /// Handle `unconditional_checked_cast_addr` instruction.
+  ///   Original: y = unconditional_checked_cast_addr x
+  ///    Adjoint: adj[x] += unconditional_checked_cast_addr adj[y]
+  void visitUnconditionalCheckedCastAddrInst(
+      UnconditionalCheckedCastAddrInst *uccai) {
+    auto *bb = uccai->getParent();
+    auto &adjDest = getAdjointBuffer(bb, uccai->getDest());
+    auto &adjSrc = getAdjointBuffer(bb, uccai->getSrc());
+    if (errorOccurred)
+      return;
+    auto destType = remapType(adjDest->getType());
+    auto castBuf = builder.createAllocStack(uccai->getLoc(), adjSrc->getType());
+    builder.createUnconditionalCheckedCastAddr(
+        uccai->getLoc(), adjDest, adjDest->getType().getASTType(), castBuf,
+        adjSrc->getType().getASTType());
+    addToAdjointBuffer(bb, uccai->getSrc(), castBuf, uccai->getLoc());
+    builder.emitDestroyAddrAndFold(uccai->getLoc(), castBuf);
+    builder.createDeallocStack(uccai->getLoc(), castBuf);
+    emitZeroIndirect(destType.getASTType(), adjDest, uccai->getLoc());
+  }
+
 #define NOT_DIFFERENTIABLE(INST, DIAG) \
   void visit##INST##Inst(INST##Inst *inst) { \
     getContext().emitNondifferentiabilityError( \

test/AutoDiff/simple_math.swift

@@ -315,4 +315,13 @@ SimpleMathTests.test("SubsetIndices") {
   expectEqual(4, gradWRTNonDiff { x, y in x + Float(y) })
 }
 
+SimpleMathTests.test("ForceUnwrapping") {
+  func bla<T: Differentiable & FloatingPoint>(_ t: T) -> (T, Float) where T == T.TangentVector {
+    gradient(at: t, Float(1)) { (x, y) in
+      (x as! Float) * y
+    }
+  }
+  expectEqual((1, 2), bla(Float(2)))
+}
+
 runAllTests()