averello/swift-mirror
1
0
Fork 0
mirror of https://github.com/apple/swift.git synced 2025-12-21 12:14:44 +01:00
Code Issues Packages Projects Releases Wiki Activity

[AutoDiff] Add differentiability kind to differentiability witnesses and mangle them.

Browse Source 
Differentiability witnesses are now keyed by the original function name, the differentiability kind, and the autodiff config.

Updated SIL syntax:
```
differentiability-kind ::= 'forward' | 'reverse' | 'normal' | 'linear'
sil-differentiability-witness ::=
    'sil_differentiability_witness'
    sil-linkage?
    '[' differentiability-kind ']'
    '[' 'parameters' sil-differentiability-witness-function-index-list ']'
    '[' 'results' sil-differentiability-witness-function-index-list ']'
    generic-parameter-clause?
    sil-function-name ':' sil-type
    sil-differentiability-witness-body?
sil-instruction ::=
    'differentiability_witness_function'
    '[' sil-differentiability-witness-function-kind ']'
    '[' differentiability-kind ']'
    '[' 'parameters' sil-differentiability-witness-function-index-list ']'
    '[' 'results' sil-differentiability-witness-function-index-list ']'
    generic-parameter-clause?
    sil-function-name ':' sil-type
```
```console
sil_differentiability_witness [reverse] [parameters 0 1] [results 0] <T where T: Differentiable> @foo : <T> $(T) -> T
differentiability_witness_function [vjp] [reverse] [parameters 0] [results 0] <T where T: Differentiable> @foo : $(T) -> T
```

New mangling:
```swift
  global ::= global generic-signature? 'WJ' DIFFERENTIABILITY-KIND INDEX-SUBSET 'p' INDEX-SUBSET 'r' // differentiability witness
```
```console
$s13test_mangling3fooyS2f_S2ftFWJrSpSr ---> reverse differentiability witness for test_mangling.foo(Swift.Float, Swift.Float, Swift.Float) -> Swift.Float with respect to parameters {0} and results {0}
```

Resolves rdar://74380324.
This commit is contained in:
Richard Wei
2021-02-17 03:15:52 -05:00
parent 37d52fbb16
commit e494df2ee6
53 changed files with 618 additions and 290 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
lib/AST/AutoDiff.cpp
View File

@@ -559,13 +559,34 @@ TangentPropertyInfo TangentStoredPropertyRequest::evaluate(
return TangentPropertyInfo(tanField);
}
void SILDifferentiabilityWitnessKey::print(llvm::raw_ostream &s) const {
s << "(original=@" << originalFunctionName << " kind=";
switch (kind) {
case DifferentiabilityKind::NonDifferentiable:
s << "nondifferentiable";
break;
case DifferentiabilityKind::Forward:
s << "forward";
break;
case DifferentiabilityKind::Reverse:
s << "reverse";
break;
case DifferentiabilityKind::Normal:
s << "normal";
break;
case DifferentiabilityKind::Linear:
s << "linear";
break;
}
s << " config=" << config << ')';
}
Demangle::AutoDiffFunctionKind Demangle::getAutoDiffFunctionKind(
AutoDiffDerivativeFunctionKind kind) {
switch (kind) {
case AutoDiffDerivativeFunctionKind::JVP:
return Demangle::AutoDiffFunctionKind::JVP;
case AutoDiffDerivativeFunctionKind::VJP:
return Demangle::AutoDiffFunctionKind::VJP;
case AutoDiffDerivativeFunctionKind::VJP: return Demangle::AutoDiffFunctionKind::VJP;
}
}
@@ -578,3 +599,18 @@ Demangle::AutoDiffFunctionKind Demangle::getAutoDiffFunctionKind(
return Demangle::AutoDiffFunctionKind::Pullback;
}
}
Demangle::MangledDifferentiabilityKind
Demangle::getMangledDifferentiabilityKind(DifferentiabilityKind kind) {
using namespace Demangle;
switch (kind) {
#define SIMPLE_CASE(CASE) \
case DifferentiabilityKind::CASE: return MangledDifferentiabilityKind::CASE;
SIMPLE_CASE(NonDifferentiable)
SIMPLE_CASE(Forward)
SIMPLE_CASE(Reverse)
SIMPLE_CASE(Normal)
SIMPLE_CASE(Linear)
#undef SIMPLE_CASE
}
}