averello/swift-mirror
1
0
Fork 0
mirror of https://github.com/apple/swift.git synced 2025-12-14 20:36:38 +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

42
lib/Demangling/NodePrinter.cpp
View File

@@ -570,6 +570,7 @@ private:
case Node::Kind::AutoDiffSelfReorderingReabstractionThunk:
case Node::Kind::AutoDiffSubsetParametersThunk:
case Node::Kind::AutoDiffFunctionKind:
case Node::Kind::DifferentiabilityWitness:
case Node::Kind::IndexSubset:
return false;
}
@@ -1850,6 +1851,47 @@ NodePointer NodePrinter::print(NodePointer Node, bool asPrefixContext) {
}
return nullptr;
}
case Node::Kind::DifferentiabilityWitness: {
auto kindNodeIndex = Node->getNumChildren() - (
Node->getLastChild()->getKind() == Node::Kind::DependentGenericSignature
? 4 : 3);
auto kind =
(MangledDifferentiabilityKind)Node->getChild(kindNodeIndex)->getIndex();
switch (kind) {
case MangledDifferentiabilityKind::Forward:
Printer << "forward-mode";
break;
case MangledDifferentiabilityKind::Reverse:
Printer << "reverse-mode";
break;
case MangledDifferentiabilityKind::Normal:
Printer << "normal";
break;
case MangledDifferentiabilityKind::Linear:
Printer << "linear";
break;
case MangledDifferentiabilityKind::NonDifferentiable:
assert(false && "Impossible case");
}
Printer << " differentiability witness for ";
unsigned idx = 0;
for (auto numChildren = Node->getNumChildren();
idx < numChildren &&
Node->getChild(idx)->getKind() != Node::Kind::Index; ++idx)
print(Node->getChild(idx));
++idx; // kind (handled earlier)
Printer << " with respect to parameters ";
print(Node->getChild(idx++)); // parameter indices
Printer << " and results ";
print(Node->getChild(idx++));
if (idx < Node->getNumChildren()) {
auto *genSig = Node->getChild(idx);
assert(genSig->getKind() == Node::Kind::DependentGenericSignature);
Printer << " with ";
print(genSig);
}
return nullptr;
}
case Node::Kind::IndexSubset: {
Printer << '{';
auto text = Node->getText();