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

37
lib/SIL/IR/SILPrinter.cpp
View File

@@ -2442,6 +2442,23 @@ public:
*this << "transpose";
break;
}
*this << "] [";
switch (dwfi->getWitness()->getKind()) {
case DifferentiabilityKind::Forward:
*this << "forward";
break;
case DifferentiabilityKind::Reverse:
*this << "reverse";
break;
case DifferentiabilityKind::Normal:
*this << "normal";
break;
case DifferentiabilityKind::Linear:
*this << "linear";
break;
case DifferentiabilityKind::NonDifferentiable:
llvm_unreachable("Impossible case");
}
*this << "] [parameters";
for (auto i : witness->getParameterIndices()->getIndices())
*this << ' ' << i;
@@ -3342,8 +3359,26 @@ void SILDifferentiabilityWitness::print(llvm::raw_ostream &OS,
// ([serialized])?
if (isSerialized())
OS << "[serialized] ";
// Kind
OS << '[';
switch (getKind()) {
case DifferentiabilityKind::Forward:
OS << "forward";
break;
case DifferentiabilityKind::Reverse:
OS << "reverse";
break;
case DifferentiabilityKind::Normal:
OS << "normal";
break;
case DifferentiabilityKind::Linear:
OS << "linear";
break;
case DifferentiabilityKind::NonDifferentiable:
llvm_unreachable("Impossible case");
}
// [parameters ...]
OS << "[parameters ";
OS << "] [parameters ";
interleave(
getParameterIndices()->getIndices(), [&](unsigned index) { OS << index; },
[&] { OS << ' '; });