[AutoDiff upstream] Add @differentiable declaration attribute type-checking.

[dan-zheng]

The @differentiable attribute marks a function as differentiable.

Example:

@differentiable(wrt: x, jvp: derivativeFoo where T: Differentiable)
func id<T>(_ x: T) -> T { x }

The @differentiable attribute has an optional wrt: clause specifying the
parameters that are differentiated "with respect to", i.e. the differentiability
parameters. The differentiability parameters must conform to the
Differentiable protocol.

If the wrt: clause is unspecified, the differentiability parameters are
currently inferred to be all parameters that conform to Differentiable.

The @differentiable attribute also has optional jvp: and vjp: labels
for registering derivative functions. These labels are deprecated in favor of
the @derivative attribute and will be removed soon.

The @differentiable attribute also has an optional where clause, specifying
extra differentiability requirements for generic functions.

The @differentiable attribute is gated by the
-enable-experimental-differentiable-programming flag.

---

Type-checking rules (summary):

• @differentiable attribute must be declared on a function-like "original"
  declaration: func, init, subscript, var (computed properties only).
• Parsed differentiability parameters must be valid (if they exist).
• Parsed where clause must be valid (if it exists).
• Differentiability parameters must all conform to Differentiable.
• Original result must all conform to Differentiable.
• If JVP/VJP functions are specified, they must match the expected type.
  • @differentiable(jvp:vjp:) for derivative registration is deprecated in
    favor of @derivative attribute, and will be removed soon.
• Duplicate @differentiable attributes with the same differentiability
  parameters are invalid.
• For protocol requirements and class members with @differentiable attribute,
  conforming types and subclasses must have the same @differentiable attribute
  (or one with a superset of differentiability parameter indices) on
  implementing/overriding declarations.

These rules are consistent with the differentiable programming manifesto.
The main proposed rules are implemented.

Code changes:

• Add DifferentiableAttributeTypeCheckRequest.
  • Currently, the request returns differentiability parameter indices,
    while also resolving JVPFunction, VJPFunction, and
    DerivativeGenericSignature and mutating them in-place in
    DifferentiableAttr. This works fine for now.
• Add "is type-checked" bit to DifferentiableAttr.
  • Alternatively, I tried changing CacheKind::SeparatelyCached to
    CacheKind::Cached, but it did not seem to work: @differentiable
    attributes in non-primary-files were left unchecked.

Upstream disorganized tests as-is from tensorflow branch.

Resolves TF-828: upstream @differentiable attribute type-checking.

[dan-zheng]

Let's see if tests pass.
@swift-ci Please smoke test

[dan-zheng]

Note: I anticipate that some code owners might care strongly about preserving return finalize(anyRenaming, optionalAdjustments) here. It doesn't matter too much to me, so I upstreamed code from tensorflow branch as-is.

Aside: this logic for diagnosing unmet @differentiable attribute requirement doesn't work when associated type inference is involved in witness checking. TF-1014 tracks that issue.

[dan-zheng]

Note: I anticipate that some code owners might care strongly about preserving return finalize(anyRenaming, optionalAdjustments) here.

Slightly mitigating the above point: I refactored swift::matchWitness logic for checking @differentiable attributes into a matchWitnessDifferentiableAttr helper. So the code changes to swift::matchWitness are minimal.

[rxwei]

LGTM. Please clarify on the PR title that this is "Adding @differentiable declaration attribute type-checking" to avoid any confusion with the @differentiable function type attribute.

[dan-zheng]

Currently, DifferentiableAttributeTypeCheckRequest takes a DifferentiableAttr *and returns an IndexSubset * representing parameter indices. It also resolves two FuncDecl * (JVPFunction and VJPFunction) and a GenericSignature (DerivativeGenericSignature), mutating them in-place in the DifferentiableAttr *.

This works fine for now. I thought of two ways to refactor the request to avoid mutation:

• Make DifferentiableAttributeTypeCheckRequest return a tuple/struct of the resolved components.
  • This caused many request cycles, since DifferentiableAttr::get{JVPFunction,VJPFunction,DerivativeGenericSignature} now all trigger the request. I haven't debugged further.
• Make an individual request for resolving each @differentiable attribute component: original AbstractFuncDecl(s), parameter indices, JVP/VJP FuncDecl (if specified), derivative GenericSignature (if specified).
  • I haven't tried this. It seems like a fair amount of work, and the upside is unclear.

[dan-zheng]

History: #28017 moved @differentiable attribute type-checking logic from AttributeChecker::visitDifferentiableAttr to DifferentiableAttributeTypeCheckRequest::evaluate.

The current request approach (return IndexSubset * but mutate other components in-place) is the first one that worked.

[dan-zheng]

Using CacheKind::Cached seems preferable so that DifferentiableAttr is truly stateless, adhering to the request evaluator caching vision.

I tried changing CacheKind::SeparatelyCached to CacheKind::Cached (some time ago on tensorflow branch), but it did not seem to work: @differentiable attributes in non-primary-files were left unchecked.

I'm not sure why CacheKind::Cached did not work. I did notice that quite a few other requests (e.g. InterfaceTypeRequest) used CacheKind::SeparatelyCached for some reason - perhaps that's just for request cycle breaking?

[dan-zheng]

@CodaFi: would you like to review this PR, or suggest a reviewer? I'm asking you since you seem like a request evaluator owner.

This patch adds a DifferentiableAttributeTypeCheckRequest for type-checking @differentiable attributes. A request is necessary (instead of just using AttributeChecker) to type-check @differentiable attributes in non-primary files.

I created two discussion threads regarding possible improvements:

• Improving DifferentiableAttributeTypeCheckRequest, which currently mutates DifferentiableAttr.
• Using CacheKind::Cached for DifferentiableAttributeTypeCheckRequest doesn't work.

---

Any improvement suggestions would be appreciated! If you'd like me to split this patch into two smaller ones (idea), I'm happy to do so.

Otherwise, I'll merge this patch within 2 days so that it's in a development snapshot release before our master -> tensorflow branch merge rotation next week.

[dan-zheng]

@swift-ci Please test

[dan-zheng]

Merging now, happy to address further feedback later!