[AutoDiff] Implement active Optional differentiation (#74977)

Fixes #74972

Author: asl

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -1803,6 +1803,57 @@ class PullbackCloner::Implementation final
     builder.emitZeroIntoBuffer(uccai->getLoc(), adjDest, IsInitialization);
   }
 
+  /// Handle `enum` instruction.
+  ///   Original: y = enum $Enum, #Enum.some!enumelt, x
+  ///    Adjoint: adj[x] += adj[y]
+  void visitEnumInst(EnumInst *ei) {
+    SILBasicBlock *bb = ei->getParent();
+    SILLocation loc = ei->getLoc();
+    auto *optionalEnumDecl = getASTContext().getOptionalDecl();
+
+    // Only `Optional`-typed operands are supported for now. Diagnose all other
+    // enum operand types.
+    if (ei->getType().getEnumOrBoundGenericEnum() != optionalEnumDecl) {
+      LLVM_DEBUG(getADDebugStream()
+                 << "Unsupported enum type in PullbackCloner: " << *ei);
+      getContext().emitNondifferentiabilityError(
+          ei, getInvoker(),
+          diag::autodiff_expression_not_differentiable_note);
+      errorOccurred = true;
+      return;
+    }
+
+    auto adjOpt = getAdjointValue(bb, ei);
+    auto adjStruct = materializeAdjointDirect(adjOpt, loc);
+    StructDecl *adjStructDecl =
+        adjStruct->getType().getStructOrBoundGenericStruct();
+
+    VarDecl *adjOptVar = nullptr;
+    if (adjStructDecl) {
+      ArrayRef<VarDecl *> properties = adjStructDecl->getStoredProperties();
+      adjOptVar = properties.size() == 1 ? properties[0] : nullptr;
+    }
+
+    EnumDecl *adjOptDecl =
+        adjOptVar ? adjOptVar->getTypeInContext()->getEnumOrBoundGenericEnum()
+                  : nullptr;
+
+    // Optional<T>.TangentVector should be a struct with a single
+    // Optional<T.TangentVector> property. This is an implementation detail of
+    // OptionalDifferentiation.swift
+    // TODO: Maybe it would be better to have getters / setters here that we
+    // can call and hide this implementation detail?
+    if (!adjOptDecl || adjOptDecl != optionalEnumDecl)
+      llvm_unreachable("Unexpected type of Optional.TangentVector");
+
+    auto *adjVal = builder.createStructExtract(loc, adjStruct, adjOptVar);
+
+    EnumElementDecl *someElemDecl = getASTContext().getOptionalSomeDecl();
+    auto *adjData = builder.createUncheckedEnumData(loc, adjVal, someElemDecl);
+
+    addAdjointValue(bb, ei->getOperand(), makeConcreteAdjointValue(adjData), loc);
+  }
+
   /// Handle a sequence of `init_enum_data_addr` and `inject_enum_addr`
   /// instructions.
   ///

test/AutoDiff/SILOptimizer/activity_analysis.swift

@@ -197,11 +197,8 @@ func checked_cast_addr_nonactive_result<T: Differentiable>(_ x: T) -> T {
 // CHECK:   checked_cast_addr_br take_always T in %3 : $*T to Float in %5 : $*Float, bb1, bb2
 // CHECK: }
 
-// expected-error @+1 {{function is not differentiable}}
 @differentiable(reverse)
-// expected-note @+1 {{when differentiating this function definition}}
 func checked_cast_addr_active_result<T: Differentiable>(x: T) -> T {
-  // expected-note @+1 {{expression is not differentiable}}
   if let y = x as? Float {
     // Use `y: Float?` value in an active way.
     return y as! T
@@ -777,12 +774,9 @@ func testClassModifyAccessor(_ c: inout C) {
 // Enum differentiation
 //===----------------------------------------------------------------------===//
 
-// expected-error @+1 {{function is not differentiable}}
 @differentiable(reverse)
-// expected-note @+1 {{when differentiating this function definition}}
 func testActiveOptional(_ x: Float) -> Float {
   var maybe: Float? = 10
-  // expected-note @+1 {{expression is not differentiable}}
   maybe = x
   return maybe!
 }

test/AutoDiff/SILOptimizer/differentiation_diagnostics.swift

@@ -147,12 +147,9 @@ class C<T: Differentiable>: Differentiable {
 // Enum differentiation
 //===----------------------------------------------------------------------===//
 
-// expected-error @+1 {{function is not differentiable}}
 @differentiable(reverse)
-// expected-note @+1 {{when differentiating this function definition}}
 func usesOptionals(_ x: Float) -> Float {
   var maybe: Float? = 10
-  // expected-note @+1 {{expression is not differentiable}}
   maybe = x
   return maybe!
 }

test/AutoDiff/validation-test/optional.swift

@@ -21,6 +21,15 @@ func optional_nil_coalescing(_ maybeX: Float?) -> Float {
 }
 */
 
+OptionalTests.test("Active") {
+  @differentiable(reverse)
+  func square(y: Float) -> Float? {
+    return y * y
+  }
+
+  expectEqual(gradient(at: 10, of: {y in square(y:y)!}), .init(20.0))
+}
+
 OptionalTests.test("Let") {
   @differentiable(reverse)
   func optional_let(_ maybeX: Float?) -> Float {