[AutoDiff] Handle materializing adjoints with non-differentiable fields

Fixes #66522

Author: jkshtj

include/swift/SILOptimizer/Differentiation/AdjointValue.h

@@ -19,6 +19,8 @@
 #define SWIFT_SILOPTIMIZER_UTILS_DIFFERENTIATION_ADJOINTVALUE_H
 
 #include "swift/AST/Decl.h"
+#include "swift/SIL/SILDebugVariable.h"
+#include "swift/SIL/SILLocation.h"
 #include "swift/SIL/SILValue.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/Debug.h"
@@ -38,10 +40,25 @@ 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 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.
+  AddElement,
 };
 
 class AdjointValue;
 
+struct AddElementValue;
+
 class AdjointValueBase {
   friend class AdjointValue;
 
@@ -60,9 +77,13 @@ class AdjointValueBase {
   union Value {
     unsigned numAggregateElements;
     SILValue concrete;
+    AddElementValue *addElementValue;
+
     Value(unsigned numAggregateElements)
         : numAggregateElements(numAggregateElements) {}
     Value(SILValue v) : concrete(v) {}
+    Value(AddElementValue *addElementValue)
+        : addElementValue(addElementValue) {}
     Value() {}
   } value;
 
@@ -86,6 +107,11 @@ class AdjointValueBase {
 
   explicit AdjointValueBase(SILType type, llvm::Optional<DebugInfo> debugInfo)
       : kind(AdjointValueKind::Zero), type(type), debugInfo(debugInfo) {}
+
+  explicit AdjointValueBase(SILType type, AddElementValue *addElementValue,
+                            llvm::Optional<DebugInfo> debugInfo)
+      : kind(AdjointValueKind::AddElement), type(type), debugInfo(debugInfo),
+        value(addElementValue) {}
 };
 
 /// A symbolic adjoint value that wraps a `SILValue`, a zero, or an aggregate
@@ -127,6 +153,14 @@ class AdjointValue final {
     return new (buf) AdjointValueBase(type, elements, debugInfo);
   }
 
+  static AdjointValue
+  createAddElement(llvm::BumpPtrAllocator &allocator, SILType type,
+                   AddElementValue *addElementValue,
+                   llvm::Optional<DebugInfo> debugInfo = llvm::None) {
+    auto *buf = allocator.Allocate<AdjointValueBase>();
+    return new (buf) AdjointValueBase(type, addElementValue, debugInfo);
+  }
+
   AdjointValueKind getKind() const { return base->kind; }
   SILType getType() const { return base->type; }
   CanType getSwiftType() const { return getType().getASTType(); }
@@ -140,6 +174,9 @@ class AdjointValue final {
   bool isZero() const { return getKind() == AdjointValueKind::Zero; }
   bool isAggregate() const { return getKind() == AdjointValueKind::Aggregate; }
   bool isConcrete() const { return getKind() == AdjointValueKind::Concrete; }
+  bool isAddElement() const {
+    return getKind() == AdjointValueKind::AddElement;
+  }
 
   unsigned getNumAggregateElements() const {
     assert(isAggregate());
@@ -162,41 +199,60 @@ class AdjointValue final {
     return base->value.concrete;
   }
 
-  void print(llvm::raw_ostream &s) const {
-    switch (getKind()) {
-    case AdjointValueKind::Zero:
-      s << "Zero[" << getType() << ']';
-      break;
-    case AdjointValueKind::Aggregate:
-      s << "Aggregate[" << getType() << "](";
-      if (auto *decl =
-              getType().getASTType()->getStructOrBoundGenericStruct()) {
-        interleave(
-            llvm::zip(decl->getStoredProperties(), getAggregateElements()),
-            [&s](std::tuple<VarDecl *, const AdjointValue &> elt) {
-              s << std::get<0>(elt)->getName() << ": ";
-              std::get<1>(elt).print(s);
-            },
-            [&s] { s << ", "; });
-      } else if (getType().is<TupleType>()) {
-        interleave(
-            getAggregateElements(),
-            [&s](const AdjointValue &elt) { elt.print(s); },
-            [&s] { s << ", "; });
-      } else {
-        llvm_unreachable("Invalid aggregate");
-      }
-      s << ')';
-      break;
-    case AdjointValueKind::Concrete:
-      s << "Concrete[" << getType() << "](" << base->value.concrete << ')';
-      break;
-    }
+  AddElementValue *getAddElementValue() const {
+    assert(isAddElement());
+    return base->value.addElementValue;
   }
 
+  void print(llvm::raw_ostream &s) const;
+
   SWIFT_DEBUG_DUMP { print(llvm::dbgs()); };
 };
 
+/// @brief The underlying value for an `AddElement` adjoint.
+struct AddElementValue final {
+  AdjointValue baseAdjoint;
+  AdjointValue eltToAdd;
+
+private:
+  union FieldLocator {
+    VarDecl *field;
+    unsigned int index;
+
+    FieldLocator(VarDecl *field) : field(field) {}
+    FieldLocator(unsigned int index) : index(index) {}
+  } fieldLocator;
+
+public:
+  AddElementValue(AdjointValue baseAdjoint, AdjointValue eltToAdd,
+                  VarDecl *field)
+      : baseAdjoint(baseAdjoint), eltToAdd(eltToAdd), fieldLocator(field) {
+    assert(baseAdjoint.getKind() == AdjointValueKind::Zero);
+  }
+
+  AddElementValue(AdjointValue baseAdjoint, AdjointValue eltToAdd,
+                  unsigned int index)
+      : baseAdjoint(baseAdjoint), eltToAdd(eltToAdd), fieldLocator(index) {
+    assert(baseAdjoint.getKind() == AdjointValueKind::Zero);
+  }
+
+  bool isStructAdjoint() const {
+    return !baseAdjoint.getType().is<TupleType>();
+  }
+
+  bool isTupleAdjoint() const { return baseAdjoint.getType().is<TupleType>(); }
+
+  VarDecl *getFieldDecl() const {
+    assert(isStructAdjoint());
+    return this->fieldLocator.field;
+  }
+
+  unsigned int getFieldIndex() const {
+    assert(isTupleAdjoint());
+    return this->fieldLocator.index;
+  }
+};
+
 inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
                                      const AdjointValue &adjVal) {
   adjVal.print(os);

lib/SILOptimizer/Differentiation/AdjointValue.cpp

@@ -0,0 +1,68 @@
+//===--- AdjointValue.cpp - Helper class for differentiation ----*- C++
+//-*---===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2019 - 2020 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// AdjointValue - a symbolic representation for adjoint values enabling
+// efficient differentiation by avoiding zero materialization.
+//
+//===----------------------------------------------------------------------===//
+
+#define DEBUG_TYPE "differentiation"
+
+#include "swift/SILOptimizer/Differentiation/AdjointValue.h"
+
+void swift::autodiff::AdjointValue::print(llvm::raw_ostream &s) const {
+  switch (getKind()) {
+  case AdjointValueKind::Zero:
+    s << "Zero[" << getType() << ']';
+    break;
+  case AdjointValueKind::Aggregate:
+    s << "Aggregate[" << getType() << "](";
+    if (auto *decl = getType().getASTType()->getStructOrBoundGenericStruct()) {
+      interleave(
+          llvm::zip(decl->getStoredProperties(), getAggregateElements()),
+          [&s](std::tuple<VarDecl *, const AdjointValue &> elt) {
+            s << std::get<0>(elt)->getName() << ": ";
+            std::get<1>(elt).print(s);
+          },
+          [&s] { s << ", "; });
+    } else if (getType().is<TupleType>()) {
+      interleave(
+          getAggregateElements(),
+          [&s](const AdjointValue &elt) { elt.print(s); }, [&s] { s << ", "; });
+    } else {
+      llvm_unreachable("Invalid aggregate");
+    }
+    s << ')';
+    break;
+  case AdjointValueKind::Concrete:
+    s << "Concrete[" << getType() << "](" << base->value.concrete << ')';
+    break;
+  case AdjointValueKind::AddElement:
+    auto *addElementValue = getAddElementValue();
+    auto baseAdjoint = addElementValue->baseAdjoint;
+    auto eltToAdd = addElementValue->eltToAdd;
+
+    s << "AddElement[";
+    if (addElementValue->isTupleAdjoint()) {
+      s << "Zero[" << baseAdjoint.getType() << "].#n["
+        << addElementValue->getFieldIndex() << "] += Concrete["
+        << eltToAdd.getType() << "]";
+    } else {
+      s << "Zero[" << baseAdjoint.getType() << "].#field["
+        << addElementValue->getFieldDecl()->getNameStr() << "] += Concrete["
+        << eltToAdd.getType() << "]";
+    }
+    s << "]";
+    break;
+  }
+}

lib/SILOptimizer/Differentiation/CMakeLists.txt

@@ -1,5 +1,6 @@
 target_sources(swiftSILOptimizer PRIVATE
   ADContext.cpp
+  AdjointValue.cpp
   Common.cpp
   DifferentiationInvoker.cpp
   JVPCloner.cpp

lib/SILOptimizer/Differentiation/JVPCloner.cpp

@@ -228,11 +228,12 @@ class JVPCloner::Implementation final
       auto zeroVal = emitZeroDirect(val.getSwiftType(), loc);
       return zeroVal;
     }
+    case AdjointValueKind::Concrete:
+      return val.getConcreteValue();
     case AdjointValueKind::Aggregate:
+    case AdjointValueKind::AddElement:
       llvm_unreachable(
           "Tuples and structs are not supported in forward mode yet.");
-    case AdjointValueKind::Concrete:
-      return val.getConcreteValue();
     }
     llvm_unreachable("Invalid adjoint value kind"); // silences MSVC C4715
   }