mirror of
https://github.com/apple/swift.git
synced 2025-12-21 12:14:44 +01:00
e494df2ee6
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.
81 lines
3.6 KiB
C++
81 lines
3.6 KiB
C++
//===--- SILDifferentiabilityWitness.cpp - Differentiability witnesses ----===//
|
|
//
|
|
// This source file is part of the Swift.org open source project
|
|
//
|
|
// Copyright (c) 2014 - 2019 Apple Inc. and the Swift project authors
|
|
// Licensed under Apache License v2.0 with Runtime Library Exception
|
|
//
|
|
// See https://swift.org/LICENSE.txt for license information
|
|
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#define DEBUG_TYPE "sil-differentiability-witness"
|
|
|
|
#include "swift/AST/ASTMangler.h"
|
|
#include "swift/SIL/SILDifferentiabilityWitness.h"
|
|
#include "swift/SIL/SILModule.h"
|
|
|
|
using namespace swift;
|
|
|
|
SILDifferentiabilityWitness *SILDifferentiabilityWitness::createDeclaration(
|
|
SILModule &module, SILLinkage linkage, SILFunction *originalFunction,
|
|
DifferentiabilityKind kind, IndexSubset *parameterIndices,
|
|
IndexSubset *resultIndices, GenericSignature derivativeGenSig,
|
|
const DeclAttribute *attribute) {
|
|
auto *diffWitness = new (module) SILDifferentiabilityWitness(
|
|
module, linkage, originalFunction, kind, parameterIndices, resultIndices,
|
|
derivativeGenSig, /*jvp*/ nullptr, /*vjp*/ nullptr,
|
|
/*isDeclaration*/ true, /*isSerialized*/ false, attribute);
|
|
// Register the differentiability witness in the module.
|
|
Mangle::ASTMangler mangler;
|
|
auto mangledKey = mangler.mangleSILDifferentiabilityWitness(
|
|
diffWitness->getOriginalFunction()->getName(),
|
|
diffWitness->getKind(), diffWitness->getConfig());
|
|
assert(!module.DifferentiabilityWitnessMap.count(mangledKey) &&
|
|
"Cannot create duplicate differentiability witness in a module");
|
|
module.DifferentiabilityWitnessMap[mangledKey] = diffWitness;
|
|
module.DifferentiabilityWitnessesByFunction[originalFunction->getName()]
|
|
.push_back(diffWitness);
|
|
module.getDifferentiabilityWitnessList().push_back(diffWitness);
|
|
return diffWitness;
|
|
}
|
|
|
|
SILDifferentiabilityWitness *SILDifferentiabilityWitness::createDefinition(
|
|
SILModule &module, SILLinkage linkage, SILFunction *originalFunction,
|
|
DifferentiabilityKind kind, IndexSubset *parameterIndices,
|
|
IndexSubset *resultIndices, GenericSignature derivativeGenSig,
|
|
SILFunction *jvp, SILFunction *vjp, bool isSerialized,
|
|
const DeclAttribute *attribute) {
|
|
auto *diffWitness = new (module) SILDifferentiabilityWitness(
|
|
module, linkage, originalFunction, kind, parameterIndices, resultIndices,
|
|
derivativeGenSig, jvp, vjp, /*isDeclaration*/ false, isSerialized,
|
|
attribute);
|
|
// Register the differentiability witness in the module.
|
|
Mangle::ASTMangler mangler;
|
|
auto mangledKey = mangler.mangleSILDifferentiabilityWitness(
|
|
diffWitness->getOriginalFunction()->getName(),
|
|
diffWitness->getKind(), diffWitness->getConfig());
|
|
assert(!module.DifferentiabilityWitnessMap.count(mangledKey) &&
|
|
"Cannot create duplicate differentiability witness in a module");
|
|
module.DifferentiabilityWitnessMap[mangledKey] = diffWitness;
|
|
module.DifferentiabilityWitnessesByFunction[originalFunction->getName()]
|
|
.push_back(diffWitness);
|
|
module.getDifferentiabilityWitnessList().push_back(diffWitness);
|
|
return diffWitness;
|
|
}
|
|
|
|
void SILDifferentiabilityWitness::convertToDefinition(SILFunction *jvp,
|
|
SILFunction *vjp,
|
|
bool isSerialized) {
|
|
assert(IsDeclaration);
|
|
IsDeclaration = false;
|
|
JVP = jvp;
|
|
VJP = vjp;
|
|
IsSerialized = isSerialized;
|
|
}
|
|
|
|
SILDifferentiabilityWitnessKey SILDifferentiabilityWitness::getKey() const {
|
|
return {getOriginalFunction()->getName(), getKind(), getConfig()};
|
|
}
|