[AutoDiff] Fix stdlib @derivative(of: subscript) crash. (#28935)

Fix `@derivative(of: subscript)` crash for SIMD{n} types.

Simplify `findAbstractFunctionDecl`.
If `lookupContext` is a type context, use its `self` type for member lookup.

Resolves TF-1090.

Exposes TF-1094:
- During `.swiftinterface` compilation, unserialized `@derivative` function is
  not registered for serialized `@differentiable` function.
- Workaround: mark `SIMD{n}._vjpSubscript(_:)` as `@inlinable` (serialized).

Author: dan-zheng

lib/Sema/TypeCheckAttr.cpp

@@ -3195,20 +3195,15 @@ static AbstractFunctionDecl *findAbstractFunctionDecl(
 
   // Perform lookup.
   LookupResult results;
+  // If `baseType` is not null but `lookupContext` is a type context, set
+  // `baseType` to the `self` type of `lookupContext` to perform member lookup.
+  if (!baseType && lookupContext->isTypeContext())
+    baseType = lookupContext->getSelfTypeInContext();
   if (baseType) {
     results = TypeChecker::lookupMember(lookupContext, baseType, funcName);
   } else {
     results = TypeChecker::lookupUnqualified(lookupContext, funcName,
                                              funcNameLoc, lookupOptions);
-
-    // If looking up an operator within a type context, look specifically within
-    // the type context.
-    // This tries to resolve unqualified operators, like `+`.
-    if (funcName.isOperator() && lookupContext->isTypeContext()) {
-      if (auto tmp = TypeChecker::lookupMember(
-              lookupContext, lookupContext->getSelfTypeInContext(), funcName))
-        results = tmp;
-    }
   }
 
   // Initialize error flags.

stdlib/public/core/SIMDVectorTypes.swift.gyb

@@ -46,7 +46,7 @@ public struct SIMD${n}<Scalar>: SIMD where Scalar: SIMDScalar {
 
   /// Accesses the scalar at the specified position.
   // SWIFT_ENABLE_TENSORFLOW
-  @differentiable(vjp: _vjpSubscript
+  @differentiable(
     where Scalar : EuclideanDifferentiable & BinaryFloatingPoint,
           Scalar.TangentVector : BinaryFloatingPoint)
   public subscript(index: Int) -> Scalar {
@@ -206,9 +206,11 @@ extension SIMD${n} : EuclideanDifferentiable
 extension SIMD${n}
   where Scalar : EuclideanDifferentiable & BinaryFloatingPoint,
         Scalar.TangentVector : BinaryFloatingPoint {
-  @usableFromInline
+  // NOTE(TF-1094): serialized `@derivative` for `.swiftinterface` compilation.
+  @inlinable
+  @derivative(of: subscript(_:))
   internal func _vjpSubscript(index: Int)
-  -> (Scalar, (Scalar.TangentVector) -> TangentVector) {
+  -> (value: Scalar, pullback: (Scalar.TangentVector) -> TangentVector) {
     return (self[index], { v in
       var zeros = Self.zero
       zeros[index] = Scalar(v)

test/AutoDiff/differentiability_witness_swiftinterface.swift

@@ -0,0 +1,36 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -typecheck -emit-module-interface-path - %s -swift-version 5 -enable-library-evolution -module-name Module > %t/Module.swiftinterface
+// RUN: %target-swift-frontend -emit-silgen %t/Module.swiftinterface | %FileCheck  %s --check-prefix=CHECK-SILGEN
+// RUN: not %target-swift-frontend -compile-module-from-interface %t/Module.swiftinterface -o %t/Module.swiftmodule 2>&1 | %FileCheck %s --check-prefix=CHECK-COMPILE
+
+// TF-1094: Derivative registration fails for `.swiftinterface` compilation when
+// original `@differentiable` function is serialized but `@derivative` function
+// is unserialized.
+
+@inlinable // serialized
+@differentiable
+@_silgen_name("foo")
+public func foo(_ x: Float) -> Float {
+  fatalError()
+}
+
+@usableFromInline // not serialized
+@derivative(of: foo)
+@_silgen_name("vjp_foo")
+func vjpFoo(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  return (x, { $0 })
+}
+
+// Missing differentiability witness VJP entry because `func vjpFoo` is bodiless
+// in the `.swiftinterface` file and not lowered to a SIL function.
+
+// CHECK-SILGEN-LABEL: // differentiability witness for foo
+// CHECK-SILGEN-NEXT: sil_differentiability_witness [serialized] [parameters 0] [results 0] @foo : $@convention(thin) (Float) -> Float {
+// CHECK-SILGEN-NEXT: }
+
+// CHECK-SILGEN-LABEL: sil [serialized] [ossa] @foo
+// CHECK-SILGEN-NOT: sil {{.*}} @vjp_foo
+
+// CHECK-COMPILE: Module.swiftinterface:5:2: error: function is not differentiable
+// CHECK-COMPILE: Module.swiftinterface:7:24: note: when differentiating this function definition
+// CHECK-COMPILE: Module.swiftinterface:9:1: note: missing return for differentiation