[AutoDiff] Ensure autodiff code does not ignore getSingleTerminatorOperands return value (#64200)

Also extend activity analysis to handle `try_apply` normal result properly. Add testcase from #63728

Author: asl

include/swift/SIL/SILArgument.h

@@ -178,7 +178,7 @@ class SILArgument : public ValueBase {
   /// Note: this peeks through any projections or cast implied by the
   /// terminator. e.g. the incoming value for a switch_enum payload argument is
   /// the enum itself (the operand of the switch_enum).
-  bool getSingleTerminatorOperands(
+  [[nodiscard]] bool getSingleTerminatorOperands(
       SmallVectorImpl<SILValue> &returnedSingleTermOperands) const;
 
   /// Returns true if we were able to find single terminator operand values for
@@ -188,7 +188,7 @@ class SILArgument : public ValueBase {
   /// Note: this peeks through any projections or cast implied by the
   /// terminator. e.g. the incoming value for a switch_enum payload argument is
   /// the enum itself (the operand of the switch_enum).
-  bool getSingleTerminatorOperands(
+  [[nodiscard]] bool getSingleTerminatorOperands(
       SmallVectorImpl<std::pair<SILBasicBlock *, SILValue>>
           &returnedSingleTermOperands) const;
 
@@ -303,7 +303,7 @@ class SILPhiArgument : public SILArgument {
   /// Note: this peeks through any projections or cast implied by the
   /// terminator. e.g. the incoming value for a switch_enum payload argument is
   /// the enum itself (the operand of the switch_enum).
-  bool getSingleTerminatorOperands(
+  [[nodiscard]] bool getSingleTerminatorOperands(
       SmallVectorImpl<SILValue> &returnedSingleTermOperands) const;
 
   /// Returns true if we were able to find single terminator operand values for
@@ -313,7 +313,7 @@ class SILPhiArgument : public SILArgument {
   /// Note: this peeks through any projections or cast implied by the
   /// terminator. e.g. the incoming value for a switch_enum payload argument is
   /// the enum itself (the operand of the switch_enum).
-  bool getSingleTerminatorOperands(
+  [[nodiscard]] bool getSingleTerminatorOperands(
       SmallVectorImpl<std::pair<SILBasicBlock *, SILValue>>
           &returnedSingleTermOperands) const;
 

lib/SILOptimizer/Analysis/DifferentiableActivityAnalysis.cpp

@@ -303,15 +303,25 @@ void DifferentiableActivityInfo::setUsefulAndPropagateToOperands(
     return;
   }
   setUseful(value, dependentVariableIndex);
+
   // If the given value is a basic block argument, propagate usefulness to
   // incoming values.
   if (auto *bbArg = dyn_cast<SILPhiArgument>(value)) {
     SmallVector<SILValue, 4> incomingValues;
-    bbArg->getSingleTerminatorOperands(incomingValues);
-    for (auto incomingValue : incomingValues)
-      setUsefulAndPropagateToOperands(incomingValue, dependentVariableIndex);
-    return;
+    if (bbArg->getSingleTerminatorOperands(incomingValues)) {
+      for (auto incomingValue : incomingValues)
+        setUsefulAndPropagateToOperands(incomingValue, dependentVariableIndex);
+      return;
+    } else if (bbArg->isTerminatorResult()) {
+      if (TryApplyInst *tai = dyn_cast<TryApplyInst>(bbArg->getTerminatorForResult())) {
+        propagateUseful(tai, dependentVariableIndex);
+        return;
+      } else
+        llvm::report_fatal_error("unknown terminator with result");
+   } else
+      llvm::report_fatal_error("do not know how to handle this incoming bb argument");
   }
+  
   auto *inst = value->getDefiningInstruction();
   if (!inst)
     return;

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -2515,12 +2515,11 @@ void PullbackCloner::Implementation::visitSILBasicBlock(SILBasicBlock *bb) {
 
     // Get predecessor terminator operands.
     SmallVector<std::pair<SILBasicBlock *, SILValue>, 4> incomingValues;
-    bbArg->getSingleTerminatorOperands(incomingValues);
-
-    // Returns true if the given terminator instruction is a `switch_enum` on
-    // an `Optional`-typed value. `switch_enum` instructions require
-    // special-case adjoint value propagation for the operand.
-    auto isSwitchEnumInstOnOptional =
+    if (bbArg->getSingleTerminatorOperands(incomingValues)) {
+      // Returns true if the given terminator instruction is a `switch_enum` on
+      // an `Optional`-typed value. `switch_enum` instructions require
+      // special-case adjoint value propagation for the operand.
+      auto isSwitchEnumInstOnOptional =
         [&ctx = getASTContext()](TermInst *termInst) {
           if (!termInst)
             return false;
@@ -2531,49 +2530,52 @@ void PullbackCloner::Implementation::visitSILBasicBlock(SILBasicBlock *bb) {
           return false;
         };
 
-    // Check the tangent value category of the active basic block argument.
-    switch (getTangentValueCategory(bbArg)) {
-    // If argument has a loadable tangent value category: materialize adjoint
-    // value of the argument, create a copy, and set the copy as the adjoint
-    // value of incoming values.
-    case SILValueCategory::Object: {
-      auto bbArgAdj = getAdjointValue(bb, bbArg);
-      auto concreteBBArgAdj = materializeAdjointDirect(bbArgAdj, pbLoc);
-      auto concreteBBArgAdjCopy =
+      // Check the tangent value category of the active basic block argument.
+      switch (getTangentValueCategory(bbArg)) {
+        // If argument has a loadable tangent value category: materialize adjoint
+        // value of the argument, create a copy, and set the copy as the adjoint
+        // value of incoming values.
+      case SILValueCategory::Object: {
+        auto bbArgAdj = getAdjointValue(bb, bbArg);
+        auto concreteBBArgAdj = materializeAdjointDirect(bbArgAdj, pbLoc);
+        auto concreteBBArgAdjCopy =
           builder.emitCopyValueOperation(pbLoc, concreteBBArgAdj);
-      for (auto pair : incomingValues) {
-        auto *predBB = std::get<0>(pair);
-        auto incomingValue = std::get<1>(pair);
-        // Handle `switch_enum` on `Optional`.
-        auto termInst = bbArg->getSingleTerminator();
-        if (isSwitchEnumInstOnOptional(termInst)) {
-          accumulateAdjointForOptional(bb, incomingValue, concreteBBArgAdjCopy);
-        } else {
-          blockTemporaries[getPullbackBlock(predBB)].insert(
+        for (auto pair : incomingValues) {
+          pair.second->dump();
+          auto *predBB = std::get<0>(pair);
+          auto incomingValue = std::get<1>(pair);
+          // Handle `switch_enum` on `Optional`.
+          auto termInst = bbArg->getSingleTerminator();
+          if (isSwitchEnumInstOnOptional(termInst)) {
+            accumulateAdjointForOptional(bb, incomingValue, concreteBBArgAdjCopy);
+          } else {
+            blockTemporaries[getPullbackBlock(predBB)].insert(
               concreteBBArgAdjCopy);
-          setAdjointValue(predBB, incomingValue,
-                          makeConcreteAdjointValue(concreteBBArgAdjCopy));
+            setAdjointValue(predBB, incomingValue,
+                            makeConcreteAdjointValue(concreteBBArgAdjCopy));
+          }
         }
+        break;
       }
-      break;
-    }
-    // If argument has an address tangent value category: materialize adjoint
-    // value of the argument, create a copy, and set the copy as the adjoint
-    // value of incoming values.
-    case SILValueCategory::Address: {
-      auto bbArgAdjBuf = getAdjointBuffer(bb, bbArg);
-      for (auto pair : incomingValues) {
-        auto incomingValue = std::get<1>(pair);
-        // Handle `switch_enum` on `Optional`.
-        auto termInst = bbArg->getSingleTerminator();
-        if (isSwitchEnumInstOnOptional(termInst))
-          accumulateAdjointForOptional(bb, incomingValue, bbArgAdjBuf);
-        else
-          addToAdjointBuffer(bb, incomingValue, bbArgAdjBuf, pbLoc);
+      // If argument has an address tangent value category: materialize adjoint
+      // value of the argument, create a copy, and set the copy as the adjoint
+      // value of incoming values.
+      case SILValueCategory::Address: {
+        auto bbArgAdjBuf = getAdjointBuffer(bb, bbArg);
+        for (auto pair : incomingValues) {
+          auto incomingValue = std::get<1>(pair);
+          // Handle `switch_enum` on `Optional`.
+          auto termInst = bbArg->getSingleTerminator();
+          if (isSwitchEnumInstOnOptional(termInst))
+            accumulateAdjointForOptional(bb, incomingValue, bbArgAdjBuf);
+          else
+            addToAdjointBuffer(bb, incomingValue, bbArgAdjBuf, pbLoc);
+        }
+        break;
       }
-      break;
-    }
-    }
+      }
+    } else
+      llvm::report_fatal_error("do not know how to handle this incoming bb argument");
   }
 
   // 3. Build the pullback successor cases for the `switch_enum`

test/AutoDiff/compiler_crashers_fixed/issue-63728-try-apply-activity.swift

@@ -0,0 +1,59 @@
+// RUN: not %target-swift-frontend -emit-sil -verify %s
+
+// The testcase from https://github.com/apple/swift/issues/63728 is not valid
+// (the function is not differentiable), however, it should not cause verifier errors
+// Here the root case is lack of activity analysis for `try_apply` terminators
+
+import _Differentiation
+
+func a() throws {
+    let keyPaths = (readable: [String: KeyPath<T, Double>](), writable: [String: WritableKeyPath<T, Double>]())
+    @differentiable(reverse)
+    func f(p: PAndT) -> Double {
+        var mutableP = p
+        let s = p.p.e
+        var sArray: [[Double]] = []
+        sArray.append((s["a"]!.asArray()).map {$0.value})
+        mutableP.s = w(mutableP.s, at: keyPaths.writable["a"]!, with: sArray[0][0])
+        return mutableP.s[keyPath: keyPaths.writable["a"]!]
+    }
+}
+
+public struct S<I: SProtocol, D> {
+    public func asArray() -> [(index: I, value: D)] {
+        return [(index: I, value: D)]()
+    }
+}
+struct T: Differentiable {}
+struct P: Differentiable {
+    public var e: F<Double, Double>
+}
+
+struct PAndT: Differentiable{
+    @differentiable(reverse) public var p: P
+    @differentiable(reverse) public var s: T
+}
+
+public struct F<I: SProtocol, D>{
+    public func asArray() -> [(index: I, value: D)] {return [(index: I, value: D)]()}
+    public subscript(_ name: String) -> S<I, D>? {get { return self.sGet(name) }}
+    func sGet(_ name: String) -> S<I, D>? { fatalError("") }
+}
+
+public protocol ZProtocol: Differentiable {var z: () -> TangentVector { get }}
+public protocol SProtocol: Hashable {}
+
+//extension S: Differentiable where D: ZProtocol, D.TangentVector == D {}
+
+extension F: Differentiable where D: ZProtocol, D.TangentVector == D {}
+public extension ZProtocol {var z: () -> TangentVector {{ Self.TangentVector.zero }}}
+
+extension F: Equatable where I: Equatable, D: Equatable {}
+//extension S: Equatable where I: Equatable, D: Equatable {}
+extension Double: SProtocol, ZProtocol {}
+
+@differentiable(reverse where O: Differentiable, M: ZProtocol)
+func w<O, M>(_ o: O, at m: WritableKeyPath<O, M>, with v: M) -> O {return o}
+
+@derivative(of: w)
+func vjpw<O, M>(_ o: O, at m: WritableKeyPath<O, M>, with v: M) -> (value: O, pullback: (O.TangentVector) -> (O.TangentVector, M.TangentVector)) where O: Differentiable, M: ZProtocol{fatalError("")}