[AutoDiff] Handle materializing adjoints with non-differentiable fields

Fixes #66522

Author: jkshtj

include/swift/SILOptimizer/Differentiation/AdjointValue.h

@@ -41,9 +41,10 @@ enum AdjointValueKind {
   /// A concrete SIL value.
   Concrete,
 
-  /// A special adjoint, made up of 2 adjoints -- a base adjoint and an element
-  /// adjoint to add to it. This case exists to avoid eager materialization of
-  /// `Aggregate` and `Zero` adjoints upon addition.
+  /// A special adjoint, made up of 2 adjoints -- an aggregate base adjoint and
+  /// an element adjoint to add to one of its fields. This case exists to avoid
+  /// eager materialization of a base adjoint upon addition with one of its
+  /// fields.
   AddElement,
 };
 
@@ -201,49 +202,64 @@ class AdjointValue final {
   SWIFT_DEBUG_DUMP { print(llvm::dbgs()); };
 };
 
-/// @brief The underlying value for an `AddElement` adjoint.
-struct AddElementValue final {
-  AdjointValue baseAdjoint;
-  AdjointValue eltToAdd;
+/// An abstraction that represents the field locator in
+/// an `AddElement` adjoint kind. Depending on the aggregate
+/// kind - tuple or struct, of the `baseAdjoint` in an
+/// `AddElement` adjoint, the field locator may be an `unsigned int`
+/// or a `VarDecl *`.
+struct FieldLocator final {
+  FieldLocator(VarDecl *field) : inner(field) {}
+  FieldLocator(unsigned int index) : inner(index) {}
+
+  friend AddElementValue;
 
 private:
-  union FieldLocator {
-    VarDecl *field;
-    unsigned int index;
+  bool isTupleFieldLocator() const {
+    return std::holds_alternative<unsigned int>(inner);
+  }
 
-    FieldLocator(VarDecl *field) : field(field) {}
-    FieldLocator(unsigned int index) : index(index) {}
-  } fieldLocator;
+  const static constexpr std::true_type TUPLE_FIELD_LOCATOR_TAG =
+      std::true_type{};
+  const static constexpr std::false_type STRUCT_FIELD_LOCATOR_TAG =
+      std::false_type{};
 
-public:
-  AddElementValue(AdjointValue baseAdjoint, AdjointValue eltToAdd,
-                  VarDecl *field)
-      : baseAdjoint(baseAdjoint), eltToAdd(eltToAdd), fieldLocator(field) {
-    assert(baseAdjoint.getKind() == AdjointValueKind::Zero ||
-           baseAdjoint.getKind() == AdjointValueKind::Aggregate);
+  unsigned int getInner(std::true_type) const {
+    return std::get<unsigned int>(inner);
   }
 
-  AddElementValue(AdjointValue baseAdjoint, AdjointValue eltToAdd,
-                  unsigned int index)
-      : baseAdjoint(baseAdjoint), eltToAdd(eltToAdd), fieldLocator(index) {
-    assert(baseAdjoint.getKind() == AdjointValueKind::Zero ||
-           baseAdjoint.getKind() == AdjointValueKind::Aggregate);
+  VarDecl *getInner(std::false_type) const {
+    return std::get<VarDecl *>(inner);
   }
 
-  bool isStructAdjoint() const {
-    return !baseAdjoint.getType().is<TupleType>();
+  std::variant<unsigned int, VarDecl *> inner;
+};
+
+/// The underlying value for an `AddElement` adjoint.
+struct AddElementValue final {
+  AdjointValue baseAdjoint;
+  AdjointValue eltToAdd;
+  FieldLocator fieldLocator;
+
+  AddElementValue(AdjointValue baseAdjoint, AdjointValue eltToAdd,
+                  FieldLocator fieldLocator)
+      : baseAdjoint(baseAdjoint), eltToAdd(eltToAdd),
+        fieldLocator(fieldLocator) {
+    assert(baseAdjoint.getType().is<TupleType>() ||
+           baseAdjoint.getType().getStructOrBoundGenericStruct() != nullptr);
   }
 
-  bool isTupleAdjoint() const { return baseAdjoint.getType().is<TupleType>(); }
+  bool isTupleAdjoint() const { return fieldLocator.isTupleFieldLocator(); }
+
+  bool isStructAdjoint() const { return !isTupleAdjoint(); }
 
   VarDecl *getFieldDecl() const {
     assert(isStructAdjoint());
-    return this->fieldLocator.field;
+    return this->fieldLocator.getInner(FieldLocator::STRUCT_FIELD_LOCATOR_TAG);
   }
 
   unsigned int getFieldIndex() const {
     assert(isTupleAdjoint());
-    return this->fieldLocator.index;
+    return this->fieldLocator.getInner(FieldLocator::TUPLE_FIELD_LOCATOR_TAG);
   }
 };
 

lib/SILOptimizer/Differentiation/AdjointValue.cpp

@@ -51,19 +51,20 @@ 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[";
+    baseAdjoint.print(s);
+
+    s << ", Field(";
     if (addElementValue->isTupleAdjoint()) {
-      s << baseAdjointKind << "[" << baseAdjoint.getType() << "].#n["
-        << addElementValue->getFieldIndex() << "] += Concrete["
-        << eltToAdd.getType() << "]";
+      s << addElementValue->getFieldIndex();
     } else {
-      s << baseAdjointKind << "[" << baseAdjoint.getType() << "].#field["
-        << addElementValue->getFieldDecl()->getNameStr() << "] += Concrete["
-        << eltToAdd.getType() << "]";
+      s << addElementValue->getFieldDecl()->getNameStr();
     }
+    s << "), ";
+
+    eltToAdd.print(s);
+
     s << "]";
     break;
   }