[AutoDiff] Improve invalid stored property projection diagnostics.

Use TangentStoredPropertyRequest in differentiation transform.

Improve non-differentiability diagnostics regarding invalid stored
property projection instructions:
`struct_extract`, `struct_element_addr`, `ref_element_addr`.

Diagnose the following cases:
- Original property's type does not conform to `Differentiable`.
- Base type's `TangentVector` is not a struct.
- Tangent property not found: base type's `TangentVector` does not have a
  stored property with the same name as the original property.
- Tangent property's type is not equal to the original property's
  `TangentVector` type.
- Tangent property is not a stored property.

Resolves TF-969 and TF-970.

Author: dan-zheng

include/swift/AST/DiagnosticsSIL.def

@@ -504,9 +504,26 @@ NOTE(autodiff_loadable_value_addressonly_tangent_unsupported,none,
      "properties", (Type, Type))
 NOTE(autodiff_enums_unsupported,none,
      "differentiating enum values is not yet supported", ())
+NOTE(autodiff_stored_property_parent_not_differentiable,none,
+     "cannot differentiate access to property '%0.%1' because '%0' does not "
+     "conform to 'Differentiable'", (StringRef, StringRef))
+NOTE(autodiff_stored_property_not_differentiable,none,
+     "cannot differentiate access to property '%0.%1' because property type %2 "
+     "does not conform to 'Differentiable'", (StringRef, StringRef, Type))
+NOTE(autodiff_stored_property_tangent_not_struct,none,
+     "cannot differentiate access to property '%0.%1' because "
+     "'%0.TangentVector' is not a struct", (StringRef, StringRef))
 NOTE(autodiff_stored_property_no_corresponding_tangent,none,
-     "property cannot be differentiated because '%0.TangentVector' does not "
-     "have a member named '%1'", (StringRef, StringRef))
+     "cannot differentiate access to property '%0.%1' because "
+     "'%0.TangentVector' does not have a stored property named '%1'",
+     (StringRef, StringRef))
+NOTE(autodiff_tangent_property_wrong_type,none,
+     "cannot differentiate access to property '%0.%1' because "
+     "'%0.TangentVector.%1' does not have expected type %2",
+     (StringRef, StringRef, /*originalPropertyTanType*/ Type))
+NOTE(autodiff_tangent_property_not_stored,none,
+     "cannot differentiate access to property '%0.%1' because "
+     "'%0.TangentVector.%1' is not a stored property", (StringRef, StringRef))
 NOTE(autodiff_coroutines_not_supported,none,
      "differentiation of coroutine calls is not yet supported", ())
 NOTE(autodiff_cannot_differentiate_writes_to_global_variables,none,

include/swift/SIL/SILInstruction.h

@@ -5758,6 +5758,13 @@ class FieldIndexCacheBase : public SingleValueInstruction {
     return s;
   }
 
+  static bool classof(const SILNode *node) {
+    SILNodeKind kind = node->getKind();
+    return kind == SILNodeKind::StructExtractInst ||
+           kind == SILNodeKind::StructElementAddrInst ||
+           kind == SILNodeKind::RefElementAddrInst;
+  }
+
 private:
   unsigned cacheFieldIndex();
 };

include/swift/SILOptimizer/Differentiation/ADContext.h

@@ -253,11 +253,9 @@ ADContext::emitNondifferentiabilityError(SILValue value,
     getADDebugStream() << "For value:\n" << value;
     getADDebugStream() << "With invoker:\n" << invoker << '\n';
   });
-  auto valueLoc = value.getLoc().getSourceLoc();
   // If instruction does not have a valid location, use the function location
   // as a fallback. Improves diagnostics in some cases.
-  if (valueLoc.isInvalid())
-    valueLoc = value->getFunction()->getLocation().getSourceLoc();
+  auto valueLoc = getValidLocation(value).getSourceLoc();
   return emitNondifferentiabilityError(valueLoc, invoker, diag,
                                        std::forward<U>(args)...);
 }
@@ -272,12 +270,10 @@ ADContext::emitNondifferentiabilityError(SILInstruction *inst,
     getADDebugStream() << "For instruction:\n" << *inst;
     getADDebugStream() << "With invoker:\n" << invoker << '\n';
   });
-  auto instLoc = inst->getLoc().getSourceLoc();
   // If instruction does not have a valid location, use the function location
   // as a fallback. Improves diagnostics for `ref_element_addr` generated in
   // synthesized stored property getters.
-  if (instLoc.isInvalid())
-    instLoc = inst->getFunction()->getLocation().getSourceLoc();
+  auto instLoc = getValidLocation(inst).getSourceLoc();
   return emitNondifferentiabilityError(instLoc, invoker, diag,
                                        std::forward<U>(args)...);
 }

include/swift/SILOptimizer/Differentiation/Common.h

@@ -17,22 +17,27 @@
 #ifndef SWIFT_SILOPTIMIZER_UTILS_DIFFERENTIATION_COMMON_H
 #define SWIFT_SILOPTIMIZER_UTILS_DIFFERENTIATION_COMMON_H
 
+#include "swift/AST/DiagnosticsSIL.h"
+#include "swift/AST/Expr.h"
 #include "swift/AST/SemanticAttrs.h"
 #include "swift/SIL/SILDifferentiabilityWitness.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILModule.h"
 #include "swift/SIL/TypeSubstCloner.h"
 #include "swift/SILOptimizer/Analysis/ArraySemantic.h"
 #include "swift/SILOptimizer/Analysis/DifferentiableActivityAnalysis.h"
+#include "swift/SILOptimizer/Differentiation/DifferentiationInvoker.h"
 
 namespace swift {
 
+namespace autodiff {
+
+class ADContext;
+
 //===----------------------------------------------------------------------===//
 // Helpers
 //===----------------------------------------------------------------------===//
 
-namespace autodiff {
-
 /// Prints an "[AD] " prefix to `llvm::dbgs()` and returns the debug stream.
 /// This is being used to print short debug messages within the AD pass.
 raw_ostream &getADDebugStream();
@@ -136,6 +141,34 @@ template <class Inst> Inst *peerThroughFunctionConversions(SILValue value) {
   return nullptr;
 }
 
+//===----------------------------------------------------------------------===//
+// Diagnostic utilities
+//===----------------------------------------------------------------------===//
+
+// Returns `v`'s location if it is valid. Otherwise, returns `v`'s function's
+// location as as a fallback. Used for diagnostics.
+SILLocation getValidLocation(SILValue v);
+
+// Returns `inst`'s location if it is valid. Otherwise, returns `inst`'s
+// function's location as as a fallback. Used for diagnostics.
+SILLocation getValidLocation(SILInstruction *inst);
+
+//===----------------------------------------------------------------------===//
+// Tangent property lookup utilities
+//===----------------------------------------------------------------------===//
+
+/// Returns the tangent stored property of `originalField`. On error, emits
+/// diagnostic and returns nullptr.
+VarDecl *getTangentStoredProperty(ADContext &context, VarDecl *originalField,
+                                  SILLocation loc,
+                                  DifferentiationInvoker invoker);
+
+/// Returns the tangent stored property of the original stored property
+/// referenced by `inst`. On error, emits diagnostic and returns nullptr.
+VarDecl *getTangentStoredProperty(ADContext &context,
+                                  FieldIndexCacheBase *projectionInst,
+                                  DifferentiationInvoker invoker);
+
 //===----------------------------------------------------------------------===//
 // Code emission utilities
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/Differentiation/Common.cpp

@@ -17,6 +17,8 @@
 #define DEBUG_TYPE "differentiation"
 
 #include "swift/SILOptimizer/Differentiation/Common.h"
+#include "swift/AST/TypeCheckRequests.h"
+#include "swift/SILOptimizer/Differentiation/ADContext.h"
 
 namespace swift {
 namespace autodiff {
@@ -244,6 +246,97 @@ void collectMinimalIndicesForFunctionCall(
   }));
 }
 
+//===----------------------------------------------------------------------===//
+// Diagnostic utilities
+//===----------------------------------------------------------------------===//
+
+SILLocation getValidLocation(SILValue v) {
+  auto loc = v.getLoc();
+  if (loc.isNull() || loc.getSourceLoc().isInvalid())
+    loc = v->getFunction()->getLocation();
+  return loc;
+}
+
+SILLocation getValidLocation(SILInstruction *inst) {
+  auto loc = inst->getLoc();
+  if (loc.isNull() || loc.getSourceLoc().isInvalid())
+    loc = inst->getFunction()->getLocation();
+  return loc;
+}
+
+//===----------------------------------------------------------------------===//
+// Tangent property lookup utilities
+//===----------------------------------------------------------------------===//
+
+VarDecl *getTangentStoredProperty(ADContext &context, VarDecl *originalField,
+                                  SILLocation loc,
+                                  DifferentiationInvoker invoker) {
+  auto &astCtx = context.getASTContext();
+  auto tanFieldInfo = evaluateOrDefault(
+      astCtx.evaluator, TangentStoredPropertyRequest{originalField},
+      TangentPropertyInfo(nullptr));
+  // If no error, return the tangent property.
+  if (tanFieldInfo)
+    return tanFieldInfo.tangentProperty;
+  // Otherwise, diagnose error and return nullptr.
+  assert(tanFieldInfo.error);
+  auto *parentDC = originalField->getDeclContext();
+  assert(parentDC->isTypeContext());
+  auto parentDeclName = parentDC->getSelfNominalTypeDecl()->getNameStr();
+  auto fieldName = originalField->getNameStr();
+  auto sourceLoc = loc.getSourceLoc();
+  switch (tanFieldInfo.error->kind) {
+  case TangentPropertyInfo::Error::Kind::NoDerivativeOriginalProperty:
+    llvm_unreachable(
+        "`@noDerivative` stored property accesses should not be "
+        "differentiated; activity analysis should not mark as varied");
+  case TangentPropertyInfo::Error::Kind::NominalParentNotDifferentiable:
+    context.emitNondifferentiabilityError(
+        sourceLoc, invoker,
+        diag::autodiff_stored_property_parent_not_differentiable,
+        parentDeclName, fieldName);
+    break;
+  case TangentPropertyInfo::Error::Kind::OriginalPropertyNotDifferentiable:
+    context.emitNondifferentiabilityError(
+        sourceLoc, invoker, diag::autodiff_stored_property_not_differentiable,
+        parentDeclName, fieldName, originalField->getInterfaceType());
+    break;
+  case TangentPropertyInfo::Error::Kind::ParentTangentVectorNotStruct:
+    context.emitNondifferentiabilityError(
+        sourceLoc, invoker, diag::autodiff_stored_property_tangent_not_struct,
+        parentDeclName, fieldName);
+    break;
+  case TangentPropertyInfo::Error::Kind::TangentPropertyNotFound:
+    context.emitNondifferentiabilityError(
+        sourceLoc, invoker,
+        diag::autodiff_stored_property_no_corresponding_tangent, parentDeclName,
+        fieldName);
+    break;
+  case TangentPropertyInfo::Error::Kind::TangentPropertyWrongType:
+    context.emitNondifferentiabilityError(
+        sourceLoc, invoker, diag::autodiff_tangent_property_wrong_type,
+        parentDeclName, fieldName, tanFieldInfo.error->getType());
+    break;
+  case TangentPropertyInfo::Error::Kind::TangentPropertyNotStored:
+    context.emitNondifferentiabilityError(
+        sourceLoc, invoker, diag::autodiff_tangent_property_not_stored,
+        parentDeclName, fieldName);
+    break;
+  }
+  return nullptr;
+}
+
+VarDecl *getTangentStoredProperty(ADContext &context,
+                                  FieldIndexCacheBase *projectionInst,
+                                  DifferentiationInvoker invoker) {
+  assert(isa<StructExtractInst>(projectionInst) ||
+         isa<StructElementAddrInst>(projectionInst) ||
+         isa<RefElementAddrInst>(projectionInst));
+  auto loc = getValidLocation(projectionInst);
+  return getTangentStoredProperty(context, projectionInst->getField(), loc,
+                                  invoker);
+}
+
 //===----------------------------------------------------------------------===//
 // Code emission utilities
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/Differentiation/JVPEmitter.cpp

@@ -547,29 +547,12 @@ CLONE_AND_EMIT_TANGENT(StructExtract, sei) {
   assert(!sei->getField()->getAttrs().hasAttribute<NoDerivativeAttr>() &&
          "`struct_extract` with `@noDerivative` field should not be "
          "differentiated; activity analysis should not marked as varied.");
-
   auto diffBuilder = getDifferentialBuilder();
-  ;
-  auto tangentVectorTy = getRemappedTangentType(sei->getOperand()->getType());
-  auto *tangentVectorDecl = tangentVectorTy.getStructOrBoundGenericStruct();
-
   // Find the corresponding field in the tangent space.
-  VarDecl *tanField = nullptr;
-  // If the tangent space is the original struct, then field is the same.
-  if (tangentVectorDecl == sei->getStructDecl())
-    tanField = sei->getField();
-  // Otherwise, look up the field by name.
-  else {
-    auto tanFieldLookup =
-        tangentVectorDecl->lookupDirect(sei->getField()->getName());
-    if (tanFieldLookup.empty()) {
-      context.emitNondifferentiabilityError(
-          sei, invoker, diag::autodiff_stored_property_no_corresponding_tangent,
-          sei->getStructDecl()->getNameStr(), sei->getField()->getNameStr());
-      errorOccurred = true;
-      return;
-    }
-    tanField = cast<VarDecl>(tanFieldLookup.front());
+  auto *tanField = getTangentStoredProperty(context, sei, invoker);
+  if (!tanField) {
+    errorOccurred = true;
+    return;
   }
   // Emit tangent `struct_extract`.
   auto tanStruct =
@@ -590,32 +573,14 @@ CLONE_AND_EMIT_TANGENT(StructElementAddr, seai) {
   assert(!seai->getField()->getAttrs().hasAttribute<NoDerivativeAttr>() &&
          "`struct_element_addr` with `@noDerivative` field should not be "
          "differentiated; activity analysis should not marked as varied.");
-
   auto diffBuilder = getDifferentialBuilder();
   auto *bb = seai->getParent();
-  auto tangentVectorTy = getRemappedTangentType(seai->getOperand()->getType());
-  auto *tangentVectorDecl = tangentVectorTy.getStructOrBoundGenericStruct();
-
   // Find the corresponding field in the tangent space.
-  VarDecl *tanField = nullptr;
-  // If the tangent space is the original struct, then field is the same.
-  if (tangentVectorDecl == seai->getStructDecl())
-    tanField = seai->getField();
-  // Otherwise, look up the field by name.
-  else {
-    auto tanFieldLookup =
-        tangentVectorDecl->lookupDirect(seai->getField()->getName());
-    if (tanFieldLookup.empty()) {
-      context.emitNondifferentiabilityError(
-          seai, invoker,
-          diag::autodiff_stored_property_no_corresponding_tangent,
-          seai->getStructDecl()->getNameStr(), seai->getField()->getNameStr());
-      errorOccurred = true;
-      return;
-    }
-    tanField = cast<VarDecl>(tanFieldLookup.front());
+  auto *tanField = getTangentStoredProperty(context, seai, invoker);
+  if (!tanField) {
+    errorOccurred = true;
+    return;
   }
-
   // Emit tangent `struct_element_addr`.
   auto tanOperand = getTangentBuffer(bb, seai->getOperand());
   auto tangentInst =