From 85c828e704c6e5490381baf2d43bcbc61178fb08 Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Fri, 29 Nov 2019 00:44:34 -0500 Subject: [PATCH 1/6] [AutoDiff] Differentiate partially-applied reabstraction thunks. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Support differentiation of partially-applied reabstraction thunks: ``` %generic_fn = function_ref @generic : $<τ_0_0> (@in_guaranteed τ_0_0) -> @out τ_0_0 %specialized_fn = partial_apply %generic_fn() %reabs_thunk = function_ref @reabs_thunk : $(Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float %thunked_fn = partial_apply %reabs_thunk(%specialized_fn) %diff_fn = differentiable_function %thunked_fn ``` The enables differentiation of direct references to generic functions. - SIL: add special reabstraction thunk case to `SILFunctionType::getAutoDiffDerivativeFunctionType`. - Reabstraction thunks have a function as the last argument. - Reabstraction thunk JVPs/VJPs have a `@differentiable` function as the last argument. - Differentiation transform: handle partially-applied reabstraction thunks in `reapplyFunctionConversion`. Resolves TF-201. --- include/swift/AST/Types.h | 3 +- lib/SIL/SILFunctionType.cpp | 26 +++-- lib/SIL/SILInstructions.cpp | 7 +- lib/SIL/SILVerifier.cpp | 10 +- .../Mandatory/Differentiation.cpp | 99 ++++++++++++------- test/AutoDiff/forward_mode_runtime.swift | 13 +-- test/AutoDiff/generics.swift | 43 +++++--- test/AutoDiff/simple_math.swift | 15 ++- 8 files changed, 137 insertions(+), 79 deletions(-) diff --git a/include/swift/AST/Types.h b/include/swift/AST/Types.h index a33671e37a642..5028a725c4231 100644 --- a/include/swift/AST/Types.h +++ b/include/swift/AST/Types.h @@ -4308,7 +4308,8 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode, IndexSubset *parameterIndices, unsigned resultIndex, AutoDiffDerivativeFunctionKind kind, Lowering::TypeConverter &TC, LookupConformanceFn lookupConformance, - CanGenericSignature derivativeFunctionGenericSignature = nullptr); + CanGenericSignature derivativeFunctionGenericSignature = nullptr, + bool isReabstractionThunk = false); /// Returns the type of the transpose function. CanSILFunctionType getAutoDiffTransposeFunctionType( diff --git a/lib/SIL/SILFunctionType.cpp b/lib/SIL/SILFunctionType.cpp index 02bd60fad2a9b..982f0a7bdb29f 100644 --- a/lib/SIL/SILFunctionType.cpp +++ b/lib/SIL/SILFunctionType.cpp @@ -216,7 +216,7 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType( IndexSubset *parameterIndices, unsigned resultIndex, AutoDiffDerivativeFunctionKind kind, TypeConverter &TC, LookupConformanceFn lookupConformance, - CanGenericSignature derivativeFnGenSig) { + CanGenericSignature derivativeFnGenSig, bool isReabstractionThunk) { // JVP: (T...) -> ((R...), // (T.TangentVector...) -> (R.TangentVector...)) // VJP: (T...) -> ((R...), @@ -341,6 +341,20 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType( } } + SmallVector newParameters; + newParameters.reserve(getNumParameters()); + newParameters.append(getParameters().begin(), getParameters().end()); + // Reabstraction thunks have a function-typed argument (the function to + // reabstract) as their last argument. Reabstraction thunk JVPs/VJPs have a + // `@differentiable` function-typed last argument instead. + if (isReabstractionThunk) { + auto fnParam = newParameters.back(); + auto fnParamType = dyn_cast(fnParam.getInterfaceType()); + assert(fnParamType); + auto diffFnType = fnParamType->getWithDifferentiability( + DifferentiabilityKind::Normal, parameterIndices); + newParameters.back() = fnParam.getWithInterfaceType(diffFnType); + } SmallVector newResults; newResults.reserve(getNumResults() + 1); for (auto &result : getResults()) { @@ -350,11 +364,11 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType( } newResults.push_back({closureType->getCanonicalType(derivativeFnGenSig), ResultConvention::Owned}); - return SILFunctionType::get(derivativeFnGenSig, getExtInfo(), - getCoroutineKind(), getCalleeConvention(), - getParameters(), getYields(), newResults, - getOptionalErrorResult(), getSubstitutions(), isGenericSignatureImplied(), ctx, - getWitnessMethodConformanceOrInvalid()); + return SILFunctionType::get( + derivativeFnGenSig, getExtInfo(), getCoroutineKind(), + getCalleeConvention(), newParameters, getYields(), newResults, + getOptionalErrorResult(), getSubstitutions(), isGenericSignatureImplied(), + ctx, getWitnessMethodConformanceOrInvalid()); } CanSILFunctionType SILFunctionType::getAutoDiffTransposeFunctionType( diff --git a/lib/SIL/SILInstructions.cpp b/lib/SIL/SILInstructions.cpp index 3441ee91cacc3..9c0b800c76cbb 100644 --- a/lib/SIL/SILInstructions.cpp +++ b/lib/SIL/SILInstructions.cpp @@ -768,9 +768,12 @@ SILType DifferentiabilityWitnessFunctionInst::getDifferentiabilityWitnessType( auto *parameterIndices = witness->getParameterIndices(); auto *resultIndices = witness->getResultIndices(); if (auto derivativeKind = witnessKind.getAsDerivativeFunctionKind()) { + bool isReabstractionThunk = + witness->getOriginalFunction()->isThunk() == IsReabstractionThunk; auto diffFnTy = fnTy->getAutoDiffDerivativeFunctionType( - parameterIndices, *resultIndices->begin(), *derivativeKind, module.Types, - LookUpConformanceInModule(module.getSwiftModule()), witnessCanGenSig); + parameterIndices, *resultIndices->begin(), *derivativeKind, + module.Types, LookUpConformanceInModule(module.getSwiftModule()), + witnessCanGenSig, isReabstractionThunk); return SILType::getPrimitiveObjectType(diffFnTy); } assert(witnessKind == diff --git a/lib/SIL/SILVerifier.cpp b/lib/SIL/SILVerifier.cpp index ecfeb533ff6ce..e582cf242cc9e 100644 --- a/lib/SIL/SILVerifier.cpp +++ b/lib/SIL/SILVerifier.cpp @@ -5412,7 +5412,9 @@ void SILDifferentiabilityWitness::verify(const SILModule &M) const { // parameter/result conventions in lowered SIL. if (M.getStage() == SILStage::Lowered) return; - auto origFnType = getOriginalFunction()->getLoweredFunctionType(); + auto *origFn = getOriginalFunction(); + auto origFnType = origFn->getLoweredFunctionType(); + bool origIsReabstractionThunk = origFn->isThunk() == IsReabstractionThunk; CanGenericSignature derivativeCanGenSig; if (auto derivativeGenSig = getDerivativeGenericSignature()) derivativeCanGenSig = derivativeGenSig->getCanonicalSignature(); @@ -5438,7 +5440,8 @@ void SILDifferentiabilityWitness::verify(const SILModule &M) const { auto expectedJVPType = origFnType->getAutoDiffDerivativeFunctionType( getParameterIndices(), /*resultIndex*/ *getResultIndices()->begin(), AutoDiffDerivativeFunctionKind::JVP, M.Types, - LookUpConformanceInModule(M.getSwiftModule()), derivativeCanGenSig); + LookUpConformanceInModule(M.getSwiftModule()), derivativeCanGenSig, + origIsReabstractionThunk); requireSameType(jvp->getLoweredFunctionType(), expectedJVPType, "JVP type does not match expected JVP type"); } @@ -5448,7 +5451,8 @@ void SILDifferentiabilityWitness::verify(const SILModule &M) const { auto expectedVJPType = origFnType->getAutoDiffDerivativeFunctionType( getParameterIndices(), /*resultIndex*/ *getResultIndices()->begin(), AutoDiffDerivativeFunctionKind::VJP, M.Types, - LookUpConformanceInModule(M.getSwiftModule()), derivativeCanGenSig); + LookUpConformanceInModule(M.getSwiftModule()), derivativeCanGenSig, + origIsReabstractionThunk); requireSameType(vjp->getLoweredFunctionType(), expectedVJPType, "VJP type does not match expected VJP type"); } diff --git a/lib/SILOptimizer/Mandatory/Differentiation.cpp b/lib/SILOptimizer/Mandatory/Differentiation.cpp index 2f61489ae4938..7f01d9c0d8a64 100644 --- a/lib/SILOptimizer/Mandatory/Differentiation.cpp +++ b/lib/SILOptimizer/Mandatory/Differentiation.cpp @@ -236,12 +236,16 @@ template static Inst *peerThroughFunctionConversions(SILValue value) { if (auto *inst = dyn_cast(value)) return inst; - if (auto *thinToThick = dyn_cast(value)) - return peerThroughFunctionConversions(thinToThick->getOperand()); - if (auto *convertFn = dyn_cast(value)) - return peerThroughFunctionConversions(convertFn->getOperand()); - if (auto *partialApply = dyn_cast(value)) - return peerThroughFunctionConversions(partialApply->getCallee()); + if (auto *cvi = dyn_cast(value)) + return peerThroughFunctionConversions(cvi->getOperand()); + if (auto *bbi = dyn_cast(value)) + return peerThroughFunctionConversions(bbi->getOperand()); + if (auto *tttfi = dyn_cast(value)) + return peerThroughFunctionConversions(tttfi->getOperand()); + if (auto *cfi = dyn_cast(value)) + return peerThroughFunctionConversions(cfi->getOperand()); + if (auto *pai = dyn_cast(value)) + return peerThroughFunctionConversions(pai->getCallee()); return nullptr; } @@ -1524,11 +1528,10 @@ static bool diagnoseUnsupportedControlFlow(ADContext &context, } /// Check whether the given requirements are satisfied, with the given -/// derivative generic signature (containing requirements), original function, -/// and substitution map. Returns true if error is emitted. +/// derivative generic signature (containing requirements), and substitution +/// map. Returns true if error is emitted. static bool diagnoseUnsatisfiedRequirements(ADContext &context, GenericSignature derivativeGenSig, - SILFunction *original, SubstitutionMap substMap, DifferentiationInvoker invoker, SourceLoc loc) { @@ -1707,20 +1710,30 @@ static void copyParameterArgumentsForApply( /// recursively converts the new function just like how the old function is /// converted. If the new function's generic signature is specified, it is used /// to create substitution maps for reapplied `partial_apply` instructions. -static SILValue -reapplyFunctionConversion( - SILValue newFunc, SILValue oldFunc, SILValue oldConvertedFunc, - SILBuilder &builder, SILLocation loc, +static SILValue reapplyFunctionConversion( + ADContext &context, SILValue newFunc, SILValue oldFunc, + SILValue oldConvertedFunc, SILBuilder &builder, SILLocation loc, SmallVectorImpl &newBuffersToDealloc, - GenericSignature newFuncGenSig = GenericSignature()) { + GenericSignature newFuncGenSig = GenericSignature(), + IndexSubset *reabstractionThunkIndices = nullptr) { // If the old func is the new func, then there's no conversion. if (oldFunc == oldConvertedFunc) return newFunc; // Handle a few instruction cases. + // copy_value + if (auto *cvi = dyn_cast(oldConvertedFunc)) { + auto innerNewFunc = reapplyFunctionConversion( + context, newFunc, oldFunc, cvi->getOperand(), builder, loc, + newBuffersToDealloc, newFuncGenSig); + // Note: no `copy_value` is needed for the re-converted function because the + // caller of `reapplyFunctionConversion` should consume the re-converted + // function. + return innerNewFunc; + } // thin_to_thick_function if (auto *tttfi = dyn_cast(oldConvertedFunc)) { auto innerNewFunc = reapplyFunctionConversion( - newFunc, oldFunc, tttfi->getOperand(), builder, loc, + context, newFunc, oldFunc, tttfi->getOperand(), builder, loc, newBuffersToDealloc, newFuncGenSig); auto operandFnTy = innerNewFunc->getType().castTo(); auto thickTy = operandFnTy->getWithRepresentation( @@ -1735,9 +1748,18 @@ reapplyFunctionConversion( SmallVector newArgsToDestroy; copyParameterArgumentsForApply(pai, newArgs, newArgsToDestroy, newBuffersToDealloc); + if (reabstractionThunkIndices) { + assert(newArgs.size() == 1 && + "Expected reabstraction thunk to be partially applied with only " + "one argument"); + auto *dfi = context.createDifferentiableFunction( + builder, loc, reabstractionThunkIndices, newArgs.back()); + context.getDifferentiableFunctionInsts().push_back(dfi); + newArgs.back() = dfi; + } auto innerNewFunc = reapplyFunctionConversion( - newFunc, oldFunc, pai->getCallee(), builder, loc, newBuffersToDealloc, - newFuncGenSig); + context, newFunc, oldFunc, pai->getCallee(), builder, loc, + newBuffersToDealloc, newFuncGenSig); // If new function's generic signature is specified, use it to create // substitution map for reapplied `partial_apply` instruction. auto substMap = !newFuncGenSig @@ -1900,14 +1922,15 @@ emitDerivativeFunctionReference( // If the original callee is a `partial_apply` or `apply` instruction, use // its substitution map instead. auto substMap = original->getFunction()->getForwardingSubstitutionMap(); - if (auto *pai = dyn_cast(original)) { + if (auto *pai = + peerThroughFunctionConversions(original)) { substMap = pai->getSubstitutionMap(); - } else if (auto *ai = dyn_cast(original)) { + } else if (auto *ai = peerThroughFunctionConversions(original)) { substMap = ai->getSubstitutionMap(); } if (diagnoseUnsatisfiedRequirements( - context, minimalWitness->getDerivativeGenericSignature(), - originalFn, substMap, invoker, original.getLoc().getSourceLoc())) + context, minimalWitness->getDerivativeGenericSignature(), substMap, + invoker, original.getLoc().getSourceLoc())) return None; DifferentiabilityWitnessFunctionKind witnessKind; switch (kind) { @@ -1920,19 +1943,17 @@ emitDerivativeFunctionReference( } auto *derivativeFnRef = builder.createDifferentiabilityWitnessFunction( loc, witnessKind, minimalWitness); - // FIXME(TF-201): Handle direct differentiation of reabstraction thunks. - // Tentative solution: clone a new reabstraction thunk where function - // argument has a `@differentiable` function type. - if (originalFn->isThunk() == IsReabstractionThunk) { - // Handle here. - } + IndexSubset *reabstractionThunkIndices = nullptr; + if (originalFn->isThunk() == IsReabstractionThunk) + reabstractionThunkIndices = desiredIndices.parameters; auto convertedRef = - reapplyFunctionConversion(derivativeFnRef, originalFRI, original, - builder, loc, newBuffersToDealloc, + reapplyFunctionConversion(context, derivativeFnRef, originalFRI, + original, builder, loc, newBuffersToDealloc, derivativeFnRef->getType() .getASTType() ->castTo() - ->getSubstGenericSignature()); + ->getSubstGenericSignature(), + reabstractionThunkIndices); return std::make_pair( convertedRef, SILAutoDiffIndices(desiredIndices.source, @@ -1978,8 +1999,8 @@ emitDerivativeFunctionReference( requirementDeclRef.asAutoDiffDerivativeFunction(autoDiffFuncId), SILType::getPrimitiveObjectType(assocType)); auto convertedRef = - reapplyFunctionConversion(ref, witnessMethod, original, builder, loc, - newBuffersToDealloc); + reapplyFunctionConversion(context, ref, witnessMethod, original, + builder, loc, newBuffersToDealloc); return std::make_pair(convertedRef, minimalIndices); } @@ -2023,8 +2044,8 @@ emitDerivativeFunctionReference( methodDeclRef.asAutoDiffDerivativeFunction(autoDiffFuncId), SILType::getPrimitiveObjectType(assocType)); auto convertedRef = - reapplyFunctionConversion(ref, classMethodInst, original, builder, loc, - newBuffersToDealloc); + reapplyFunctionConversion(context, ref, classMethodInst, original, + builder, loc, newBuffersToDealloc); return std::make_pair(convertedRef, minimalIndices); } @@ -7159,7 +7180,8 @@ static SILFunction *createEmptyVJP(ADContext &context, SILFunction *original, auto vjpType = originalTy->getAutoDiffDerivativeFunctionType( indices.parameters, indices.source, AutoDiffDerivativeFunctionKind::VJP, module.Types, LookUpConformanceInModule(module.getSwiftModule()), - vjpGenericSig); + vjpGenericSig, + /*isReabstractionThunk*/ original->isThunk() == IsReabstractionThunk); SILOptFunctionBuilder fb(context.getTransform()); auto *vjp = fb.createFunction(linkage, vjpName, vjpType, vjpGenericEnv, @@ -7204,9 +7226,10 @@ static SILFunction *createEmptyJVP(ADContext &context, SILFunction *original, ? jvpGenericSig->getGenericEnvironment() : nullptr; auto jvpType = originalTy->getAutoDiffDerivativeFunctionType( - indices.parameters, indices.source, - AutoDiffDerivativeFunctionKind::JVP, module.Types, - LookUpConformanceInModule(module.getSwiftModule()), jvpGenericSig); + indices.parameters, indices.source, AutoDiffDerivativeFunctionKind::JVP, + module.Types, LookUpConformanceInModule(module.getSwiftModule()), + jvpGenericSig, + /*isReabstractionThunk*/ original->isThunk() == IsReabstractionThunk); SILOptFunctionBuilder fb(context.getTransform()); auto *jvp = fb.createFunction(linkage, jvpName, jvpType, jvpGenericEnv, diff --git a/test/AutoDiff/forward_mode_runtime.swift b/test/AutoDiff/forward_mode_runtime.swift index 46a4b94cbb41a..3f05c9aee2ec1 100644 --- a/test/AutoDiff/forward_mode_runtime.swift +++ b/test/AutoDiff/forward_mode_runtime.swift @@ -318,16 +318,13 @@ ForwardModeTests.test("TupleMutation") { } expectEqual(405, derivative(at: 3, in: nested)) - // FIXME(TF-201): Update after reabstraction thunks can be directly differentiated. - /* - func generic(_ x: T) -> T { - var tuple = (x, x) - tuple.0 += x - tuple.1 += x - return tuple.0 + tuple.0 + func generic(_ x: T) -> T { + var tuple = (T.zero, T.zero) + tuple.0 = x + tuple.1 = x + return tuple.0 } expectEqual(1, derivative(at: 3.0, in: generic)) - */ } // Tests TF-321. diff --git a/test/AutoDiff/generics.swift b/test/AutoDiff/generics.swift index b1605635d5bed..50560144d907a 100644 --- a/test/AutoDiff/generics.swift +++ b/test/AutoDiff/generics.swift @@ -16,31 +16,50 @@ _ = gradient(at: Float(1), in: { x in identity(x) }) // CHECK-SIL-NEXT: [[EMIT_ZERO_INDIRECT:%.*]] = apply [[ZERO_WITNESS]]<τ_0_0.TangentVector>([[ORIG_COTAN]], [[ORIG_COTAN_METATYPE]]) // CHECK-SIL: } -struct Tensor : Differentiable { - // NOTE: `value` must have type with known size (e.g. `Float`, not `Scalar`) - // until differentiation has indirect passing support. - var value: Float - init(_ value: Float) { self.value = value } +// Test TF-201: differentiate direct references to generic function: +// partially-applied reabstraction thunk. + +_ = gradient(at: Float(1), in: identity) + +protocol DifferentiableAdditiveArithmetic: Differentiable & AdditiveArithmetic { + @differentiable + static func + (lhs: Self, rhs: Self) -> Self } +extension Float: DifferentiableAdditiveArithmetic {} +func generic(_ x: T) -> T { + x + x + x +} +_ = gradient(at: Float(10), in: generic) -func generic(_ x: Tensor) -> Float { - return x.value + x.value +struct Wrapper : Differentiable { + var value: Scalar + init(_ value: Scalar) { self.value = value } +} +func generic(_ x: Wrapper) -> T { + return x.value +} +_ = gradient(at: Wrapper(1), in: generic) + +func generic2(_ x: T, _ y: Float, _ z: U) -> T { + return x +} +func foo(_ x: Wrapper) { + _ = gradient(at: Float(1), 2, x, in: generic2) } -_ = gradient(at: Tensor(1), in: generic) // Test case where associated derivative function's requirements are met. -extension Tensor where Scalar : Numeric { +extension Wrapper where Scalar : Numeric { @differentiable(wrt: self where Scalar : Differentiable & FloatingPoint) - func mean() -> Tensor { + func mean() -> Wrapper { return self } @differentiable(wrt: self where Scalar : Differentiable & FloatingPoint) - func variance() -> Tensor { + func variance() -> Wrapper { return mean() // ok } } -_ = pullback(at: Tensor(1), in: { $0.variance() }) +_ = pullback(at: Wrapper(1), in: { $0.variance() }) // Tests TF-277. protocol Layer : Differentiable { diff --git a/test/AutoDiff/simple_math.swift b/test/AutoDiff/simple_math.swift index 167ec761160fd..f50cbaa99f33d 100644 --- a/test/AutoDiff/simple_math.swift +++ b/test/AutoDiff/simple_math.swift @@ -1,7 +1,7 @@ -// RUN: %target-swift-frontend -Xllvm -sil-print-after=differentiation %s -emit-sil -o /dev/null 2>&1 | %FileCheck %s // RUN: %target-run-simple-swift // NOTE(TF-813): verify that enabling forward-mode does not affect reverse-mode. // RUN: %target_run_simple_swift_forward_mode_differentiation +// RUN: %target-swift-frontend -Xllvm -sil-print-after=differentiation %s -emit-sil -o /dev/null 2>&1 | %FileCheck %s // REQUIRES: executable_test import StdlibUnittest @@ -132,16 +132,13 @@ SimpleMathTests.test("TupleMutation") { } expectEqual(405, gradient(at: 3, in: nested)) - // FIXME(TF-201): Update after reabstraction thunks can be directly differentiated. - /* - func generic(_ x: T) -> T { - var tuple = (x, x) - tuple.0 += x - tuple.1 += x - return tuple.0 + tuple.0 + func generic(_ x: T) -> T { + var tuple = (T.zero, T.zero) + tuple.0 = x + tuple.1 = x + return tuple.0 } expectEqual(1, gradient(at: 3.0, in: generic)) - */ } // Tests TF-321. From 35a6fe3854c59170e4dcdf8419c58276ac0140ff Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Wed, 4 Dec 2019 23:48:44 -0800 Subject: [PATCH 2/6] Refactor tests to avoid orthogonal issue TF-1033. --- ...differential-ownership-uninit-memory.swift | 21 +++++++++++++++++++ test/AutoDiff/forward_mode_runtime.swift | 14 ++++++++++++- test/AutoDiff/simple_math.swift | 14 ++++++++++++- 3 files changed, 47 insertions(+), 2 deletions(-) create mode 100644 test/AutoDiff/compiler_crashers/tf1033-differential-ownership-uninit-memory.swift diff --git a/test/AutoDiff/compiler_crashers/tf1033-differential-ownership-uninit-memory.swift b/test/AutoDiff/compiler_crashers/tf1033-differential-ownership-uninit-memory.swift new file mode 100644 index 0000000000000..4a24ec33d8857 --- /dev/null +++ b/test/AutoDiff/compiler_crashers/tf1033-differential-ownership-uninit-memory.swift @@ -0,0 +1,21 @@ +// RUN: not --crash %target-swift-emit-sil -enable-experimental-forward-mode-differentiation %s +// REQUIRES: asserts + +// TF-1033: Ownership verification error in differential function generated by +// the differentiation transform. +// Differential generation sees a `copy_addr [take]` in original function and +// emits a `copy_addr [take]` in the differential function, but this is +// problematic because the source buffer in the differential function is not +// fully initialized. + +@_silgen_name("tuple") +@differentiable +func tupleInitialNonactive(_ x: T) -> T { + var tuple = (T.zero, T.zero) + tuple.0 = x + return tuple.0 +} + +// SIL memory lifetime failure in @AD__tuple__differential_src_0_wrt_0: memory is not initialized, but should +// memory location: %15 = tuple_element_addr %14 : $*(τ_0_0.TangentVector, τ_0_0.TangentVector), 0 // user: %16 +// at instruction: copy_addr [take] %7 to %10 : $*τ_0_0.TangentVector // id: %11 diff --git a/test/AutoDiff/forward_mode_runtime.swift b/test/AutoDiff/forward_mode_runtime.swift index 3f05c9aee2ec1..29530a56309d8 100644 --- a/test/AutoDiff/forward_mode_runtime.swift +++ b/test/AutoDiff/forward_mode_runtime.swift @@ -319,12 +319,24 @@ ForwardModeTests.test("TupleMutation") { expectEqual(405, derivative(at: 3, in: nested)) func generic(_ x: T) -> T { + var tuple = (x, x) + return tuple.0 + } + expectEqual(1, derivative(at: 3.0, in: generic)) + + // FIXME(TF-1033): Fix forward-mode ownership error for tuple with non-active + // initial values. + /* + func genericInitialNonactive( + _ x: T + ) -> T { var tuple = (T.zero, T.zero) tuple.0 = x tuple.1 = x return tuple.0 } - expectEqual(1, derivative(at: 3.0, in: generic)) + expectEqual(1, derivative(at: 3.0, in: genericInitialNonactive)) + */ } // Tests TF-321. diff --git a/test/AutoDiff/simple_math.swift b/test/AutoDiff/simple_math.swift index f50cbaa99f33d..7e4a829a833ba 100644 --- a/test/AutoDiff/simple_math.swift +++ b/test/AutoDiff/simple_math.swift @@ -133,12 +133,24 @@ SimpleMathTests.test("TupleMutation") { expectEqual(405, gradient(at: 3, in: nested)) func generic(_ x: T) -> T { + var tuple = (x, x) + return tuple.0 + } + expectEqual(1, gradient(at: 3.0, in: generic)) + + // FIXME(TF-1033): Fix forward-mode ownership error for tuple with non-active + // initial values. + /* + func genericInitialNonactive( + _ x: T + ) -> T { var tuple = (T.zero, T.zero) tuple.0 = x tuple.1 = x return tuple.0 } - expectEqual(1, gradient(at: 3.0, in: generic)) + expectEqual(1, gradient(at: 3.0, in: genericInitialNonactive)) + */ } // Tests TF-321. From c7c0ebe682b2b7087b7bc081d061b7b1acbe3502 Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Thu, 5 Dec 2019 01:53:51 -0800 Subject: [PATCH 3/6] Minor renaming for clarity. --- .../Utils/Differentiation/ADContext.h | 3 ++- lib/SILOptimizer/Mandatory/Differentiation.cpp | 16 ++++++++-------- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/include/swift/SILOptimizer/Utils/Differentiation/ADContext.h b/include/swift/SILOptimizer/Utils/Differentiation/ADContext.h index 56ef213c4ee6a..7dd465aa10694 100644 --- a/include/swift/SILOptimizer/Utils/Differentiation/ADContext.h +++ b/include/swift/SILOptimizer/Utils/Differentiation/ADContext.h @@ -136,7 +136,8 @@ class ADContext { } /// Adds the given `differentiable_function` instruction to the worklist. - void addDifferentiableFunctionInst(DifferentiableFunctionInst* dfi) { + void + addDifferentiableFunctionInstToWorklist(DifferentiableFunctionInst *dfi) { differentiableFunctionInsts.push_back(dfi); } diff --git a/lib/SILOptimizer/Mandatory/Differentiation.cpp b/lib/SILOptimizer/Mandatory/Differentiation.cpp index ca8c98e7fd761..b952ccd658000 100644 --- a/lib/SILOptimizer/Mandatory/Differentiation.cpp +++ b/lib/SILOptimizer/Mandatory/Differentiation.cpp @@ -1291,7 +1291,7 @@ static SILValue reapplyFunctionConversion( "one argument"); auto *dfi = context.createDifferentiableFunction( builder, loc, reabstractionThunkIndices, newArgs.back()); - context.addDifferentiableFunctionInst(dfi); + context.addDifferentiableFunctionInstToWorklist(dfi); newArgs.back() = dfi; } auto innerNewFunc = reapplyFunctionConversion( @@ -2654,7 +2654,7 @@ class VJPEmitter final getBuilder(), loc, indices.parameters, original); // Record the `differentiable_function` instruction. - context.addDifferentiableFunctionInst(diffFuncInst); + context.addDifferentiableFunctionInstToWorklist(diffFuncInst); // TODO(TF-689): Make `differentiable_function` store result indices and // remove `ADContext::resultIndices`. context.setResultIndex(diffFuncInst, activeResultIndices.front()); @@ -2753,7 +2753,7 @@ class VJPEmitter final // instruction to the `differentiable_function` worklist. TypeSubstCloner::visitDifferentiableFunctionInst(dfi); auto *newDFI = cast(getOpValue(dfi)); - context.addDifferentiableFunctionInst(newDFI); + context.addDifferentiableFunctionInstToWorklist(newDFI); } }; } // end anonymous namespace @@ -4339,7 +4339,7 @@ class JVPEmitter final builder, loc, indices.parameters, original); // Record the `differentiable_function` instruction. - context.addDifferentiableFunctionInst(diffFuncInst); + context.addDifferentiableFunctionInstToWorklist(diffFuncInst); // TODO(TF-689): Make `differentiable_function` store result indices and // remove `ADContext::resultIndices`. context.setResultIndex(diffFuncInst, activeResultIndices.front()); @@ -4478,7 +4478,7 @@ class JVPEmitter final // instruction to the `differentiable_function` worklist. TypeSubstCloner::visitDifferentiableFunctionInst(dfi); auto *newDFI = cast(getOpValue(dfi)); - context.addDifferentiableFunctionInst(newDFI); + context.addDifferentiableFunctionInstToWorklist(newDFI); } }; } // end anonymous namespace @@ -7361,7 +7361,7 @@ SILValue DifferentiationTransformer::promoteToDifferentiableFunction( retInst->eraseFromParent(); context.recordGeneratedFunction(newThunk); - context.addDifferentiableFunctionInst(dfi); + context.addDifferentiableFunctionInstToWorklist(dfi); if (processDifferentiableFunctionInst(dfi)) return nullptr; } @@ -7485,7 +7485,7 @@ SILValue DifferentiationTransformer::promoteToDifferentiableFunction( builder, loc, parameterIndices, origFnCopy, std::make_pair(derivativeFns[0], derivativeFns[1])); context.setResultIndex(dfi, resultIndex); - context.addDifferentiableFunctionInst(dfi); + context.addDifferentiableFunctionInstToWorklist(dfi); return newDFI; } @@ -7603,7 +7603,7 @@ void Differentiation::run() { for (SILBasicBlock &bb : f) { for (SILInstruction &i : bb) { if (auto *dfi = dyn_cast(&i)) - context.addDifferentiableFunctionInst(dfi); + context.addDifferentiableFunctionInstToWorklist(dfi); // Reject uncanonical `linear_function` instructions. // FIXME(SR-11850): Add support for linear map transposition. else if (auto *lfi = dyn_cast(&i)) { From 8f0659d32f5ee1dcf724dd9389b6e8c01ddd2db1 Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Thu, 5 Dec 2019 10:56:43 -0800 Subject: [PATCH 4/6] Address review feedback. - Add doc comment for `SILFunctionType::getAutoDiffDerivativeFunctionType` and explain reabstraction thunk special case. - Change `reapplyFunctionConversion` to always take parameter indices. - Reabstraction thunk special logic is triggered when the original callee is a reabstraction thunk `function_ref`. --- include/swift/AST/Types.h | 65 +++++++++++++++++- lib/SIL/SILFunctionType.cpp | 8 ++- .../Mandatory/Differentiation.cpp | 66 +++++++++++-------- test/AutoDiff/generics.swift | 4 +- test/AutoDiff/reabstraction_thunk.swift | 16 +++++ 5 files changed, 124 insertions(+), 35 deletions(-) create mode 100644 test/AutoDiff/reabstraction_thunk.swift diff --git a/include/swift/AST/Types.h b/include/swift/AST/Types.h index 5028a725c4231..468b7aebb109b 100644 --- a/include/swift/AST/Types.h +++ b/include/swift/AST/Types.h @@ -4303,7 +4303,70 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode, CanSILFunctionType getWithoutDifferentiability(); - /// Returns the type of the derivative function. + /// Returns the type of the derivative function for the given parameter + /// indices, result index, derivative function kind, derivative function + /// generic signature (optional), and other auxiliary parameters. + /// + /// Preconditions: + /// - Parameters corresponding to parameter indices must conform to + /// `Differentiable`. + /// - The result corresponding to the result index must conform to + /// `Differentiable`. + /// + /// Typing rules, given: + /// - Original function type: $(T0, T1, ...) -> (R0, R1, ...) + /// + /// Terminology: + /// - The derivative of a `Differentiable`-conforming type has the + /// `TangentVector` associated type. `TangentVector` is abbreviated as `Tan` + /// below. + /// - "wrt" parameters refers to parameters indicated by the parameter + /// indices. + /// - "wrt" result refers to the result indicated by the result index. + /// + /// JVP derivative type: + /// - Takes original parameters. + /// - Returns original results, followed by a differential function, which + /// takes "wrt" parameter derivatives and returns a "wrt" result derivative. + /// + /// $(T0, ...) -> (R0, ..., (T0.Tan, T1.Tan, ...) -> R0.Tan) + /// ^~~~~~~ ^~~~~~~~~~~~~~~~~~~ ^~~~~~ + /// original results | derivatives wrt params | derivative wrt result + /// + /// VJP derivative type: + /// - Takes original parameters. + /// - Returns original results, followed by a pullback function, which + /// takes a "wrt" result derivative and returns "wrt" parameter derivatives. + /// + /// $(T0, ...) -> (R0, ..., (R0.Tan) -> (T0.Tan, T1.Tan, ...)) + /// ^~~~~~~ ^~~~~~ ^~~~~~~~~~~~~~~~~~~ + /// original results | derivative wrt result | derivatives wrt params + /// + /// The JVP/VJP generic signature is a "constrained" version of the given + /// `derivativeFunctionGenericSignature` if specified. Otherwise, it is a + /// "constrained" version of the original generic signature. A "constrained" + /// generic signature requires all "wrt" parameters to conform to + /// `Differentiable`; this is important for correctness. + /// + /// Other properties of the original function type are copied exactly: + /// `ExtInfo`, coroutine kind, callee convention, yields, optional error + /// result, witness method conformance, etc. + /// + /// Special cases: + /// - Reabstraction thunks have special derivative type calculation. The + /// original function-typed last parameter is transformed into a + /// `@differentiable` function-typed parameter in the derivative type. This + /// is necessary for the differentiation transform to support reabstraction + /// thunk differentiation because the function argument is opaque and cannot + /// be differentiated. Instead, the argument is made `@differentiable` and + /// reabstraction thunk JVP/VJP callers are reponsible for passing a + /// `@differentiable` function. + /// + /// Caveats: + /// - We may support multiple result indices instead of a single result index + /// eventually. At the SIL level, this enables differentiating wrt multiple + /// function results. At the Swift level, this enables differentiating wrt + /// multiple tuple elements for tuple-returning functions. CanSILFunctionType getAutoDiffDerivativeFunctionType( IndexSubset *parameterIndices, unsigned resultIndex, AutoDiffDerivativeFunctionKind kind, Lowering::TypeConverter &TC, diff --git a/lib/SIL/SILFunctionType.cpp b/lib/SIL/SILFunctionType.cpp index 982f0a7bdb29f..cc5c5a7af4ebb 100644 --- a/lib/SIL/SILFunctionType.cpp +++ b/lib/SIL/SILFunctionType.cpp @@ -344,10 +344,12 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType( SmallVector newParameters; newParameters.reserve(getNumParameters()); newParameters.append(getParameters().begin(), getParameters().end()); - // Reabstraction thunks have a function-typed argument (the function to - // reabstract) as their last argument. Reabstraction thunk JVPs/VJPs have a - // `@differentiable` function-typed last argument instead. + // Reabstraction thunks have a function-typed parameter (the function to + // reabstract) as their last parameter. Reabstraction thunk JVPs/VJPs have a + // `@differentiable` function-typed last parameter instead. if (isReabstractionThunk) { + assert(!parameterIndices->contains(getNumParameters() - 1) && + "Function-typed parameter should not be wrt"); auto fnParam = newParameters.back(); auto fnParamType = dyn_cast(fnParam.getInterfaceType()); assert(fnParamType); diff --git a/lib/SILOptimizer/Mandatory/Differentiation.cpp b/lib/SILOptimizer/Mandatory/Differentiation.cpp index b952ccd658000..08ddf345d6382 100644 --- a/lib/SILOptimizer/Mandatory/Differentiation.cpp +++ b/lib/SILOptimizer/Mandatory/Differentiation.cpp @@ -1241,18 +1241,20 @@ static void copyParameterArgumentsForApply( } } -/// When a function value is used in an instruction (usually `apply`), there's -/// some conversion instruction in between, e.g. `thin_to_thick_function`. Given +/// When a function value is used in an instruction (usually `apply`), there may +/// be conversion instructions in between, e.g. `thin_to_thick_function`. Given /// a new function value and an old function value, this helper function /// recursively converts the new function just like how the old function is -/// converted. If the new function's generic signature is specified, it is used +/// converted. +/// +/// If the new function's generic signature is specified, it is used /// to create substitution maps for reapplied `partial_apply` instructions. static SILValue reapplyFunctionConversion( ADContext &context, SILValue newFunc, SILValue oldFunc, SILValue oldConvertedFunc, SILBuilder &builder, SILLocation loc, SmallVectorImpl &newBuffersToDealloc, - GenericSignature newFuncGenSig = GenericSignature(), - IndexSubset *reabstractionThunkIndices = nullptr) { + IndexSubset *parameterIndices, + GenericSignature newFuncGenSig = GenericSignature()) { // If the old func is the new func, then there's no conversion. if (oldFunc == oldConvertedFunc) return newFunc; @@ -1261,7 +1263,7 @@ static SILValue reapplyFunctionConversion( if (auto *cvi = dyn_cast(oldConvertedFunc)) { auto innerNewFunc = reapplyFunctionConversion( context, newFunc, oldFunc, cvi->getOperand(), builder, loc, - newBuffersToDealloc, newFuncGenSig); + newBuffersToDealloc, parameterIndices, newFuncGenSig); // Note: no `copy_value` is needed for the re-converted function because the // caller of `reapplyFunctionConversion` should consume the re-converted // function. @@ -1271,7 +1273,7 @@ static SILValue reapplyFunctionConversion( if (auto *tttfi = dyn_cast(oldConvertedFunc)) { auto innerNewFunc = reapplyFunctionConversion( context, newFunc, oldFunc, tttfi->getOperand(), builder, loc, - newBuffersToDealloc, newFuncGenSig); + newBuffersToDealloc, parameterIndices, newFuncGenSig); auto operandFnTy = innerNewFunc->getType().castTo(); auto thickTy = operandFnTy->getWithRepresentation( SILFunctionTypeRepresentation::Thick); @@ -1285,18 +1287,28 @@ static SILValue reapplyFunctionConversion( SmallVector newArgsToDestroy; copyParameterArgumentsForApply(pai, newArgs, newArgsToDestroy, newBuffersToDealloc); - if (reabstractionThunkIndices) { + auto innerNewFunc = reapplyFunctionConversion( + context, newFunc, oldFunc, pai->getCallee(), builder, loc, + newBuffersToDealloc, parameterIndices, newFuncGenSig); + // Reabstraction thunk `partial_apply` reapplications require special + // support. Reabstraction thunk JVP/VJP expects a `@differentiable` + // function-typed argument to avoid opaque function non-differentiability + // errors. Thus, `partial_apply` reapplications must first form a + // `differentiable_function` of the function-typed thunk argument. + auto isReabstractionThunkCallee = [&]() -> bool { + auto *fri = dyn_cast(oldFunc); + return fri && fri->getReferencedFunctionOrNull()->isThunk() == + IsReabstractionThunk; + }; + if (isReabstractionThunkCallee()) { assert(newArgs.size() == 1 && "Expected reabstraction thunk to be partially applied with only " "one argument"); auto *dfi = context.createDifferentiableFunction( - builder, loc, reabstractionThunkIndices, newArgs.back()); + builder, loc, parameterIndices, newArgs.back()); context.addDifferentiableFunctionInstToWorklist(dfi); newArgs.back() = dfi; } - auto innerNewFunc = reapplyFunctionConversion( - context, newFunc, oldFunc, pai->getCallee(), builder, loc, - newBuffersToDealloc, newFuncGenSig); // If new function's generic signature is specified, use it to create // substitution map for reapplied `partial_apply` instruction. auto substMap = !newFuncGenSig @@ -1482,17 +1494,13 @@ emitDerivativeFunctionReference( } auto *derivativeFnRef = builder.createDifferentiabilityWitnessFunction( loc, witnessKind, minimalWitness); - IndexSubset *reabstractionThunkIndices = nullptr; - if (originalFn->isThunk() == IsReabstractionThunk) - reabstractionThunkIndices = desiredIndices.parameters; - auto convertedRef = - reapplyFunctionConversion(context, derivativeFnRef, originalFRI, - original, builder, loc, newBuffersToDealloc, - derivativeFnRef->getType() - .getASTType() - ->castTo() - ->getSubstGenericSignature(), - reabstractionThunkIndices); + auto convertedRef = reapplyFunctionConversion( + context, derivativeFnRef, originalFRI, original, builder, loc, + newBuffersToDealloc, desiredIndices.parameters, + derivativeFnRef->getType() + .getASTType() + ->castTo() + ->getSubstGenericSignature()); return std::make_pair( convertedRef, SILAutoDiffIndices(desiredIndices.source, @@ -1537,9 +1545,9 @@ emitDerivativeFunctionReference( loc, witnessMethod->getLookupType(), witnessMethod->getConformance(), requirementDeclRef.asAutoDiffDerivativeFunction(autoDiffFuncId), SILType::getPrimitiveObjectType(assocType)); - auto convertedRef = - reapplyFunctionConversion(context, ref, witnessMethod, original, - builder, loc, newBuffersToDealloc); + auto convertedRef = reapplyFunctionConversion( + context, ref, witnessMethod, original, builder, loc, + newBuffersToDealloc, desiredIndices.parameters); return std::make_pair(convertedRef, minimalIndices); } @@ -1582,9 +1590,9 @@ emitDerivativeFunctionReference( loc, classMethodInst->getOperand(), methodDeclRef.asAutoDiffDerivativeFunction(autoDiffFuncId), SILType::getPrimitiveObjectType(assocType)); - auto convertedRef = - reapplyFunctionConversion(context, ref, classMethodInst, original, - builder, loc, newBuffersToDealloc); + auto convertedRef = reapplyFunctionConversion( + context, ref, classMethodInst, original, builder, loc, + newBuffersToDealloc, desiredIndices.parameters); return std::make_pair(convertedRef, minimalIndices); } diff --git a/test/AutoDiff/generics.swift b/test/AutoDiff/generics.swift index 50560144d907a..fc27aa4bffea5 100644 --- a/test/AutoDiff/generics.swift +++ b/test/AutoDiff/generics.swift @@ -16,8 +16,8 @@ _ = gradient(at: Float(1), in: { x in identity(x) }) // CHECK-SIL-NEXT: [[EMIT_ZERO_INDIRECT:%.*]] = apply [[ZERO_WITNESS]]<τ_0_0.TangentVector>([[ORIG_COTAN]], [[ORIG_COTAN_METATYPE]]) // CHECK-SIL: } -// Test TF-201: differentiate direct references to generic function: -// partially-applied reabstraction thunk. +// Test TF-201: differentiate direct references to generic function. +// This involves reabstraction thunk differentiation. _ = gradient(at: Float(1), in: identity) diff --git a/test/AutoDiff/reabstraction_thunk.swift b/test/AutoDiff/reabstraction_thunk.swift new file mode 100644 index 0000000000000..6d8834c77df32 --- /dev/null +++ b/test/AutoDiff/reabstraction_thunk.swift @@ -0,0 +1,16 @@ +// RUN: %target-swift-frontend -emit-sil -verify %s | %FileCheck %s + +// Test reabstraction thunk differentiation. +// Check that reabstraction thunk JVP/VJP functions have a `@differentiable` +// function-typed argument. + +@_silgen_name("id") +func id(_ x: T) -> T { x } +let _: @differentiable (Float) -> Float = id + +// CHECK-LABEL: sil_differentiability_witness shared [serialized] [parameters 0] [results 0] @$sS2fIegnr_S2fIegyd_TR : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float { +// CHECK: jvp: @AD__$sS2fIegnr_S2fIegyd_TR__jvp_src_0_wrt_0 : $@convention(thin) (Float, @guaranteed @differentiable @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) +// CHECK: vjp: @AD__$sS2fIegnr_S2fIegyd_TR__vjp_src_0_wrt_0 : $@convention(thin) (Float, @guaranteed @differentiable @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) +// CHECK: } + +// CHECK-LABEL: sil hidden [serializable] @AD__$sS2fIegnr_S2fIegyd_TR__vjp_src_0_wrt_0 : $@convention(thin) (Float, @guaranteed @differentiable @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) { From f6c56f1599d04990e16ecf5bd3c6d50cc9e99618 Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Thu, 5 Dec 2019 16:03:22 -0800 Subject: [PATCH 5/6] Reference TF-1036 optimizations. --- include/swift/AST/Types.h | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/include/swift/AST/Types.h b/include/swift/AST/Types.h index 468b7aebb109b..9e732f7975d79 100644 --- a/include/swift/AST/Types.h +++ b/include/swift/AST/Types.h @@ -4361,6 +4361,12 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode, /// be differentiated. Instead, the argument is made `@differentiable` and /// reabstraction thunk JVP/VJP callers are reponsible for passing a /// `@differentiable` function. + /// - TODO(TF-1036): Investigate more efficient reabstraction thunk + /// derivative approaches. The last argument can simply be a + /// corresponding derivative function, instead of a `@differentiable` + /// function - this is more direct. It may be possible to implement + /// reabstraction thunk derivatives using "reabstraction thunks for + /// the original function's derivative", avoiding extra code generation. /// /// Caveats: /// - We may support multiple result indices instead of a single result index From e870b6d232e7df80b7edc6b68a6ad94c9d3da350 Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Thu, 5 Dec 2019 23:18:26 -0800 Subject: [PATCH 6/6] Fix test based on upstream changes. --- test/AutoDiff/reabstraction_thunk.swift | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/AutoDiff/reabstraction_thunk.swift b/test/AutoDiff/reabstraction_thunk.swift index 6d8834c77df32..af35ae289f551 100644 --- a/test/AutoDiff/reabstraction_thunk.swift +++ b/test/AutoDiff/reabstraction_thunk.swift @@ -8,9 +8,9 @@ func id(_ x: T) -> T { x } let _: @differentiable (Float) -> Float = id -// CHECK-LABEL: sil_differentiability_witness shared [serialized] [parameters 0] [results 0] @$sS2fIegnr_S2fIegyd_TR : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float { +// CHECK-LABEL: sil_differentiability_witness private [parameters 0] [results 0] @$sS2fIegnr_S2fIegyd_TR : $@convention(thin) (Float, @guaranteed @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> Float { // CHECK: jvp: @AD__$sS2fIegnr_S2fIegyd_TR__jvp_src_0_wrt_0 : $@convention(thin) (Float, @guaranteed @differentiable @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) // CHECK: vjp: @AD__$sS2fIegnr_S2fIegyd_TR__vjp_src_0_wrt_0 : $@convention(thin) (Float, @guaranteed @differentiable @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) // CHECK: } -// CHECK-LABEL: sil hidden [serializable] @AD__$sS2fIegnr_S2fIegyd_TR__vjp_src_0_wrt_0 : $@convention(thin) (Float, @guaranteed @differentiable @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) { +// CHECK-LABEL: sil private @AD__$sS2fIegnr_S2fIegyd_TR__vjp_src_0_wrt_0 : $@convention(thin) (Float, @guaranteed @differentiable @callee_guaranteed (@in_guaranteed Float) -> @out Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) {