[AutoDiff] Handle materializing adjoints with non-differentiable fields

Fixes #66522

Author: jkshtj

include/swift/SILOptimizer/Differentiation/AdjointValue.h

@@ -42,16 +42,8 @@ enum AdjointValueKind {
   Concrete,
 
   /// A special adjoint, made up of 2 adjoints -- a base adjoint and an element
-  /// adjoint to add to it. This case exists due to the existence of custom
-  /// tangent vectors which may comprise of non-differentiable fields and may
-  /// be used in the adjoint of the `struct_extract` and `tuple_extract` SIL
-  /// instructions.
-  ///
-  /// The adjoints for such tangent vectors are not pieceswise materializable,
-  /// i.e., cannot be materialized by materializing individual fields. Therefore
-  /// when used w/ a `struct_extact`/`tuple_extract` they must be materialized
-  /// by first creating a zero tangent vector of the base adjoint and then
-  /// in-place adding element adjoint to the specified field.
+  /// adjoint to add to it. This case exists to avoid eager materialization of
+  /// `Aggregate` and `Zero` adjoints upon addition.
   AddElement,
 };
 
@@ -227,13 +219,15 @@ struct AddElementValue final {
   AddElementValue(AdjointValue baseAdjoint, AdjointValue eltToAdd,
                   VarDecl *field)
       : baseAdjoint(baseAdjoint), eltToAdd(eltToAdd), fieldLocator(field) {
-    assert(baseAdjoint.getKind() == AdjointValueKind::Zero);
+    assert(baseAdjoint.getKind() == AdjointValueKind::Zero ||
+           baseAdjoint.getKind() == AdjointValueKind::Aggregate);
   }
 
   AddElementValue(AdjointValue baseAdjoint, AdjointValue eltToAdd,
                   unsigned int index)
       : baseAdjoint(baseAdjoint), eltToAdd(eltToAdd), fieldLocator(index) {
-    assert(baseAdjoint.getKind() == AdjointValueKind::Zero);
+    assert(baseAdjoint.getKind() == AdjointValueKind::Zero ||
+           baseAdjoint.getKind() == AdjointValueKind::Aggregate);
   }
 
   bool isStructAdjoint() const {

lib/SILOptimizer/Differentiation/AdjointValue.cpp

@@ -51,14 +51,16 @@ void swift::autodiff::AdjointValue::print(llvm::raw_ostream &s) const {
     auto *addElementValue = getAddElementValue();
     auto baseAdjoint = addElementValue->baseAdjoint;
     auto eltToAdd = addElementValue->eltToAdd;
+    auto baseAdjointKind =
+        baseAdjoint.getKind() == AdjointValueKind::Zero ? "Zero" : "Aggregate";
 
     s << "AddElement[";
     if (addElementValue->isTupleAdjoint()) {
-      s << "Zero[" << baseAdjoint.getType() << "].#n["
+      s << baseAdjointKind << "[" << baseAdjoint.getType() << "].#n["
         << addElementValue->getFieldIndex() << "] += Concrete["
         << eltToAdd.getType() << "]";
     } else {
-      s << "Zero[" << baseAdjoint.getType() << "].#field["
+      s << baseAdjointKind << "[" << baseAdjoint.getType() << "].#field["
         << addElementValue->getFieldDecl()->getNameStr() << "] += Concrete["
         << eltToAdd.getType() << "]";
     }