[SR-13166] Default derivative implementations for protocol requirements

[dan-zheng]

Previous ID	SR-13166
Radar	rdar://problem/69987698
Original Reporter	@dan-zheng
Type	New Feature

Additional Detail from JIRA

Votes	1
Component/s	Compiler
Labels	New Feature, AutoDiff
Assignee	@rxwei
Priority	Medium

md5: b0201d55480662f7a4792a7ddd4694d2

Sub-Tasks:

• [SR-13168] SILGen: add SIL default witness table entries for default derivatives of protocol requirements #55610
• [SR-13169] Sema: lift derivative registration restriction for protocol requirements #55611
• [SR-13167] Parse: add parsing/syntax support for @differentiable(default, ...) #55609

Issue Description:

Overview

Default derivative implementations enables protocol requirements (like requirements from AdditiveArithmetic, FloatingPoint, ElementaryFunctions, etc) to be differentiable by default.

See "default derivatives and transposes" from the Differentiable Programming Manifesto for more info.

Example:

// In the standard library:
// public protocol AdditiveArithmetic: Equatable {
//   static func +(lhs: Self, rhs: Self) -> Self
//   ...
// }

extension AdditiveArithmetic where Self: Differentiable {
  @derivative(of: +)
  static func vjpAdd(_ lhs: Self, _ rhs: Self) ->
    (value: Self, pullback: (TangentVector) -> (TangentVector, TangentVector)) {
    return (lhs + rhs, { v in (v, v) })
  }
}

This is important for idiomatic protocol-oriented programming to avoid tons of code duplication. Without this feature, all conforming types of the protocol must:

• Define a @differentiable concrete implementation of the original protocol requirement, if none already exists (because a default implementation was used.)

• Register a derivative function for the concrete implementation.

---

Details

There are two cases to consider:

1. Non-@differentiable protocol requirement.

protocol P {
  func foo(_ x: Float) -> Float
}
extension P {
  @derivative(of: foo)
  func vjpFoo(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
    return (x, { $0 })
  }
}

2. @differentiable protocol requirement.

protocol P {
  @differentiable
  func foo(_ x: Float) -> Float
}
extension P {
  @derivative(of: foo)
  func vjpFoo(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
    return (x, { $0 })
  }
}

Supporting this may require lifting the current restriction that "all implementations of @differentiable protocol requirements must themselves be marked as @differentiable":

protocol AdditiveArithmetic: Equatable {
  static func +(lhs: Self, rhs: Self) -> Self
}

extension AdditiveArithmetic where Self: Differentiable {
  @derivative(of: +)
  static func vjpAdd(_ lhs: Self, _ rhs: Self) ->
    (value: Self, pullback: (TangentVector) -> (TangentVector, TangentVector)) {
    return (lhs + rhs, { v in (v, v) })
  }
}

struct Foo: AdditiveArithmetic {
  static func +(lhs: Self, rhs: Self) -> Self {
    lhs
  }
}

deriv.swift:13:8: error: type 'Foo' does not conform to protocol 'AdditiveArithmetic'
struct Foo: AdditiveArithmetic {
       ^
deriv.swift:14:15: note: candidate is missing attribute '@differentiable(wrt: (lhs, rhs) where Self : Differentiable)'
  static func +(lhs: Self, rhs: Self) -> Self {
              ^

SIL default witness table support may be needed.

[rxwei]

@swift-ci create

[kovdan01]

I started looking at this, and it looks like that the following should be done if we want to allow custom default derivative implementations for protocol requirements.

1. In findAutoDiffOriginalFunctionDecl, collect not only validCandidates, but also smth like validProtocolCandidates. Obviously, isValidOriginalCandidate lambda in typeCheckDerivativeAttr should be changed to allow ProtocolDecl.

2. If there are validCandidates, they takes precedence over protocol candidates. If we have only validProtocolCandidates, we must be inside an extension which has a function with @derivative attr for a protocol requirement. Indeed, if we are inside a struct conforming to a protocol, it must have a requirement implementation itself, so we'll have non-empty validCandidates.

3. In findMinimalDerivativeConfiguration, if nothing is found and the decl context is a nominal type decl, try to look through protocols which this NTD conforms to. If we find a corresponding requirement and findMinimalDerivativeConfiguration is non-null - we should take it. Also, we should emit a differentiability witness accordingly.

A quick attempt to implement that stopped at the following error:

error: expression is not differentiable
note: function call is not differentiable because generic requirements are not met: 'Self : ProtocolRequirementDerivative'

For the code snippet below:

import _Differentiation
protocol ProtocolRequirementDerivative {
  func requirement(_ x: Float) -> Float
}
extension ProtocolRequirementDerivative {
  @derivative(of: requirement)
  func vjpRequirement(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
    return (value: x, pullback: { v in 42 })
  }
}
struct StructProtocolRequirementDerivative : ProtocolRequirementDerivative {
  // This intentionally does not correspond to the custom derivative defined
  // in the extension above in order to ensure that we pick it up.
  func requirement(_ x: Float) -> Float { x }
}
let s = StructProtocolRequirementDerivative()
print(gradient(at: 24, of: s.requirement))

I'll continue investigation later.