Skip to content

Commit 66316fd

Browse files
marcrasimarcrasi
authored
add parameter differentiability to SILFunctionType (#21142)
1 parent f47da5a commit 66316fd

File tree

9 files changed

+206
-72
lines changed

9 files changed

+206
-72
lines changed

‎include/swift/AST/Types.h

Lines changed: 46 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3476,13 +3476,33 @@ inline bool isGuaranteedParameter(ParameterConvention conv) {
34763476
llvm_unreachable("bad convention kind");
34773477
}
34783478

3479+
/// SWIFT_ENABLE_TENSORFLOW
3480+
/// Determines whether a differentiable function type is differentiable with
3481+
/// respect to this parameter.
3482+
enum class SILParameterDifferentiability : unsigned {
3483+
/// The function type is differentiable with respect to this parameter, or
3484+
/// differentiability is not applicable because the function is not
3485+
/// differentiable.
3486+
DifferentiableOrNotApplicable,
3487+
3488+
/// The function type is not differentiable with respect to this parameter.
3489+
NotDifferentiable,
3490+
};
3491+
34793492
/// A parameter type and the rules for passing it.
34803493
class SILParameterInfo {
34813494
llvm::PointerIntPair<CanType, 3, ParameterConvention> TypeAndConvention;
3495+
3496+
// SWIFT_ENABLE_TENSORFLOW
3497+
SILParameterDifferentiability Differentiability : 1;
34823498
public:
34833499
SILParameterInfo() = default;//: Ty(), Convention((ParameterConvention)0) {}
3484-
SILParameterInfo(CanType type, ParameterConvention conv)
3485-
: TypeAndConvention(type, conv) {
3500+
// SWIFT_ENABLE_TENSORFLOW
3501+
SILParameterInfo(
3502+
CanType type, ParameterConvention conv,
3503+
SILParameterDifferentiability differentiability =
3504+
SILParameterDifferentiability::DifferentiableOrNotApplicable)
3505+
: TypeAndConvention(type, conv), Differentiability(differentiability) {
34863506
assert(type->isLegalSILType() && "SILParameterInfo has illegal SIL type");
34873507
}
34883508

@@ -3527,6 +3547,16 @@ class SILParameterInfo {
35273547
return isGuaranteedParameter(getConvention());
35283548
}
35293549

3550+
// SWIFT_ENABLE_TENSORFLOW
3551+
SILParameterDifferentiability getDifferentiability() const {
3552+
return Differentiability;
3553+
}
3554+
3555+
SILParameterInfo getWithDifferentiability(
3556+
SILParameterDifferentiability differentiability) const {
3557+
return SILParameterInfo(getType(), getConvention(), differentiability);
3558+
}
3559+
35303560
/// The SIL storage type determines the ABI for arguments based purely on the
35313561
/// formal parameter conventions. The actual SIL type for the argument values
35323562
/// may differ in canonical SIL. In particular, opaque values require indirect
@@ -3537,7 +3567,8 @@ class SILParameterInfo {
35373567

35383568
/// Return a version of this parameter info with the type replaced.
35393569
SILParameterInfo getWithType(CanType type) const {
3540-
return SILParameterInfo(type, getConvention());
3570+
// SWIFT_ENABLE_TENSORFLOW
3571+
return SILParameterInfo(type, getConvention(), getDifferentiability());
35413572
}
35423573

35433574
/// Transform this SILParameterInfo by applying the user-provided
@@ -3553,6 +3584,8 @@ class SILParameterInfo {
35533584
void profile(llvm::FoldingSetNodeID &id) {
35543585
id.AddPointer(getType().getPointer());
35553586
id.AddInteger((unsigned)getConvention());
3587+
// SWIFT_ENABLE_TENSORFLOW
3588+
id.AddInteger((unsigned)getDifferentiability());
35563589
}
35573590

35583591
void dump() const;
@@ -3566,7 +3599,10 @@ class SILParameterInfo {
35663599
}
35673600

35683601
bool operator==(SILParameterInfo rhs) const {
3569-
return getType() == rhs.getType() && getConvention() == rhs.getConvention();
3602+
// SWIFT_ENABLE_TENSORFLOW
3603+
return getType() == rhs.getType() &&
3604+
getConvention() == rhs.getConvention() &&
3605+
getDifferentiability() == rhs.getDifferentiability();
35703606
}
35713607
bool operator!=(SILParameterInfo rhs) const {
35723608
return !(*this == rhs);
@@ -4169,6 +4205,12 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
41694205
unsigned differentiationOrder, AutoDiffAssociatedFunctionKind kind,
41704206
SILModule &module, LookupConformanceFn lookupConformance);
41714207

4208+
/// Returns a bit vector that specifices which parameters you can
4209+
/// differentiate with respect to for this differentiable function type. (e.g.
4210+
/// which parameters are not @nondiff). The function type must be
4211+
/// differentiable.
4212+
SmallBitVector getDifferentiationParameterIndices() const;
4213+
41724214
/// If this is a @convention(witness_method) function with a protocol
41734215
/// constrained self parameter, return the protocol constraint for
41744216
/// the Self type.

‎lib/AST/ASTPrinter.cpp

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4244,6 +4244,14 @@ void SILParameterInfo::print(raw_ostream &OS, const PrintOptions &Opts) const {
42444244
}
42454245
void SILParameterInfo::print(ASTPrinter &Printer,
42464246
const PrintOptions &Opts) const {
4247+
/// SWIFT_ENABLE_TENSORFLOW
4248+
switch (getDifferentiability()) {
4249+
case SILParameterDifferentiability::NotDifferentiable:
4250+
Printer << "@nondiff ";
4251+
break;
4252+
default:
4253+
break;
4254+
}
42474255
Printer << getStringForParameterConvention(getConvention());
42484256
getType().print(Printer, Opts);
42494257
}

‎lib/SIL/SILFunctionType.cpp

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -200,13 +200,42 @@ CanSILFunctionType SILFunctionType::getGradientType(
200200
getWitnessMethodConformanceOrNone());
201201
}
202202

203+
llvm::SmallBitVector
204+
SILFunctionType::getDifferentiationParameterIndices() const {
205+
assert(isDifferentiable());
206+
SmallBitVector result(getNumParameters());
207+
for (auto paramAndIndex : enumerate(getParameters())) {
208+
auto &param = paramAndIndex.value();
209+
unsigned index = paramAndIndex.index();
210+
if (param.getDifferentiability() ==
211+
SILParameterDifferentiability::DifferentiableOrNotApplicable)
212+
result.set(index);
213+
}
214+
return result;
215+
}
216+
203217
CanSILFunctionType SILFunctionType::getWithDifferentiability(
204218
unsigned differentiationOrder,
205219
const SmallBitVector &parameterIndices) {
206220
// FIXME(rxwei): Handle differentiation order.
207-
// FIXME(rxwei): Handle parameter indices.
208-
return getWithExtInfo(
209-
getExtInfo().withDifferentiability(Differentiability::Bidirectional));
221+
222+
SmallVector<SILParameterInfo, 8> newParameters;
223+
for (auto paramAndIndex : enumerate(getParameters())) {
224+
auto &param = paramAndIndex.value();
225+
unsigned index = paramAndIndex.index();
226+
newParameters.push_back(param.getWithDifferentiability(
227+
index < parameterIndices.size() && parameterIndices[index]
228+
? SILParameterDifferentiability::DifferentiableOrNotApplicable
229+
: SILParameterDifferentiability::NotDifferentiable));
230+
}
231+
232+
auto newExtInfo =
233+
getExtInfo().withDifferentiability(Differentiability::Bidirectional);
234+
235+
return get(getGenericSignature(), newExtInfo, getCoroutineKind(),
236+
getCalleeConvention(), newParameters, getYields(), getResults(),
237+
getOptionalErrorResult(), getASTContext(),
238+
getWitnessMethodConformanceOrNone());
210239
}
211240

212241
CanSILFunctionType SILFunctionType::getAutoDiffAssociatedFunctionType(

‎lib/SIL/SILInstructions.cpp

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -692,16 +692,14 @@ getExtracteeType(SILValue function, Extractee extractee,
692692
break;
693693
case Extractee::JVP:
694694
resultFnTy = originalFnTy->getAutoDiffAssociatedFunctionType(
695-
SmallBitVector(originalFnTy->getNumParameters(), true),
696-
/*resultIndex*/ 0, differentiationOrder,
697-
AutoDiffAssociatedFunctionKind::JVP, module,
695+
fnTy->getDifferentiationParameterIndices(), /*resultIndex*/ 0,
696+
differentiationOrder, AutoDiffAssociatedFunctionKind::JVP, module,
698697
LookUpConformanceInModule(module.getSwiftModule()));
699698
break;
700699
case Extractee::VJP:
701700
resultFnTy = originalFnTy->getAutoDiffAssociatedFunctionType(
702-
SmallBitVector(originalFnTy->getNumParameters(), true),
703-
/*resultIndex*/ 0, differentiationOrder,
704-
AutoDiffAssociatedFunctionKind::VJP, module,
701+
fnTy->getDifferentiationParameterIndices(), /*resultIndex*/ 0,
702+
differentiationOrder, AutoDiffAssociatedFunctionKind::VJP, module,
705703
LookUpConformanceInModule(module.getSwiftModule()));
706704
break;
707705
}

‎lib/SIL/SILVerifier.cpp

Lines changed: 27 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1261,30 +1261,33 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
12611261
require(origTy, "The original function must have a function type");
12621262
require(!origTy->isDifferentiable(),
12631263
"The original function must not be @autodiff");
1264-
if (F.getModule().getStage() == SILStage::Canonical ||
1265-
adfi->hasAssociatedFunctions()) {
1266-
for (auto order : range(1, adfi->getDifferentiationOrder() + 1)) {
1267-
auto pair = adfi->getAssociatedFunctionPair(order);
1268-
auto jvpType = pair.first->getType().getAs<SILFunctionType>();
1269-
require(jvpType, "The JVP function must have a function type");
1270-
require(!jvpType->isDifferentiable(),
1271-
"The JVP function must not be @autodiff");
1272-
auto expectedJVPType = origTy->getAutoDiffAssociatedFunctionType(
1273-
adfi->getParameterIndices(), /*resultIndex*/ 0, order,
1274-
AutoDiffAssociatedFunctionKind::JVP, F.getModule(),
1275-
LookUpConformanceInModule(F.getModule().getSwiftModule()));
1276-
require(expectedJVPType == jvpType, "Unexpected JVP function type");
1277-
auto vjpType = pair.second->getType().getAs<SILFunctionType>();
1278-
require(vjpType, "The VJP function must have a function type");
1279-
require(!vjpType->isDifferentiable(),
1280-
"The VJP function must not be @autodiff");
1281-
auto expectedVJPType = origTy->getAutoDiffAssociatedFunctionType(
1282-
adfi->getParameterIndices(), /*resultIndex*/ 0, order,
1283-
AutoDiffAssociatedFunctionKind::VJP, F.getModule(),
1284-
LookUpConformanceInModule(F.getModule().getSwiftModule()));
1285-
require(expectedVJPType == vjpType, "Unexpected VJP function type");
1286-
}
1287-
}
1264+
// TODO: Temporarily disabled this check because witness thunks generate
1265+
// `autodiff_function` instructions without associated functions, and the AD
1266+
// pass does not yet fill in the associated functions.
1267+
// if (F.getModule().getStage() == SILStage::Canonical ||
1268+
// adfi->hasAssociatedFunctions()) {
1269+
// for (auto order : range(1, adfi->getDifferentiationOrder() + 1)) {
1270+
// auto pair = adfi->getAssociatedFunctionPair(order);
1271+
// auto jvpType = pair.first->getType().getAs<SILFunctionType>();
1272+
// require(jvpType, "The JVP function must have a function type");
1273+
// require(!jvpType->isDifferentiable(),
1274+
// "The JVP function must not be @autodiff");
1275+
// auto expectedJVPType = origTy->getAutoDiffAssociatedFunctionType(
1276+
// adfi->getParameterIndices(), /*resultIndex*/ 0, order,
1277+
// AutoDiffAssociatedFunctionKind::JVP, F.getModule(),
1278+
// LookUpConformanceInModule(F.getModule().getSwiftModule()));
1279+
// require(expectedJVPType == jvpType, "Unexpected JVP function type");
1280+
// auto vjpType = pair.second->getType().getAs<SILFunctionType>();
1281+
// require(vjpType, "The VJP function must have a function type");
1282+
// require(!vjpType->isDifferentiable(),
1283+
// "The VJP function must not be @autodiff");
1284+
// auto expectedVJPType = origTy->getAutoDiffAssociatedFunctionType(
1285+
// adfi->getParameterIndices(), /*resultIndex*/ 0, order,
1286+
// AutoDiffAssociatedFunctionKind::VJP, F.getModule(),
1287+
// LookUpConformanceInModule(F.getModule().getSwiftModule()));
1288+
// require(expectedVJPType == vjpType, "Unexpected VJP function type");
1289+
// }
1290+
// }
12881291
}
12891292

12901293
void checkAutoDiffFunctionExtractInst(AutoDiffFunctionExtractInst *adfei) {

‎lib/SILGen/SILGenPoly.cpp

Lines changed: 13 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -3646,41 +3646,26 @@ getWitnessFunctionRef(SILGenFunction &SGF,
36463646
SmallVectorImpl<ManagedValue> &witnessParams,
36473647
SILLocation loc,
36483648
// SWIFT_ENABLE_TENSORFLOW
3649-
AutoDiffAssociatedFunctionIdentifier *autoDiffFuncId) {
3649+
AutoDiffAssociatedFunctionIdentifier *autoDiffFuncId,
3650+
llvm::SmallBitVector loweredIndices) {
36503651
switch (witnessKind) {
36513652
case WitnessDispatchKind::Static:
36523653
// SWIFT_ENABLE_TENSORFLOW
36533654
if (autoDiffFuncId) {
3654-
// Look up the associated autodiff function and emit a ref to that.
3655-
// TODO: We should replace this with `autodiff_function_extract`, like:
3656-
// %adfunc = autodiff_function %orig
3657-
// %assocfunc = autodiff_function_extract [...] %adfunc
3658-
// This is not yet possible because `autodiff_function_extract` does not
3659-
// calculate the right type.
3660-
3661-
auto *diffAttr =
3662-
witness.getDecl()->getAttrs().getAttribute<DifferentiableAttr>();
3663-
assert(diffAttr && autoDiffFuncId->getDifferentiationOrder() == 1
3664-
&& *diffAttr->getCheckedParameterIndices() ==
3665-
*autoDiffFuncId->getParameterIndices()
3666-
&& "TODO: use `autodiff_function_extract` so that we support "
3667-
"non-manually-specified associated functions");
3668-
3669-
FuncDecl *associatedFuncDecl = nullptr;
3655+
auto originalFn = SGF.emitGlobalFunctionRef(loc, witness);
3656+
auto autoDiffFn = SGF.B.createAutoDiffFunction(
3657+
loc, loweredIndices, /*differentiationOrder*/ 1, originalFn);
3658+
AutoDiffFunctionExtractInst::Extractee extractee;
36703659
switch (autoDiffFuncId->getKind()) {
36713660
case AutoDiffAssociatedFunctionKind::JVP:
3672-
associatedFuncDecl = diffAttr->getJVPFunction();
3661+
extractee = AutoDiffFunctionExtractInst::Extractee::JVP;
36733662
break;
36743663
case AutoDiffAssociatedFunctionKind::VJP:
3675-
associatedFuncDecl = diffAttr->getVJPFunction();
3664+
extractee = AutoDiffFunctionExtractInst::Extractee::VJP;
36763665
break;
36773666
}
3678-
3679-
assert(associatedFuncDecl
3680-
&& "TODO: use `autodiff_function_extract` so that we support "
3681-
"non-manually-specified associated functions");
3682-
3683-
return SGF.emitGlobalFunctionRef(loc, SILDeclRef(associatedFuncDecl));
3667+
return SGF.B.createAutoDiffFunctionExtract(
3668+
loc, extractee, /*differentiationOrder*/ 1, autoDiffFn);
36843669
}
36853670

36863671
return SGF.emitGlobalFunctionRef(loc, witness);
@@ -3790,8 +3775,9 @@ void SILGenFunction::emitProtocolWitness(AbstractionPattern reqtOrigTy,
37903775
// the substituted signature of the witness.
37913776
auto origWitnessFTy = getWitnessFunctionType(SGM, witness, witnessKind);
37923777
// SWIFT_ENABLE_TENSORFLOW
3778+
llvm::SmallBitVector loweredIndices;
37933779
if (autoDiffFuncId) {
3794-
auto loweredIndices = autoDiffFuncId->getParameterIndices()->getLowered(
3780+
loweredIndices = autoDiffFuncId->getParameterIndices()->getLowered(
37953781
witnessSubstTy, /*selfUncurried*/ true);
37963782
origWitnessFTy = origWitnessFTy->getAutoDiffAssociatedFunctionType(
37973783
loweredIndices, /*resultIndex*/ 0,
@@ -3819,7 +3805,7 @@ void SILGenFunction::emitProtocolWitness(AbstractionPattern reqtOrigTy,
38193805
witnessKind,
38203806
// SWIFT_ENABLE_TENSORFLOW
38213807
witnessParams, loc,
3822-
autoDiffFuncId);
3808+
autoDiffFuncId, loweredIndices);
38233809

38243810
auto coroutineKind =
38253811
witnessFnRef->getType().castTo<SILFunctionType>()->getCoroutineKind();

‎lib/Sema/TypeCheckType.cpp

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2697,6 +2697,10 @@ SILParameterInfo TypeResolver::resolveSILParameter(
26972697
Type type;
26982698
bool hadError = false;
26992699

2700+
// SWIFT_ENABLE_TENSORFLOW
2701+
auto differentiability =
2702+
SILParameterDifferentiability::DifferentiableOrNotApplicable;
2703+
27002704
if (auto attrRepr = dyn_cast<AttributedTypeRepr>(repr)) {
27012705
auto attrs = attrRepr->getAttrs();
27022706

@@ -2722,6 +2726,15 @@ SILParameterInfo TypeResolver::resolveSILParameter(
27222726
checkFor(TypeAttrKind::TAK_guaranteed,
27232727
ParameterConvention::Direct_Guaranteed);
27242728

2729+
// SWIFT_ENABLE_TENSORFLOW
2730+
if (attrs.has(TAK_nondiff)) {
2731+
// TODO: We could diagnose @nondiff on a non-@autodiff function, but we'd
2732+
// have to pass function differentiability as an argument to
2733+
// `resolveSILParameter`.
2734+
attrs.clearAttribute(TAK_nondiff);
2735+
differentiability = SILParameterDifferentiability::NotDifferentiable;
2736+
}
2737+
27252738
type = resolveAttributedType(attrs, attrRepr->getTypeRepr(), options);
27262739
} else {
27272740
type = resolveType(repr, options);
@@ -2737,7 +2750,9 @@ SILParameterInfo TypeResolver::resolveSILParameter(
27372750
}
27382751

27392752
if (hadError) type = ErrorType::get(Context);
2740-
return SILParameterInfo(type->getCanonicalType(), convention);
2753+
// SWIFT_ENABLE_TENSORFLOW
2754+
return SILParameterInfo(type->getCanonicalType(), convention,
2755+
differentiability);
27412756
}
27422757

27432758
bool TypeResolver::resolveSingleSILResult(TypeRepr *repr,

‎test/AutoDiff/autodiff_sil_function_type_parse.sil

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
// RUN: %target-sil-opt -assume-parsing-unqualified-ownership-sil %s -module-name=autodiff_sil_function_type_parse | %target-sil-opt -assume-parsing-unqualified-ownership-sil -module-name=autodiff_sil_function_type_parse | %FileCheck %s
2+
3+
sil_stage raw
4+
5+
import Swift
6+
7+
sil @examplefunc : $@convention(thin) (Float, Float, Float) -> Float
8+
9+
sil @examplemethod : $@convention(method) (Float, Float, Float) -> Float
10+
11+
// CHECK-LABEL: sil @test
12+
sil @test : $@convention(thin) () -> () {
13+
bb0:
14+
%0 = function_ref @examplefunc : $@convention(thin) (Float, Float, Float) -> Float
15+
16+
%1 = autodiff_function [wrt 0 1 2] [order 1] %0 : $@convention(thin) (Float, Float, Float) -> Float
17+
// CHECK: %2 = autodiff_function_extract [vjp] [order 1] %1 : $@autodiff @convention(thin) (Float, Float, Float) -> Float
18+
%2 = autodiff_function_extract [vjp] [order 1] %1 : $@autodiff @convention(thin) (Float, Float, Float) -> Float
19+
20+
%3 = autodiff_function [wrt 0] [order 1] %0 : $@convention(thin) (Float, Float, Float) -> Float
21+
// CHECK: %4 = autodiff_function_extract [vjp] [order 1] %3 : $@autodiff @convention(thin) (Float, @nondiff Float, @nondiff Float) -> Float
22+
%4 = autodiff_function_extract [vjp] [order 1] %3 : $@autodiff @convention(thin) (Float, @nondiff Float, @nondiff Float) -> Float
23+
24+
%5 = function_ref @examplemethod : $@convention(method) (Float, Float, Float) -> Float
25+
26+
%6 = autodiff_function [wrt 0 1 2] [order 1] %5 : $@convention(method) (Float, Float, Float) -> Float
27+
// CHECK: %7 = autodiff_function_extract [vjp] [order 1] %6 : $@autodiff @convention(method) (Float, Float, Float) -> Float
28+
%7 = autodiff_function_extract [vjp] [order 1] %6 : $@autodiff @convention(method) (Float, Float, Float) -> Float
29+
30+
%8 = autodiff_function [wrt 0] [order 1] %5 : $@convention(method) (Float, Float, Float) -> Float
31+
// CHECK: %9 = autodiff_function_extract [vjp] [order 1] %8 : $@autodiff @convention(method) (Float, @nondiff Float, @nondiff Float) -> Float
32+
%9 = autodiff_function_extract [vjp] [order 1] %8 : $@autodiff @convention(method) (Float, @nondiff Float, @nondiff Float) -> Float
33+
34+
%ret = tuple ()
35+
return %ret : $()
36+
}

0 commit comments

Comments
 (0)