[AutoDiff] Handle materializing adjoints with non-differentiable fields

jkshtj · jkshtj · commit 005a208ec713 · 2023-07-14T15:26:58.000-07:00
Fixes #66522
diff --git a/lib/SILOptimizer/Differentiation/PullbackCloner.cpp b/lib/SILOptimizer/Differentiation/PullbackCloner.cpp
@@ -327,6 +327,30 @@ class PullbackCloner::Implementation final
   // Adjoint value materialization
   //--------------------------------------------------------------------------//
 
+  /// Determines whether this adjoint value can be materialized by materializing
+  /// and then combining individual fields of the adjoint. This function should
+  /// only be called for aggregate adjoints.
+  ///
+  /// Users are allowed to define custom tangent vectors which may contain
+  /// fields that do not conform to `AdditiveArithmetic` and `Differentiable`
+  /// protocols. This is fine as long as they fulfill the corresponding
+  /// `AdditiveArithmetic` and `Differentiable` protocol requirements on the
+  /// tangent vector itself.
+  ///
+  /// A user-defined tangent vector with above characteristics, cannot be
+  /// materialized by materializing individual fields as that process relies on
+  /// the `AdditiveArithmeticness` of the individual fields.
+  bool isAdjointPiecewiseMaterializable(AdjointValue val) {
+    assert(val.getKind() == AdjointValueKind::Aggregate);
+    for (auto i : range(val.getNumAggregateElements())) {
+      auto fieldCanTy = val.getAggregateElement(i).getType().getASTType();
+      if (!getTangentSpace(fieldCanTy)) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   /// Materializes an adjoint value. The type of the given adjoint value must be
   /// loadable.
   SILValue materializeAdjointDirect(AdjointValue val, SILLocation loc) {
@@ -339,17 +363,11 @@ class PullbackCloner::Implementation final
       result = recordTemporary(builder.emitZero(loc, val.getSwiftType()));
       break;
     case AdjointValueKind::Aggregate: {
-      SmallVector<SILValue, 8> elements;
-      for (auto i : range(val.getNumAggregateElements())) {
-        auto eltVal = materializeAdjointDirect(val.getAggregateElement(i), loc);
-        elements.push_back(builder.emitCopyValueOperation(loc, eltVal));
+      if (isAdjointPiecewiseMaterializable(val)) {
+        result = materializeAggregateAdjointDirectPiecewise(val, loc);
+      } else {
+        result = materializeAggregateAdjointDirect(val, loc);
       }
-      if (val.getType().is<TupleType>())
-        result = recordTemporary(
-            builder.createTuple(loc, val.getType(), elements));
-      else
-        result = recordTemporary(
-            builder.createStruct(loc, val.getType(), elements));
       break;
     }
     case AdjointValueKind::Concrete:
@@ -362,6 +380,68 @@ class PullbackCloner::Implementation final
     return result;
   }
 
+  /// Used to materialize an aggregate adjoint directly, if
+  /// `isAdjointPiecewiseMaterializable` returned false.
+  SILValue materializeAggregateAdjointDirect(AdjointValue val,
+                                             SILLocation loc) {
+    SILValue result = builder.emitZero(loc, val.getSwiftType());
+    auto isTupleType = val.getType().is<TupleType>();
+
+    SILInstruction *destructureInst;
+    if (isTupleType) {
+      destructureInst = builder.createDestructureTuple(loc, result);
+    } else {
+      destructureInst = builder.createDestructureStruct(loc, result);
+    }
+
+    // Note - Materializing the aggregate adjoints by collecting concrete
+    // field values and using the `struct` instruction rather than doing a
+    // `+=` on individual fields buffers because the TangentBuilder's
+    // `emitInPlaceAdd` method requires the destination buffer to be an address,
+    // which does not work for trivial types such as floats and doubles.
+    SmallVector<SILValue, 8> elements;
+    for (auto i : range(val.getNumAggregateElements())) {
+      auto fieldCanTy = val.getAggregateElement(i).getType().getASTType();
+
+      if (!getTangentSpace(fieldCanTy)) {
+        elements.push_back(destructureInst->getResult(i));
+      } else {
+        auto eltVal = materializeAdjointDirect(val.getAggregateElement(i), loc);
+        elements.push_back(builder.emitCopyValueOperation(loc, eltVal));
+      }
+    }
+
+    if (isTupleType)
+      result =
+          recordTemporary(builder.createTuple(loc, val.getType(), elements));
+    else
+      result =
+          recordTemporary(builder.createStruct(loc, val.getType(), elements));
+
+    return result;
+  }
+
+  /// Used to materialize an aggregate adjoint directly, if
+  /// `isAdjointPiecewiseMaterializable` returned true.
+  SILValue materializeAggregateAdjointDirectPiecewise(AdjointValue val,
+                                                      SILLocation loc) {
+    SILValue result;
+
+    SmallVector<SILValue, 8> elements;
+    for (auto i : range(val.getNumAggregateElements())) {
+      auto eltVal = materializeAdjointDirect(val.getAggregateElement(i), loc);
+      elements.push_back(builder.emitCopyValueOperation(loc, eltVal));
+    }
+    if (val.getType().is<TupleType>())
+      result =
+          recordTemporary(builder.createTuple(loc, val.getType(), elements));
+    else
+      result =
+          recordTemporary(builder.createStruct(loc, val.getType(), elements));
+
+    return result;
+  }
+
   /// Materializes an adjoint value indirectly to a SIL buffer.
   void materializeAdjointIndirect(AdjointValue val, SILValue destAddress,
                                   SILLocation loc) {
@@ -376,25 +456,10 @@ class PullbackCloner::Implementation final
     /// materialize the symbolic tuple or struct, filling the
     /// buffer.
     case AdjointValueKind::Aggregate: {
-      if (auto *tupTy = val.getSwiftType()->getAs<TupleType>()) {
-        for (auto idx : range(val.getNumAggregateElements())) {
-          auto eltTy = SILType::getPrimitiveAddressType(
-              tupTy->getElementType(idx)->getCanonicalType());
-          auto *eltBuf =
-              builder.createTupleElementAddr(loc, destAddress, idx, eltTy);
-          materializeAdjointIndirect(val.getAggregateElement(idx), eltBuf, loc);
-        }
-      } else if (auto *structDecl =
-                     val.getSwiftType()->getStructOrBoundGenericStruct()) {
-        auto fieldIt = structDecl->getStoredProperties().begin();
-        for (unsigned i = 0; fieldIt != structDecl->getStoredProperties().end();
-             ++fieldIt, ++i) {
-          auto eltBuf =
-              builder.createStructElementAddr(loc, destAddress, *fieldIt);
-          materializeAdjointIndirect(val.getAggregateElement(i), eltBuf, loc);
-        }
+      if (isAdjointPiecewiseMaterializable(val)) {
+        materializeAggregateAdjointInDirectPiecewise(val, destAddress, loc);
       } else {
-        llvm_unreachable("Not an aggregate type");
+        materializeAggregateAdjointInDirect(val, destAddress, loc);
       }
       break;
     }
@@ -408,6 +473,87 @@ class PullbackCloner::Implementation final
     }
   }
 
+  /// Used to materialize an aggregate adjoint indirectly, if
+  /// `isAdjointPiecewiseMaterializable` returned false.
+  void materializeAggregateAdjointInDirect(AdjointValue val,
+                                           SILValue destAddress,
+                                           SILLocation loc) {
+    assert(destAddress->getType().isAddress());
+    auto zeroAggAdj = builder.emitZero(loc, val.getSwiftType());
+    auto isTupleType = val.getType().is<TupleType>();
+
+    SILInstruction *destructureInst;
+    if (isTupleType)
+      destructureInst = builder.createDestructureTuple(loc, zeroAggAdj);
+    else
+      destructureInst = builder.createDestructureStruct(loc, zeroAggAdj);
+
+    if (auto *tupTy = val.getSwiftType()->getAs<TupleType>()) {
+      for (auto idx : range(val.getNumAggregateElements())) {
+        auto eltTy = SILType::getPrimitiveAddressType(
+            tupTy->getElementType(idx)->getCanonicalType());
+        auto *eltBuf =
+            builder.createTupleElementAddr(loc, destAddress, idx, eltTy);
+
+        auto fieldCanTy = val.getAggregateElement(idx).getType().getASTType();
+        if (!getTangentSpace(fieldCanTy)) {
+          builder.emitStoreValueOperation(loc, destructureInst->getResult(idx),
+                                          eltBuf,
+                                          StoreOwnershipQualifier::Init);
+        } else {
+          materializeAdjointIndirect(val.getAggregateElement(idx), eltBuf, loc);
+        }
+      }
+    } else if (auto *structDecl =
+                   val.getSwiftType()->getStructOrBoundGenericStruct()) {
+      auto fieldIt = structDecl->getStoredProperties().begin();
+      for (unsigned i = 0; fieldIt != structDecl->getStoredProperties().end();
+           ++fieldIt, ++i) {
+        auto eltBuf =
+            builder.createStructElementAddr(loc, destAddress, *fieldIt);
+
+        auto fieldCanTy = val.getAggregateElement(i).getType().getASTType();
+        if (!getTangentSpace(fieldCanTy)) {
+          builder.emitStoreValueOperation(loc, destructureInst->getResult(i),
+                                          eltBuf,
+                                          StoreOwnershipQualifier::Init);
+        } else {
+          materializeAdjointIndirect(val.getAggregateElement(i), eltBuf, loc);
+        }
+      }
+    } else {
+      llvm_unreachable("Not an aggregate type");
+    }
+  }
+
+  /// Used to materialize an aggregate adjoint indirectly, if
+  /// `isAdjointPiecewiseMaterializable` returned true.
+  void materializeAggregateAdjointInDirectPiecewise(AdjointValue val,
+                                                    SILValue destAddress,
+                                                    SILLocation loc) {
+    assert(destAddress->getType().isAddress());
+    if (auto *tupTy = val.getSwiftType()->getAs<TupleType>()) {
+      for (auto idx : range(val.getNumAggregateElements())) {
+        auto eltTy = SILType::getPrimitiveAddressType(
+            tupTy->getElementType(idx)->getCanonicalType());
+        auto *eltBuf =
+            builder.createTupleElementAddr(loc, destAddress, idx, eltTy);
+        materializeAdjointIndirect(val.getAggregateElement(idx), eltBuf, loc);
+      }
+    } else if (auto *structDecl =
+                   val.getSwiftType()->getStructOrBoundGenericStruct()) {
+      auto fieldIt = structDecl->getStoredProperties().begin();
+      for (unsigned i = 0; fieldIt != structDecl->getStoredProperties().end();
+           ++fieldIt, ++i) {
+        auto eltBuf =
+            builder.createStructElementAddr(loc, destAddress, *fieldIt);
+        materializeAdjointIndirect(val.getAggregateElement(i), eltBuf, loc);
+      }
+    } else {
+      llvm_unreachable("Not an aggregate type");
+    }
+  }
+
   //--------------------------------------------------------------------------//
   // Adjoint value mapping
   //--------------------------------------------------------------------------//
diff --git a/test/AutoDiff/compiler_crashers_fixed/issue-66522-pullback-generation-when-tangentvector-of-input-contains-nondifferentiable-fields.swift b/test/AutoDiff/compiler_crashers_fixed/issue-66522-pullback-generation-when-tangentvector-of-input-contains-nondifferentiable-fields.swift
@@ -0,0 +1,67 @@
+// RUN: %target-swift-frontend -emit-sil -O %s
+
+import _Differentiation
+
+// Issue #66522:
+// Pullback generation for a function mapping a differentiable input type to 
+// one of its differentiable fields fails when the input's tangent vector contains
+// non-differentiable fields.                          
+public struct P<Value>: Differentiable
+where
+    Value:  Differentiable,
+    Value.TangentVector == Value,
+    Value: AdditiveArithmetic {
+  // `P` is its own `TangentVector`
+  public typealias TangentVector = Self
+  
+  // Non-differentiable field in `P`'s `TangentVector`.
+  public let name: String = ""
+  var value: Value
+}
+
+extension P: Equatable, AdditiveArithmetic
+where Value: AdditiveArithmetic {
+    public static var zero: Self {fatalError()}
+    public static func + (lhs: Self, rhs: Self) -> Self {fatalError()}
+    public static func - (lhs: Self, rhs: Self) -> Self {fatalError()}
+}
+
+@differentiable(reverse)
+internal func testFunction(data: P<Double>) -> Double {
+    data.value
+}
+
+
+// WHAT?
+// This test case currently fails with the following compiler error -
+// ```
+// /Users/kshitij/workspace/swift-project/swift/test/AutoDiff/compiler_crashers_fixed/issue-66522-pullback-generation-when-tangentvector-of-input-contains-nondifferentiable-fields.swift:9:15: error: expression is not differentiable
+// public struct P<Value>: Differentiable
+//               ^
+// /Users/kshitij/workspace/swift-project/swift/test/AutoDiff/compiler_crashers_fixed/issue-66522-pullback-generation-when-tangentvector-of-input-contains-nondifferentiable-fields.swift:9:15: note: cannot differentiate access to property 'P.name' because property type 'String' does not conform to 'Differentiable'
+// public struct P<Value>: Differentiable
+//               ^
+// ```
+// 
+// WHY?
+// The `TangentVector` for `P` has been defined as `Self`.
+// The pullback generation for `P.init` fails because it now 
+// contains a projection to a non-differentiable stored
+// property, which is not allowed currently - https://github.com/apple/swift/blob/main/lib/SILOptimizer/Differentiation/PullbackCloner.cpp#L1850.
+//
+// The same issue does not happen if we let the compiler
+// generated `P`'s `TangentVector`, because `P.init` now 
+// becomes differentiable. 
+// 
+// SO WHAT? 
+// It feels like a bit of an edge case that for a user-defined tangent 
+// vector, such as that for `P`, we cannot have a function mapping 
+// individual member fields to `P`.
+//
+// NOW WHAT?
+// Is this edge case something that we need to handle?
+// 
+// @differentiable(reverse)
+// internal func testFunction2(data: Double) -> P<Double> {
+//     P(value: data)
+// }
