From 2ccc0b31efb30b32c1094e3b5c6051309a120070 Mon Sep 17 00:00:00 2001 From: Kshitij Date: Tue, 15 Aug 2023 10:38:23 -0700 Subject: [PATCH 1/3] Fixes memory leaks in autodiff linear map context allocation builtins Fixes #67323 --- include/swift/AST/Builtins.def | 4 +- include/swift/Runtime/RuntimeFunctions.def | 8 +-- lib/AST/Builtins.cpp | 9 +-- lib/IRGen/GenBuiltin.cpp | 9 +-- lib/IRGen/GenCall.cpp | 15 ++-- lib/IRGen/GenCall.h | 9 +-- .../Differentiation/VJPCloner.cpp | 29 ++++++-- stdlib/public/runtime/AutoDiffSupport.cpp | 70 +++++++++++++------ stdlib/public/runtime/AutoDiffSupport.h | 65 ++++++++++++----- test/AutoDiff/IRGen/runtime.swift | 22 +++--- test/AutoDiff/SILGen/autodiff_builtins.swift | 27 +++---- ...todiff_builtins_deprecated_abi_fails.swift | 20 ++++++ ...ential_not_leaked_in_func_with_loops.swift | 60 ++++++++++++++++ 13 files changed, 245 insertions(+), 102 deletions(-) create mode 100644 test/AutoDiff/SILGen/autodiff_builtins_deprecated_abi_fails.swift create mode 100644 test/AutoDiff/stdlib/callee_differential_not_leaked_in_func_with_loops.swift diff --git a/include/swift/AST/Builtins.def b/include/swift/AST/Builtins.def index fb271e7e1ba88..7483903238faf 100644 --- a/include/swift/AST/Builtins.def +++ b/include/swift/AST/Builtins.def @@ -985,13 +985,13 @@ BUILTIN_MISC_OPERATION_WITH_SILGEN(CreateAsyncTaskInGroup, /// is a pure value and therefore we can consider it as readnone). BUILTIN_MISC_OPERATION_WITH_SILGEN(GlobalStringTablePointer, "globalStringTablePointer", "n", Special) -// autoDiffCreateLinearMapContext: (Builtin.Word) -> Builtin.NativeObject +// autoDiffCreateLinearMapContext: (T.Type) -> Builtin.NativeObject BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffCreateLinearMapContext, "autoDiffCreateLinearMapContext", "", Special) // autoDiffProjectTopLevelSubcontext: (Builtin.NativeObject) -> Builtin.RawPointer BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffProjectTopLevelSubcontext, "autoDiffProjectTopLevelSubcontext", "n", Special) -// autoDiffAllocateSubcontext: (Builtin.NativeObject, Builtin.Word) -> Builtin.RawPointer +// autoDiffAllocateSubcontext: (Builtin.NativeObject, T.Type) -> Builtin.RawPointer BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffAllocateSubcontext, "autoDiffAllocateSubcontext", "", Special) /// Build a Builtin.Executor value from an "ordinary" serial executor diff --git a/include/swift/Runtime/RuntimeFunctions.def b/include/swift/Runtime/RuntimeFunctions.def index 812f807d971ac..8ed18fce638c7 100644 --- a/include/swift/Runtime/RuntimeFunctions.def +++ b/include/swift/Runtime/RuntimeFunctions.def @@ -2273,12 +2273,12 @@ FUNCTION(TaskGroupDestroy, ATTRS(NoUnwind), EFFECT(Concurrency)) -// AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext(size_t); +// AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext(const Metadata *); FUNCTION(AutoDiffCreateLinearMapContext, swift_autoDiffCreateLinearMapContext, SwiftCC, DifferentiationAvailability, RETURNS(RefCountedPtrTy), - ARGS(SizeTy), + ARGS(TypeMetadataPtrTy), ATTRS(NoUnwind, ArgMemOnly), EFFECT(AutoDiff)) @@ -2291,12 +2291,12 @@ FUNCTION(AutoDiffProjectTopLevelSubcontext, ATTRS(NoUnwind, ArgMemOnly), EFFECT(AutoDiff)) -// void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, size_t); +// void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, const Metadata *); FUNCTION(AutoDiffAllocateSubcontext, swift_autoDiffAllocateSubcontext, SwiftCC, DifferentiationAvailability, RETURNS(Int8PtrTy), - ARGS(RefCountedPtrTy, SizeTy), + ARGS(RefCountedPtrTy, TypeMetadataPtrTy), ATTRS(NoUnwind, ArgMemOnly), EFFECT(AutoDiff)) diff --git a/lib/AST/Builtins.cpp b/lib/AST/Builtins.cpp index 035cbf9beee5a..45fa403ebf81a 100644 --- a/lib/AST/Builtins.cpp +++ b/lib/AST/Builtins.cpp @@ -1608,8 +1608,9 @@ static ValueDecl *getBuildComplexEqualitySerialExecutorRef(ASTContext &ctx, static ValueDecl *getAutoDiffCreateLinearMapContext(ASTContext &ctx, Identifier id) { - return getBuiltinFunction( - id, {BuiltinIntegerType::getWordType(ctx)}, ctx.TheNativeObjectType); + return getBuiltinFunction(ctx, id, _thin, _generics(_unrestricted), + _parameters(_metatype(_typeparam(0))), + _nativeObject); } static ValueDecl *getAutoDiffProjectTopLevelSubcontext(ASTContext &ctx, @@ -1621,8 +1622,8 @@ static ValueDecl *getAutoDiffProjectTopLevelSubcontext(ASTContext &ctx, static ValueDecl *getAutoDiffAllocateSubcontext(ASTContext &ctx, Identifier id) { return getBuiltinFunction( - id, {ctx.TheNativeObjectType, BuiltinIntegerType::getWordType(ctx)}, - ctx.TheRawPointerType); + ctx, id, _thin, _generics(_unrestricted), + _parameters(_nativeObject, _metatype(_typeparam(0))), _rawPointer); } static ValueDecl *getPoundAssert(ASTContext &Context, Identifier Id) { diff --git a/lib/IRGen/GenBuiltin.cpp b/lib/IRGen/GenBuiltin.cpp index c697557a38684..cb0d1c1276828 100644 --- a/lib/IRGen/GenBuiltin.cpp +++ b/lib/IRGen/GenBuiltin.cpp @@ -1308,8 +1308,8 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin, } if (Builtin.ID == BuiltinValueKind::AutoDiffCreateLinearMapContext) { - auto topLevelSubcontextSize = args.claimNext(); - out.add(emitAutoDiffCreateLinearMapContext(IGF, topLevelSubcontextSize) + auto topLevelSubcontextMetaType = args.claimNext(); + out.add(emitAutoDiffCreateLinearMapContext(IGF, topLevelSubcontextMetaType) .getAddress()); return; } @@ -1325,9 +1325,10 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin, if (Builtin.ID == BuiltinValueKind::AutoDiffAllocateSubcontext) { Address allocatorAddr(args.claimNext(), IGF.IGM.RefCountedStructTy, IGF.IGM.getPointerAlignment()); - auto size = args.claimNext(); + auto subcontextMetatype = args.claimNext(); out.add( - emitAutoDiffAllocateSubcontext(IGF, allocatorAddr, size).getAddress()); + emitAutoDiffAllocateSubcontext(IGF, allocatorAddr, subcontextMetatype) + .getAddress()); return; } diff --git a/lib/IRGen/GenCall.cpp b/lib/IRGen/GenCall.cpp index c5f264d0e2c8f..23a9e537f695a 100644 --- a/lib/IRGen/GenCall.cpp +++ b/lib/IRGen/GenCall.cpp @@ -5480,10 +5480,12 @@ IRGenFunction::getFunctionPointerForResumeIntrinsic(llvm::Value *resume) { } Address irgen::emitAutoDiffCreateLinearMapContext( - IRGenFunction &IGF, llvm::Value *topLevelSubcontextSize) { + IRGenFunction &IGF, llvm::Value *topLevelSubcontextMetatype) { + topLevelSubcontextMetatype = IGF.Builder.CreateBitCast( + topLevelSubcontextMetatype, IGF.IGM.TypeMetadataPtrTy); auto *call = IGF.Builder.CreateCall( IGF.IGM.getAutoDiffCreateLinearMapContextFunctionPointer(), - {topLevelSubcontextSize}); + {topLevelSubcontextMetatype}); call->setDoesNotThrow(); call->setCallingConv(IGF.IGM.SwiftCC); return Address(call, IGF.IGM.RefCountedStructTy, @@ -5500,11 +5502,14 @@ Address irgen::emitAutoDiffProjectTopLevelSubcontext( return Address(call, IGF.IGM.Int8Ty, IGF.IGM.getPointerAlignment()); } -Address irgen::emitAutoDiffAllocateSubcontext( - IRGenFunction &IGF, Address context, llvm::Value *size) { +Address irgen::emitAutoDiffAllocateSubcontext(IRGenFunction &IGF, + Address context, + llvm::Value *subcontextMetatype) { + subcontextMetatype = + IGF.Builder.CreateBitCast(subcontextMetatype, IGF.IGM.TypeMetadataPtrTy); auto *call = IGF.Builder.CreateCall( IGF.IGM.getAutoDiffAllocateSubcontextFunctionPointer(), - {context.getAddress(), size}); + {context.getAddress(), subcontextMetatype}); call->setDoesNotThrow(); call->setCallingConv(IGF.IGM.SwiftCC); return Address(call, IGF.IGM.Int8Ty, IGF.IGM.getPointerAlignment()); diff --git a/lib/IRGen/GenCall.h b/lib/IRGen/GenCall.h index c64e97dc8a264..e6ef58f1411d9 100644 --- a/lib/IRGen/GenCall.h +++ b/lib/IRGen/GenCall.h @@ -261,12 +261,13 @@ namespace irgen { CanSILFunctionType fnType, Explosion &result, Explosion &error); - Address emitAutoDiffCreateLinearMapContext( - IRGenFunction &IGF, llvm::Value *topLevelSubcontextSize); + Address + emitAutoDiffCreateLinearMapContext(IRGenFunction &IGF, + llvm::Value *topLevelSubcontextMetatype); Address emitAutoDiffProjectTopLevelSubcontext( IRGenFunction &IGF, Address context); - Address emitAutoDiffAllocateSubcontext( - IRGenFunction &IGF, Address context, llvm::Value *size); + Address emitAutoDiffAllocateSubcontext(IRGenFunction &IGF, Address context, + llvm::Value *subcontextMetatype); FunctionPointer getFunctionPointerForDispatchCall(IRGenModule &IGM, const FunctionPointer &fn); diff --git a/lib/SILOptimizer/Differentiation/VJPCloner.cpp b/lib/SILOptimizer/Differentiation/VJPCloner.cpp index 7cb715b278553..90084d4b56026 100644 --- a/lib/SILOptimizer/Differentiation/VJPCloner.cpp +++ b/lib/SILOptimizer/Differentiation/VJPCloner.cpp @@ -17,6 +17,8 @@ #define DEBUG_TYPE "differentiation" +#include "swift/AST/Types.h" + #include "swift/SILOptimizer/Differentiation/VJPCloner.h" #include "swift/SILOptimizer/Analysis/DifferentiableActivityAnalysis.h" #include "swift/SILOptimizer/Differentiation/ADContext.h" @@ -118,15 +120,21 @@ class VJPCloner::Implementation final auto pullbackTupleType = remapASTType(pullbackInfo.getLinearMapTupleType(returnBB)->getCanonicalType()); Builder.setInsertionPoint(vjp->getEntryBlock()); - auto topLevelSubcontextSize = emitMemoryLayoutSize( - Builder, original->getLocation(), pullbackTupleType); + + auto pbTupleMetatypeType = + CanMetatypeType::get(pullbackTupleType, MetatypeRepresentation::Thick); + auto pbTupleMetatypeSILType = + SILType::getPrimitiveObjectType(pbTupleMetatypeType); + auto pbTupleMetatype = + Builder.createMetatype(original->getLocation(), pbTupleMetatypeSILType); + // Create an context. pullbackContextValue = Builder.createBuiltin( original->getLocation(), getASTContext().getIdentifier( getBuiltinName(BuiltinValueKind::AutoDiffCreateLinearMapContext)), - SILType::getNativeObjectType(getASTContext()), - SubstitutionMap(), {topLevelSubcontextSize}); + SILType::getNativeObjectType(getASTContext()), SubstitutionMap(), + {pbTupleMetatype}); borrowedPullbackContextValue = Builder.createBeginBorrow( original->getLocation(), pullbackContextValue); LLVM_DEBUG(getADDebugStream() @@ -1067,14 +1075,21 @@ EnumInst *VJPCloner::Implementation::buildPredecessorEnumValue( assert(enumEltType == rawPtrType); auto pbTupleType = remapASTType(pullbackInfo.getLinearMapTupleType(predBB)->getCanonicalType()); - SILValue pbTupleSize = - emitMemoryLayoutSize(Builder, loc, pbTupleType); + + auto pbTupleMetatypeType = + CanMetatypeType::get(pbTupleType, MetatypeRepresentation::Thick); + auto pbTupleMetatypeSILType = + SILType::getPrimitiveObjectType(pbTupleMetatypeType); + auto pbTupleMetatype = + Builder.createMetatype(original->getLocation(), pbTupleMetatypeSILType); + auto rawBufferValue = builder.createBuiltin( loc, getASTContext().getIdentifier( getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontext)), rawPtrType, SubstitutionMap(), - {borrowedPullbackContextValue, pbTupleSize}); + {borrowedPullbackContextValue, pbTupleMetatype}); + auto typedBufferValue = builder.createPointerToAddress( loc, rawBufferValue, pbTupleVal->getType().getAddressType(), diff --git a/stdlib/public/runtime/AutoDiffSupport.cpp b/stdlib/public/runtime/AutoDiffSupport.cpp index 3873e6165a7f4..c00feeef4da4c 100644 --- a/stdlib/public/runtime/AutoDiffSupport.cpp +++ b/stdlib/public/runtime/AutoDiffSupport.cpp @@ -13,7 +13,7 @@ #include "AutoDiffSupport.h" #include "swift/ABI/Metadata.h" #include "swift/Runtime/HeapObject.h" - +#include "llvm/ADT/SmallVector.h" #include using namespace swift; @@ -21,7 +21,13 @@ using namespace llvm; SWIFT_CC(swift) static void destroyLinearMapContext(SWIFT_CONTEXT HeapObject *obj) { - static_cast(obj)->~AutoDiffLinearMapContext(); + auto *linearMapContext = static_cast(obj); + + for (auto *heapObjectPtr : linearMapContext->getAllocatedHeapObjects()) { + swift_release(heapObjectPtr); + } + + linearMapContext->~AutoDiffLinearMapContext(); free(obj); } @@ -43,36 +49,54 @@ static FullMetadata linearMapContextHeapMetadata = { } }; -AutoDiffLinearMapContext::AutoDiffLinearMapContext() +AutoDiffLinearMapContext::AutoDiffLinearMapContext( + OpaqueValue *const topLevelLinearMapContextProjection) : HeapObject(&linearMapContextHeapMetadata) { + this->topLevelLinearMapContextProjection = topLevelLinearMapContextProjection; } -void *AutoDiffLinearMapContext::projectTopLevelSubcontext() const { - auto offset = alignTo( - sizeof(AutoDiffLinearMapContext), alignof(AutoDiffLinearMapContext)); - return const_cast( - reinterpret_cast(this) + offset); -} +AutoDiffLinearMapContext *swift::swift_autoDiffCreateLinearMapContext( + const Metadata *topLevelLinearMapContextMetadata) { + // Linear map context metadata must have non-null value witnesses + assert(topLevelLinearMapContextMetadata->getValueWitnesses()); -void *AutoDiffLinearMapContext::allocate(size_t size) { - return allocator.Allocate(size, alignof(AutoDiffLinearMapContext)); -} + // Allocate a box for the top-level linear map context + auto [topLevelContextHeapObjectPtr, toplevelContextProjection] = + swift_allocBox(topLevelLinearMapContextMetadata); -AutoDiffLinearMapContext *swift::swift_autoDiffCreateLinearMapContext( - size_t topLevelLinearMapStructSize) { - auto allocationSize = alignTo( - sizeof(AutoDiffLinearMapContext), alignof(AutoDiffLinearMapContext)) - + topLevelLinearMapStructSize; - auto *buffer = (AutoDiffLinearMapContext *)malloc(allocationSize); - return ::new (buffer) AutoDiffLinearMapContext; + // Create a linear map context object that stores the projection + // for the top level context + auto linearMapContext = + new AutoDiffLinearMapContext(toplevelContextProjection); + + // Stash away the `HeapObject` pointer for the allocated context + // for proper "release" during clean up. + linearMapContext->storeAllocatedHeapObjectPtr(topLevelContextHeapObjectPtr); + + // Return the newly created linear map context object + return linearMapContext; } void *swift::swift_autoDiffProjectTopLevelSubcontext( - AutoDiffLinearMapContext *allocator) { - return allocator->projectTopLevelSubcontext(); + AutoDiffLinearMapContext *linearMapContext) { + return static_cast( + linearMapContext->getTopLevelLinearMapContextProjection()); } void *swift::swift_autoDiffAllocateSubcontext( - AutoDiffLinearMapContext *allocator, size_t size) { - return allocator->allocate(size); + AutoDiffLinearMapContext *linearMapContext, + const Metadata *linearMapSubcontextMetadata) { + // Linear map context metadata must have non-null value witnesses + assert(linearMapSubcontextMetadata->getValueWitnesses()); + + // Allocate a box for the linear map subcontext + auto [subcontextHeapObjectPtr, subcontextProjection] = + swift_allocBox(linearMapSubcontextMetadata); + + // Stash away the `HeapObject` pointer for the allocated context + // for proper "release" during clean up. + linearMapContext->storeAllocatedHeapObjectPtr(subcontextHeapObjectPtr); + + // Return the subcontext projection + return static_cast(subcontextProjection); } diff --git a/stdlib/public/runtime/AutoDiffSupport.h b/stdlib/public/runtime/AutoDiffSupport.h index 4fb63d82e062b..d8ecde3b2a7bc 100644 --- a/stdlib/public/runtime/AutoDiffSupport.h +++ b/stdlib/public/runtime/AutoDiffSupport.h @@ -13,44 +13,71 @@ #ifndef SWIFT_RUNTIME_AUTODIFF_SUPPORT_H #define SWIFT_RUNTIME_AUTODIFF_SUPPORT_H -#include "swift/Runtime/HeapObject.h" #include "swift/Runtime/Config.h" -#include "llvm/Support/Allocator.h" +#include "swift/Runtime/HeapObject.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/SmallVector.h" namespace swift { - /// A data structure responsible for efficiently allocating closure contexts for /// linear maps such as pullbacks, including recursive branching trace enum /// case payloads. class AutoDiffLinearMapContext : public HeapObject { private: - /// The underlying allocator. - // TODO: Use a custom allocator so that the initial slab can be - // tail-allocated. - llvm::BumpPtrAllocator allocator; + // TODO: Commenting out BumpPtrAllocator temporarily + // until we move away from the interim solution of allocating + // boxes for linear map contexts/subcontexts. + // + // /// The underlying allocator. + // // TODO: Use a custom allocator so that the initial slab can be + // // tail-allocated. + // llvm::BumpPtrAllocator allocator; + + /// A projection/pointer to the memory storing the + /// top-level linear map context object. + OpaqueValue *topLevelLinearMapContextProjection; + + /// Storage for `HeapObject` pointers to the linear map + /// context and subcontexts allocated for derivatives with + /// loops. + llvm::SmallVector allocatedHeapObjects; public: /// Creates a linear map context. - AutoDiffLinearMapContext(); - /// Returns the address of the tail-allocated top-level subcontext. - void *projectTopLevelSubcontext() const; - /// Allocates memory for a new subcontext. - void *allocate(size_t size); + AutoDiffLinearMapContext(OpaqueValue *const); + + // TODO: Commenting out BumpPtrAllocator temporarily + // until we move away from the interim solution of allocating + // boxes for linear map contexts/subcontexts. + // + // llvm::BumpPtrAllocator& getAllocator() const { + // return const_cast(this->allocator); + // } + + OpaqueValue *getTopLevelLinearMapContextProjection() const { + return this->topLevelLinearMapContextProjection; + } + + llvm::ArrayRef getAllocatedHeapObjects() const { + return this->allocatedHeapObjects; + } + + void storeAllocatedHeapObjectPtr(HeapObject *allocatedHeapObjectPtr) { + this->allocatedHeapObjects.push_back(allocatedHeapObjectPtr); + } }; /// Creates a linear map context with a tail-allocated top-level subcontext. SWIFT_RUNTIME_EXPORT SWIFT_CC(swift) -AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext( - size_t topLevelSubcontextSize); + AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext( + const Metadata *topLevelLinearMapContextMetadata); /// Returns the address of the tail-allocated top-level subcontext. SWIFT_RUNTIME_EXPORT SWIFT_CC(swift) void *swift_autoDiffProjectTopLevelSubcontext(AutoDiffLinearMapContext *); /// Allocates memory for a new subcontext. -SWIFT_RUNTIME_EXPORT SWIFT_CC(swift) -void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, size_t size); - -} - +SWIFT_RUNTIME_EXPORT SWIFT_CC(swift) void *swift_autoDiffAllocateSubcontext( + AutoDiffLinearMapContext *, const Metadata *linearMapSubcontextMetadata); +} // namespace swift #endif /* SWIFT_RUNTIME_AUTODIFF_SUPPORT_H */ diff --git a/test/AutoDiff/IRGen/runtime.swift b/test/AutoDiff/IRGen/runtime.swift index 09c848fafcc04..a4687d4108e89 100644 --- a/test/AutoDiff/IRGen/runtime.swift +++ b/test/AutoDiff/IRGen/runtime.swift @@ -1,24 +1,18 @@ // RUN: %target-swift-frontend -parse-stdlib %s -emit-ir | %FileCheck %s import Swift -import _Differentiation - -struct ExamplePullbackStruct { - var pb0: (T.TangentVector) -> T.TangentVector -} @_silgen_name("test_context_builtins") -func test_context_builtins() { - let pbStruct = ExamplePullbackStruct(pb0: { $0 }) - let context = Builtin.autoDiffCreateLinearMapContext(Builtin.sizeof(type(of: pbStruct))) +func test_context_builtins(t: T) { + let context = Builtin.autoDiffCreateLinearMapContext(T.self) let topLevelSubctxAddr = Builtin.autoDiffProjectTopLevelSubcontext(context) - UnsafeMutableRawPointer(topLevelSubctxAddr).storeBytes(of: pbStruct, as: type(of: pbStruct)) - let newBuffer = Builtin.autoDiffAllocateSubcontext(context, Builtin.sizeof(type(of: pbStruct))) - UnsafeMutableRawPointer(newBuffer).storeBytes(of: pbStruct, as: type(of: pbStruct)) + UnsafeMutableRawPointer(topLevelSubctxAddr).storeBytes(of: t, as: T.self) + let newBuffer = Builtin.autoDiffAllocateSubcontext(context, T.self) + UnsafeMutableRawPointer(newBuffer).storeBytes(of: t, as: T.self) } -// CHECK-LABEL: define{{.*}}@test_context_builtins() +// CHECK-LABEL: define{{.*}}@test_context_builtins(ptr noalias nocapture %0, ptr %T) // CHECK: entry: -// CHECK: [[CTX:%.*]] = call swiftcc ptr @swift_autoDiffCreateLinearMapContext({{i[0-9]+}} {{.*}}) +// CHECK: [[CTX:%.*]] = call swiftcc ptr @swift_autoDiffCreateLinearMapContext(ptr %T) // CHECK: call swiftcc ptr @swift_autoDiffProjectTopLevelSubcontext(ptr [[CTX]]) -// CHECK: [[BUF:%.*]] = call swiftcc ptr @swift_autoDiffAllocateSubcontext(ptr [[CTX]], {{i[0-9]+}} {{.*}}) +// CHECK: [[BUF:%.*]] = call swiftcc ptr @swift_autoDiffAllocateSubcontext(ptr [[CTX]], ptr %T) diff --git a/test/AutoDiff/SILGen/autodiff_builtins.swift b/test/AutoDiff/SILGen/autodiff_builtins.swift index 903ba227fbf71..7f21e1f8dacce 100644 --- a/test/AutoDiff/SILGen/autodiff_builtins.swift +++ b/test/AutoDiff/SILGen/autodiff_builtins.swift @@ -84,24 +84,19 @@ func applyDerivative_f1_vjp(t0: T) -> (T // CHECK: copy_addr [take] [[D_RESULT_BUFFER_0_FOR_LOAD]] to [init] [[ORIG_RESULT_OUT_PARAM]] // CHECK: return [[PULLBACK]] -struct ExamplePullbackStruct { - var pb0: (T.TangentVector) -> T.TangentVector -} - @_silgen_name("test_context_builtins") -func test_context_builtins() { - let pbStruct = ExamplePullbackStruct(pb0: { $0 }) - let context = Builtin.autoDiffCreateLinearMapContext(Builtin.sizeof(type(of: pbStruct))) +func test_context_builtins(t: T) { + let context = Builtin.autoDiffCreateLinearMapContext(T.self) let topLevelSubctxAddr = Builtin.autoDiffProjectTopLevelSubcontext(context) - UnsafeMutableRawPointer(topLevelSubctxAddr).storeBytes(of: pbStruct, as: type(of: pbStruct)) - let newBuffer = Builtin.autoDiffAllocateSubcontext(context, Builtin.sizeof(type(of: pbStruct))) - UnsafeMutableRawPointer(newBuffer).storeBytes(of: pbStruct, as: type(of: pbStruct)) + UnsafeMutableRawPointer(topLevelSubctxAddr).storeBytes(of: t, as: T.self) + let newBuffer = Builtin.autoDiffAllocateSubcontext(context, T.self) + UnsafeMutableRawPointer(newBuffer).storeBytes(of: t, as: T.self) } -// CHECK-LABEL: sil{{.*}}@test_context_builtins -// CHECK: bb0: -// CHECK: [[CTX:%.*]] = builtin "autoDiffCreateLinearMapContext"({{%.*}} : $Builtin.Word) : $Builtin.NativeObject -// CHECK: [[BORROWED_CTX:%.*]] = begin_borrow [lexical] [[CTX]] : $Builtin.NativeObject -// CHECK: [[BUF:%.*]] = builtin "autoDiffProjectTopLevelSubcontext"([[BORROWED_CTX]] : $Builtin.NativeObject) : $Builtin.RawPointer -// CHECK: [[BUF:%.*]] = builtin "autoDiffAllocateSubcontext"([[BORROWED_CTX]] : $Builtin.NativeObject, {{.*}} : $Builtin.Word) : $Builtin.RawPointer +// CHECK-LABEL: sil{{.*}}@test_context_builtins : $@convention(thin) (@in_guaranteed T) -> () { +// CHECK: bb0({{%.*}} : $*T): +// CHECK: [[CTX:%.*]] = builtin "autoDiffCreateLinearMapContext"({{%.*}} : $@thick T.Type) : $Builtin.NativeObject // users: {{.*}} +// CHECK: [[BORROWED_CTX:%.*]] = begin_borrow [lexical] [[CTX]] : $Builtin.NativeObject // users: {{.*}} +// CHECK: [[BUF:%.*]] = builtin "autoDiffProjectTopLevelSubcontext"([[BORROWED_CTX]] : $Builtin.NativeObject) : $Builtin.RawPointer // users: {{.*}} +// CHECK: [[BUF:%.*]] = builtin "autoDiffAllocateSubcontext"([[BORROWED_CTX]] : $Builtin.NativeObject, {{.*}} : $@thick T.Type) : $Builtin.RawPointer // users: {{.*}} // CHECK: destroy_value [[CTX]] diff --git a/test/AutoDiff/SILGen/autodiff_builtins_deprecated_abi_fails.swift b/test/AutoDiff/SILGen/autodiff_builtins_deprecated_abi_fails.swift new file mode 100644 index 0000000000000..5d6b756a7eb91 --- /dev/null +++ b/test/AutoDiff/SILGen/autodiff_builtins_deprecated_abi_fails.swift @@ -0,0 +1,20 @@ +// RUN: %target-swift-frontend -emit-silgen -parse-stdlib %s +// XFAIL: * + +// SIL code with older, deprecated autodiff ABI for linear map context +// allocation should fail, complaining about receiving a 'Builtin.Word' when +// a 'T.Type' is expected. + +struct ExamplePullbackStruct { + var pb0: (T.TangentVector) -> T.TangentVector +} + +@_silgen_name("test_context_builtins_bad") +func test_context_builtins_bad() { + let pbStruct = ExamplePullbackStruct(pb0: { $0 }) + let context = Builtin.autoDiffCreateLinearMapContext(Builtin.sizeof(type(of: pbStruct))) + let topLevelSubctxAddr = Builtin.autoDiffProjectTopLevelSubcontext(context) + UnsafeMutableRawPointer(topLevelSubctxAddr).storeBytes(of: pbStruct, as: type(of: pbStruct)) + let newBuffer = Builtin.autoDiffAllocateSubcontext(context, Builtin.sizeof(type(of: pbStruct))) + UnsafeMutableRawPointer(newBuffer).storeBytes(of: pbStruct, as: type(of: pbStruct)) +} diff --git a/test/AutoDiff/stdlib/callee_differential_not_leaked_in_func_with_loops.swift b/test/AutoDiff/stdlib/callee_differential_not_leaked_in_func_with_loops.swift new file mode 100644 index 0000000000000..8de10d1d02d2b --- /dev/null +++ b/test/AutoDiff/stdlib/callee_differential_not_leaked_in_func_with_loops.swift @@ -0,0 +1,60 @@ +// RUN: %target-run-simple-swift +// REQUIRES: executable_test + +import _Differentiation +import StdlibUnittest + +// When the original function contains loops, we allocate a context object +// on the heap. This context object may store non-trivial objects, such as closures, +// that need to be freed explicitly, at program exit. This test verifies that the +// autodiff runtime destroys and deallocates any such objects. + +extension LifetimeTracked: AdditiveArithmetic { + public static var zero: LifetimeTracked { fatalError() } + public static func + (lhs: LifetimeTracked, rhs: LifetimeTracked) -> LifetimeTracked {fatalError()} + public static func - (lhs: LifetimeTracked, rhs: LifetimeTracked) -> LifetimeTracked {fatalError()} +} + +extension LifetimeTracked: Differentiable { + public typealias TangentVector = LifetimeTracked + public func move(by: LifetimeTracked) {fatalError()} +} + +extension LifetimeTracked { + // The original differentiable callee. + func callee(_: Float) -> Float { 42 } + + // The callee differential (pullback in this case), that is + // captured in the context object allocated on the heap in the + // presence of loops. + // + // If the autodiff runtime does not free this callee differential + // properly, the `LifetimeTracked` instance that it captures will + // also not be freed and we will have a detectable memory leak. + @derivative(of: callee, wrt: (self, f)) + func calleeDifferential(f: Float) -> (value: Float, pullback: (Float) -> (LifetimeTracked, Float)) { + return ( + value: f, + pullback: { x in (self, x) } + ) + } +} + +@differentiable(reverse) +func f(ltti: LifetimeTracked) -> Float { + for _ in 0..<1 { + } + return ltti.callee(0xDEADBEEF) +} + +var Tests = TestSuite("CalleeDifferentialLeakTest") + +Tests.test("dontLeakCalleeDifferential") { + do { + let ltti = LifetimeTracked(0xDEADBEEF) + let _ = valueWithPullback(at: ltti, of: f) + } + expectEqual(0, LifetimeTracked.instances) +} + +runAllTests() \ No newline at end of file From 91bafd4b8c12d3caefd4c81b66d29348787e4c30 Mon Sep 17 00:00:00 2001 From: Kshitij Date: Wed, 16 Aug 2023 11:15:35 -0700 Subject: [PATCH 2/3] Reintegrates BumpPtrAllocator in AutoDiffLinearMapContext This commit builds on top of a previous commit that fixed memory leak issues in the autodiff loop context allocation builtins. The previous commit allocated loop contexts in a box, which is known to have memory scaling issues. This commit modifies AutoDiffLinearMapContext to perform memory allocation using the BumpPtrAllocator, like it used to, while introducing new bookkeeping logic that enables AutoDiffLinearMapContext to prevent memory leaks, unlike before. --- stdlib/public/runtime/AutoDiffSupport.cpp | 73 ++++++++----------- stdlib/public/runtime/AutoDiffSupport.h | 86 +++++++++++++---------- 2 files changed, 79 insertions(+), 80 deletions(-) diff --git a/stdlib/public/runtime/AutoDiffSupport.cpp b/stdlib/public/runtime/AutoDiffSupport.cpp index c00feeef4da4c..1f81861d9501d 100644 --- a/stdlib/public/runtime/AutoDiffSupport.cpp +++ b/stdlib/public/runtime/AutoDiffSupport.cpp @@ -21,13 +21,7 @@ using namespace llvm; SWIFT_CC(swift) static void destroyLinearMapContext(SWIFT_CONTEXT HeapObject *obj) { - auto *linearMapContext = static_cast(obj); - - for (auto *heapObjectPtr : linearMapContext->getAllocatedHeapObjects()) { - swift_release(heapObjectPtr); - } - - linearMapContext->~AutoDiffLinearMapContext(); + static_cast(obj)->~AutoDiffLinearMapContext(); free(obj); } @@ -50,53 +44,48 @@ static FullMetadata linearMapContextHeapMetadata = { }; AutoDiffLinearMapContext::AutoDiffLinearMapContext( - OpaqueValue *const topLevelLinearMapContextProjection) + const Metadata *topLevelLinearMapContextMetadata) : HeapObject(&linearMapContextHeapMetadata) { - this->topLevelLinearMapContextProjection = topLevelLinearMapContextProjection; + allocatedContextObjects.push_back(AllocatedContextObjectRecord{ + topLevelLinearMapContextMetadata, projectTopLevelSubcontext()}); } -AutoDiffLinearMapContext *swift::swift_autoDiffCreateLinearMapContext( - const Metadata *topLevelLinearMapContextMetadata) { - // Linear map context metadata must have non-null value witnesses - assert(topLevelLinearMapContextMetadata->getValueWitnesses()); - - // Allocate a box for the top-level linear map context - auto [topLevelContextHeapObjectPtr, toplevelContextProjection] = - swift_allocBox(topLevelLinearMapContextMetadata); - - // Create a linear map context object that stores the projection - // for the top level context - auto linearMapContext = - new AutoDiffLinearMapContext(toplevelContextProjection); +void *AutoDiffLinearMapContext::projectTopLevelSubcontext() const { + auto offset = alignTo(sizeof(AutoDiffLinearMapContext), + alignof(AutoDiffLinearMapContext)); + return const_cast(reinterpret_cast(this) + + offset); +} - // Stash away the `HeapObject` pointer for the allocated context - // for proper "release" during clean up. - linearMapContext->storeAllocatedHeapObjectPtr(topLevelContextHeapObjectPtr); +void *AutoDiffLinearMapContext::allocateSubcontext( + const Metadata *contextObjectMetadata) { + auto size = contextObjectMetadata->vw_size(); + auto align = contextObjectMetadata->vw_alignment(); + auto *contextObjectPtr = allocator.Allocate(size, align); + allocatedContextObjects.push_back( + AllocatedContextObjectRecord{contextObjectMetadata, contextObjectPtr}); + return contextObjectPtr; +} - // Return the newly created linear map context object - return linearMapContext; +AutoDiffLinearMapContext *swift::swift_autoDiffCreateLinearMapContext( + const Metadata *topLevelLinearMapContextMetadata) { + auto topLevelLinearMapContextSize = + topLevelLinearMapContextMetadata->vw_size(); + auto allocationSize = alignTo(sizeof(AutoDiffLinearMapContext), + alignof(AutoDiffLinearMapContext)) + + topLevelLinearMapContextSize; + auto *buffer = (AutoDiffLinearMapContext *)malloc(allocationSize); + return ::new (buffer) + AutoDiffLinearMapContext(topLevelLinearMapContextMetadata); } void *swift::swift_autoDiffProjectTopLevelSubcontext( AutoDiffLinearMapContext *linearMapContext) { - return static_cast( - linearMapContext->getTopLevelLinearMapContextProjection()); + return static_cast(linearMapContext->projectTopLevelSubcontext()); } void *swift::swift_autoDiffAllocateSubcontext( AutoDiffLinearMapContext *linearMapContext, const Metadata *linearMapSubcontextMetadata) { - // Linear map context metadata must have non-null value witnesses - assert(linearMapSubcontextMetadata->getValueWitnesses()); - - // Allocate a box for the linear map subcontext - auto [subcontextHeapObjectPtr, subcontextProjection] = - swift_allocBox(linearMapSubcontextMetadata); - - // Stash away the `HeapObject` pointer for the allocated context - // for proper "release" during clean up. - linearMapContext->storeAllocatedHeapObjectPtr(subcontextHeapObjectPtr); - - // Return the subcontext projection - return static_cast(subcontextProjection); + return linearMapContext->allocateSubcontext(linearMapSubcontextMetadata); } diff --git a/stdlib/public/runtime/AutoDiffSupport.h b/stdlib/public/runtime/AutoDiffSupport.h index d8ecde3b2a7bc..ede01af04337c 100644 --- a/stdlib/public/runtime/AutoDiffSupport.h +++ b/stdlib/public/runtime/AutoDiffSupport.h @@ -13,58 +13,68 @@ #ifndef SWIFT_RUNTIME_AUTODIFF_SUPPORT_H #define SWIFT_RUNTIME_AUTODIFF_SUPPORT_H +#include "swift/ABI/Metadata.h" #include "swift/Runtime/Config.h" #include "swift/Runtime/HeapObject.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Allocator.h" namespace swift { /// A data structure responsible for efficiently allocating closure contexts for /// linear maps such as pullbacks, including recursive branching trace enum /// case payloads. class AutoDiffLinearMapContext : public HeapObject { + /// A simple wrapper around a context object allocated by the + /// `AutoDiffLinearMapContext` type. This type knows all the "physical" + /// properties and behavior of the allocated context object by way of + /// storing the allocated type's `TypeMetadata`. It uses this information + /// to ensure that the allocated context object is destroyed/deinitialized + /// properly, upon its own destruction. + class [[nodiscard]] AllocatedContextObjectRecord final { + const Metadata *contextObjectMetadata; + OpaqueValue *contextObjectPtr; + + public: + AllocatedContextObjectRecord(const Metadata *contextObjectMetadata, + OpaqueValue *contextObjectPtr) + : contextObjectMetadata(contextObjectMetadata), + contextObjectPtr(contextObjectPtr) {} + + AllocatedContextObjectRecord(const Metadata *contextObjectMetadata, + void *contextObjectPtr) + : AllocatedContextObjectRecord( + contextObjectMetadata, + static_cast(contextObjectPtr)) {} + + ~AllocatedContextObjectRecord() { + contextObjectMetadata->vw_destroy(contextObjectPtr); + } + + size_t size() const { return contextObjectMetadata->vw_size(); } + + size_t align() const { return contextObjectMetadata->vw_alignment(); } + }; + private: - // TODO: Commenting out BumpPtrAllocator temporarily - // until we move away from the interim solution of allocating - // boxes for linear map contexts/subcontexts. - // - // /// The underlying allocator. - // // TODO: Use a custom allocator so that the initial slab can be - // // tail-allocated. - // llvm::BumpPtrAllocator allocator; - - /// A projection/pointer to the memory storing the - /// top-level linear map context object. - OpaqueValue *topLevelLinearMapContextProjection; - - /// Storage for `HeapObject` pointers to the linear map - /// context and subcontexts allocated for derivatives with - /// loops. - llvm::SmallVector allocatedHeapObjects; + /// The underlying allocator. + // TODO: Use a custom allocator so that the initial slab can be + // tail-allocated. + llvm::BumpPtrAllocator allocator; + + /// Storage for `AllocatedContextObjectRecord`s, corresponding to the + /// subcontext allocations performed by the type. + llvm::SmallVector allocatedContextObjects; public: /// Creates a linear map context. - AutoDiffLinearMapContext(OpaqueValue *const); - - // TODO: Commenting out BumpPtrAllocator temporarily - // until we move away from the interim solution of allocating - // boxes for linear map contexts/subcontexts. - // - // llvm::BumpPtrAllocator& getAllocator() const { - // return const_cast(this->allocator); - // } - - OpaqueValue *getTopLevelLinearMapContextProjection() const { - return this->topLevelLinearMapContextProjection; - } - - llvm::ArrayRef getAllocatedHeapObjects() const { - return this->allocatedHeapObjects; - } - - void storeAllocatedHeapObjectPtr(HeapObject *allocatedHeapObjectPtr) { - this->allocatedHeapObjects.push_back(allocatedHeapObjectPtr); - } + AutoDiffLinearMapContext(const Metadata *topLevelLinearMapContextMetadata); + + /// Returns the address of the tail-allocated top-level subcontext. + void *projectTopLevelSubcontext() const; + + /// Allocates memory for a new subcontext. + void *allocateSubcontext(const Metadata *contextObjectMetadata); }; /// Creates a linear map context with a tail-allocated top-level subcontext. From 8102e6aa4968eb355fdb76b39a4e6964035f9c36 Mon Sep 17 00:00:00 2001 From: Kshitij Date: Thu, 17 Aug 2023 12:11:33 -0700 Subject: [PATCH 3/3] Creates new autodiff loop context allocation builtins and preserves old ones This commit builds on top of a previous commit that fixed memory leak issues in the autodiff loop context allocation builtins. The previous commit changed the existing builtins. However, inorder to maintain ABI compatibility and not break other consumers of autodiff, this commit creates newly named builtins that don't leak memory and preserves the older builtins. As a result, now, code compiled with an older frontend and using a newer stdlib will continue to work. But code compiled with a newer frontend and using older stdlib will have undefined behavior. --- include/swift/AST/Builtins.def | 8 +-- include/swift/Runtime/RuntimeFunctions.def | 12 ++--- lib/AST/Builtins.cpp | 10 ++-- lib/IRGen/GenBuiltin.cpp | 13 ++--- lib/IRGen/GenCall.cpp | 11 ++-- lib/IRGen/GenCall.h | 10 ++-- lib/SIL/IR/OperandOwnership.cpp | 5 +- lib/SIL/IR/ValueOwnership.cpp | 4 +- lib/SIL/Utils/MemAccessUtils.cpp | 4 +- lib/SILGen/SILGenBuiltin.cpp | 16 +++--- .../Differentiation/VJPCloner.cpp | 12 ++--- .../AccessEnforcementReleaseSinking.cpp | 4 +- stdlib/public/runtime/AutoDiffSupport.cpp | 44 +++++++++++---- stdlib/public/runtime/AutoDiffSupport.h | 53 ++++++++++++++++--- test/AutoDiff/IRGen/runtime.swift | 15 +++--- test/AutoDiff/SILGen/autodiff_builtins.swift | 15 +++--- ...todiff_builtins_deprecated_abi_fails.swift | 20 ------- 17 files changed, 150 insertions(+), 106 deletions(-) delete mode 100644 test/AutoDiff/SILGen/autodiff_builtins_deprecated_abi_fails.swift diff --git a/include/swift/AST/Builtins.def b/include/swift/AST/Builtins.def index 7483903238faf..286ebdb1bfd6a 100644 --- a/include/swift/AST/Builtins.def +++ b/include/swift/AST/Builtins.def @@ -985,14 +985,14 @@ BUILTIN_MISC_OPERATION_WITH_SILGEN(CreateAsyncTaskInGroup, /// is a pure value and therefore we can consider it as readnone). BUILTIN_MISC_OPERATION_WITH_SILGEN(GlobalStringTablePointer, "globalStringTablePointer", "n", Special) -// autoDiffCreateLinearMapContext: (T.Type) -> Builtin.NativeObject -BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffCreateLinearMapContext, "autoDiffCreateLinearMapContext", "", Special) +// autoDiffCreateLinearMapContextWithType: (T.Type) -> Builtin.NativeObject +BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffCreateLinearMapContextWithType, "autoDiffCreateLinearMapContextWithType", "", Special) // autoDiffProjectTopLevelSubcontext: (Builtin.NativeObject) -> Builtin.RawPointer BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffProjectTopLevelSubcontext, "autoDiffProjectTopLevelSubcontext", "n", Special) -// autoDiffAllocateSubcontext: (Builtin.NativeObject, T.Type) -> Builtin.RawPointer -BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffAllocateSubcontext, "autoDiffAllocateSubcontext", "", Special) +// autoDiffAllocateSubcontextWithType: (Builtin.NativeObject, T.Type) -> Builtin.RawPointer +BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffAllocateSubcontextWithType, "autoDiffAllocateSubcontextWithType", "", Special) /// Build a Builtin.Executor value from an "ordinary" serial executor /// reference. diff --git a/include/swift/Runtime/RuntimeFunctions.def b/include/swift/Runtime/RuntimeFunctions.def index 8ed18fce638c7..5a788b02bf7a7 100644 --- a/include/swift/Runtime/RuntimeFunctions.def +++ b/include/swift/Runtime/RuntimeFunctions.def @@ -2273,9 +2273,9 @@ FUNCTION(TaskGroupDestroy, ATTRS(NoUnwind), EFFECT(Concurrency)) -// AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext(const Metadata *); -FUNCTION(AutoDiffCreateLinearMapContext, - swift_autoDiffCreateLinearMapContext, SwiftCC, +// AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContextWithType(const Metadata *); +FUNCTION(AutoDiffCreateLinearMapContextWithType, + swift_autoDiffCreateLinearMapContextWithType, SwiftCC, DifferentiationAvailability, RETURNS(RefCountedPtrTy), ARGS(TypeMetadataPtrTy), @@ -2291,9 +2291,9 @@ FUNCTION(AutoDiffProjectTopLevelSubcontext, ATTRS(NoUnwind, ArgMemOnly), EFFECT(AutoDiff)) -// void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, const Metadata *); -FUNCTION(AutoDiffAllocateSubcontext, - swift_autoDiffAllocateSubcontext, SwiftCC, +// void *swift_autoDiffAllocateSubcontextWithType(AutoDiffLinearMapContext *, const Metadata *); +FUNCTION(AutoDiffAllocateSubcontextWithType, + swift_autoDiffAllocateSubcontextWithType, SwiftCC, DifferentiationAvailability, RETURNS(Int8PtrTy), ARGS(RefCountedPtrTy, TypeMetadataPtrTy), diff --git a/lib/AST/Builtins.cpp b/lib/AST/Builtins.cpp index 45fa403ebf81a..fe168edc629bc 100644 --- a/lib/AST/Builtins.cpp +++ b/lib/AST/Builtins.cpp @@ -1608,9 +1608,9 @@ static ValueDecl *getBuildComplexEqualitySerialExecutorRef(ASTContext &ctx, static ValueDecl *getAutoDiffCreateLinearMapContext(ASTContext &ctx, Identifier id) { - return getBuiltinFunction(ctx, id, _thin, _generics(_unrestricted), - _parameters(_metatype(_typeparam(0))), - _nativeObject); + return getBuiltinFunction( + ctx, id, _thin, _generics(_unrestricted), + _parameters(_metatype(_typeparam(0))), _nativeObject); } static ValueDecl *getAutoDiffProjectTopLevelSubcontext(ASTContext &ctx, @@ -2967,13 +2967,13 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) { case BuiltinValueKind::HopToActor: return getHopToActor(Context, Id); - case BuiltinValueKind::AutoDiffCreateLinearMapContext: + case BuiltinValueKind::AutoDiffCreateLinearMapContextWithType: return getAutoDiffCreateLinearMapContext(Context, Id); case BuiltinValueKind::AutoDiffProjectTopLevelSubcontext: return getAutoDiffProjectTopLevelSubcontext(Context, Id); - case BuiltinValueKind::AutoDiffAllocateSubcontext: + case BuiltinValueKind::AutoDiffAllocateSubcontextWithType: return getAutoDiffAllocateSubcontext(Context, Id); } diff --git a/lib/IRGen/GenBuiltin.cpp b/lib/IRGen/GenBuiltin.cpp index cb0d1c1276828..5dfaebc1d0318 100644 --- a/lib/IRGen/GenBuiltin.cpp +++ b/lib/IRGen/GenBuiltin.cpp @@ -1307,9 +1307,10 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin, return; } - if (Builtin.ID == BuiltinValueKind::AutoDiffCreateLinearMapContext) { + if (Builtin.ID == BuiltinValueKind::AutoDiffCreateLinearMapContextWithType) { auto topLevelSubcontextMetaType = args.claimNext(); - out.add(emitAutoDiffCreateLinearMapContext(IGF, topLevelSubcontextMetaType) + out.add(emitAutoDiffCreateLinearMapContextWithType( + IGF, topLevelSubcontextMetaType) .getAddress()); return; } @@ -1322,13 +1323,13 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin, return; } - if (Builtin.ID == BuiltinValueKind::AutoDiffAllocateSubcontext) { + if (Builtin.ID == BuiltinValueKind::AutoDiffAllocateSubcontextWithType) { Address allocatorAddr(args.claimNext(), IGF.IGM.RefCountedStructTy, IGF.IGM.getPointerAlignment()); auto subcontextMetatype = args.claimNext(); - out.add( - emitAutoDiffAllocateSubcontext(IGF, allocatorAddr, subcontextMetatype) - .getAddress()); + out.add(emitAutoDiffAllocateSubcontextWithType(IGF, allocatorAddr, + subcontextMetatype) + .getAddress()); return; } diff --git a/lib/IRGen/GenCall.cpp b/lib/IRGen/GenCall.cpp index 23a9e537f695a..c92ac7e86887c 100644 --- a/lib/IRGen/GenCall.cpp +++ b/lib/IRGen/GenCall.cpp @@ -5479,12 +5479,12 @@ IRGenFunction::getFunctionPointerForResumeIntrinsic(llvm::Value *resume) { return fnPtr; } -Address irgen::emitAutoDiffCreateLinearMapContext( +Address irgen::emitAutoDiffCreateLinearMapContextWithType( IRGenFunction &IGF, llvm::Value *topLevelSubcontextMetatype) { topLevelSubcontextMetatype = IGF.Builder.CreateBitCast( topLevelSubcontextMetatype, IGF.IGM.TypeMetadataPtrTy); auto *call = IGF.Builder.CreateCall( - IGF.IGM.getAutoDiffCreateLinearMapContextFunctionPointer(), + IGF.IGM.getAutoDiffCreateLinearMapContextWithTypeFunctionPointer(), {topLevelSubcontextMetatype}); call->setDoesNotThrow(); call->setCallingConv(IGF.IGM.SwiftCC); @@ -5502,13 +5502,12 @@ Address irgen::emitAutoDiffProjectTopLevelSubcontext( return Address(call, IGF.IGM.Int8Ty, IGF.IGM.getPointerAlignment()); } -Address irgen::emitAutoDiffAllocateSubcontext(IRGenFunction &IGF, - Address context, - llvm::Value *subcontextMetatype) { +Address irgen::emitAutoDiffAllocateSubcontextWithType( + IRGenFunction &IGF, Address context, llvm::Value *subcontextMetatype) { subcontextMetatype = IGF.Builder.CreateBitCast(subcontextMetatype, IGF.IGM.TypeMetadataPtrTy); auto *call = IGF.Builder.CreateCall( - IGF.IGM.getAutoDiffAllocateSubcontextFunctionPointer(), + IGF.IGM.getAutoDiffAllocateSubcontextWithTypeFunctionPointer(), {context.getAddress(), subcontextMetatype}); call->setDoesNotThrow(); call->setCallingConv(IGF.IGM.SwiftCC); diff --git a/lib/IRGen/GenCall.h b/lib/IRGen/GenCall.h index e6ef58f1411d9..70879ce628d05 100644 --- a/lib/IRGen/GenCall.h +++ b/lib/IRGen/GenCall.h @@ -261,12 +261,14 @@ namespace irgen { CanSILFunctionType fnType, Explosion &result, Explosion &error); - Address - emitAutoDiffCreateLinearMapContext(IRGenFunction &IGF, - llvm::Value *topLevelSubcontextMetatype); + Address emitAutoDiffCreateLinearMapContextWithType( + IRGenFunction &IGF, llvm::Value *topLevelSubcontextMetatype); + Address emitAutoDiffProjectTopLevelSubcontext( IRGenFunction &IGF, Address context); - Address emitAutoDiffAllocateSubcontext(IRGenFunction &IGF, Address context, + + Address + emitAutoDiffAllocateSubcontextWithType(IRGenFunction &IGF, Address context, llvm::Value *subcontextMetatype); FunctionPointer getFunctionPointerForDispatchCall(IRGenModule &IGM, diff --git a/lib/SIL/IR/OperandOwnership.cpp b/lib/SIL/IR/OperandOwnership.cpp index 7fe14a6e059a4..faeb4a6141f5b 100644 --- a/lib/SIL/IR/OperandOwnership.cpp +++ b/lib/SIL/IR/OperandOwnership.cpp @@ -943,7 +943,7 @@ BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, InitializeDistributedRemoteActor) BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, InitializeNonDefaultDistributedActor) -BUILTIN_OPERAND_OWNERSHIP(PointerEscape, AutoDiffAllocateSubcontext) +BUILTIN_OPERAND_OWNERSHIP(PointerEscape, AutoDiffAllocateSubcontextWithType) BUILTIN_OPERAND_OWNERSHIP(PointerEscape, AutoDiffProjectTopLevelSubcontext) // FIXME: ConvertTaskToJob is documented as taking NativePointer. It's operand's @@ -955,8 +955,7 @@ BUILTIN_OPERAND_OWNERSHIP(BitwiseEscape, BuildComplexEqualitySerialExecutorRef) BUILTIN_OPERAND_OWNERSHIP(BitwiseEscape, BuildDefaultActorExecutorRef) BUILTIN_OPERAND_OWNERSHIP(BitwiseEscape, BuildMainActorExecutorRef) -BUILTIN_OPERAND_OWNERSHIP(TrivialUse, AutoDiffCreateLinearMapContext) - +BUILTIN_OPERAND_OWNERSHIP(TrivialUse, AutoDiffCreateLinearMapContextWithType) #undef BUILTIN_OPERAND_OWNERSHIP #define SHOULD_NEVER_VISIT_BUILTIN(ID) \ diff --git a/lib/SIL/IR/ValueOwnership.cpp b/lib/SIL/IR/ValueOwnership.cpp index 5cdc0b01c15e8..40c1bbcd05726 100644 --- a/lib/SIL/IR/ValueOwnership.cpp +++ b/lib/SIL/IR/ValueOwnership.cpp @@ -571,9 +571,9 @@ CONSTANT_OWNERSHIP_BUILTIN(None, InitializeDefaultActor) CONSTANT_OWNERSHIP_BUILTIN(None, DestroyDefaultActor) CONSTANT_OWNERSHIP_BUILTIN(None, InitializeDistributedRemoteActor) CONSTANT_OWNERSHIP_BUILTIN(None, InitializeNonDefaultDistributedActor) -CONSTANT_OWNERSHIP_BUILTIN(Owned, AutoDiffCreateLinearMapContext) +CONSTANT_OWNERSHIP_BUILTIN(Owned, AutoDiffCreateLinearMapContextWithType) CONSTANT_OWNERSHIP_BUILTIN(None, AutoDiffProjectTopLevelSubcontext) -CONSTANT_OWNERSHIP_BUILTIN(None, AutoDiffAllocateSubcontext) +CONSTANT_OWNERSHIP_BUILTIN(None, AutoDiffAllocateSubcontextWithType) CONSTANT_OWNERSHIP_BUILTIN(None, GetCurrentExecutor) CONSTANT_OWNERSHIP_BUILTIN(None, ResumeNonThrowingContinuationReturning) CONSTANT_OWNERSHIP_BUILTIN(None, ResumeThrowingContinuationReturning) diff --git a/lib/SIL/Utils/MemAccessUtils.cpp b/lib/SIL/Utils/MemAccessUtils.cpp index 6cf1a89f14f3a..b672f69550c50 100644 --- a/lib/SIL/Utils/MemAccessUtils.cpp +++ b/lib/SIL/Utils/MemAccessUtils.cpp @@ -2567,8 +2567,8 @@ static void visitBuiltinAddress(BuiltinInst *builtin, case BuiltinValueKind::CancelAsyncTask: case BuiltinValueKind::CreateAsyncTask: case BuiltinValueKind::CreateAsyncTaskInGroup: - case BuiltinValueKind::AutoDiffCreateLinearMapContext: - case BuiltinValueKind::AutoDiffAllocateSubcontext: + case BuiltinValueKind::AutoDiffCreateLinearMapContextWithType: + case BuiltinValueKind::AutoDiffAllocateSubcontextWithType: case BuiltinValueKind::InitializeDefaultActor: case BuiltinValueKind::InitializeDistributedRemoteActor: case BuiltinValueKind::InitializeNonDefaultDistributedActor: diff --git a/lib/SILGen/SILGenBuiltin.cpp b/lib/SILGen/SILGenBuiltin.cpp index d3b3833aef385..e07bc18b0868a 100644 --- a/lib/SILGen/SILGenBuiltin.cpp +++ b/lib/SILGen/SILGenBuiltin.cpp @@ -1708,16 +1708,15 @@ static ManagedValue emitBuiltinHopToActor(SILGenFunction &SGF, SILLocation loc, return ManagedValue::forObjectRValueWithoutOwnership(SGF.emitEmptyTuple(loc)); } -static ManagedValue emitBuiltinAutoDiffCreateLinearMapContext( +static ManagedValue emitBuiltinAutoDiffCreateLinearMapContextWithType( SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs, ArrayRef args, SGFContext C) { ASTContext &ctx = SGF.getASTContext(); auto *builtinApply = SGF.B.createBuiltin( loc, - ctx.getIdentifier( - getBuiltinName(BuiltinValueKind::AutoDiffCreateLinearMapContext)), - SILType::getNativeObjectType(ctx), - subs, + ctx.getIdentifier(getBuiltinName( + BuiltinValueKind::AutoDiffCreateLinearMapContextWithType)), + SILType::getNativeObjectType(ctx), subs, /*args*/ {args[0].getValue()}); return SGF.emitManagedRValueWithCleanup(builtinApply); } @@ -1736,16 +1735,15 @@ static ManagedValue emitBuiltinAutoDiffProjectTopLevelSubcontext( return ManagedValue::forObjectRValueWithoutOwnership(builtinApply); } -static ManagedValue emitBuiltinAutoDiffAllocateSubcontext( +static ManagedValue emitBuiltinAutoDiffAllocateSubcontextWithType( SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs, ArrayRef args, SGFContext C) { ASTContext &ctx = SGF.getASTContext(); auto *builtinApply = SGF.B.createBuiltin( loc, ctx.getIdentifier( - getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontext)), - SILType::getRawPointerType(ctx), - subs, + getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontextWithType)), + SILType::getRawPointerType(ctx), subs, /*args*/ {args[0].borrow(SGF, loc).getValue(), args[1].getValue()}); return ManagedValue::forObjectRValueWithoutOwnership(builtinApply); } diff --git a/lib/SILOptimizer/Differentiation/VJPCloner.cpp b/lib/SILOptimizer/Differentiation/VJPCloner.cpp index 90084d4b56026..50098dda2c173 100644 --- a/lib/SILOptimizer/Differentiation/VJPCloner.cpp +++ b/lib/SILOptimizer/Differentiation/VJPCloner.cpp @@ -131,8 +131,8 @@ class VJPCloner::Implementation final // Create an context. pullbackContextValue = Builder.createBuiltin( original->getLocation(), - getASTContext().getIdentifier( - getBuiltinName(BuiltinValueKind::AutoDiffCreateLinearMapContext)), + getASTContext().getIdentifier(getBuiltinName( + BuiltinValueKind::AutoDiffCreateLinearMapContextWithType)), SILType::getNativeObjectType(getASTContext()), SubstitutionMap(), {pbTupleMetatype}); borrowedPullbackContextValue = Builder.createBeginBorrow( @@ -156,8 +156,8 @@ class VJPCloner::Implementation final return builtinAutoDiffAllocateSubcontextGenericSignature; auto &ctx = getASTContext(); auto *decl = cast(getBuiltinValueDecl( - ctx, ctx.getIdentifier( - getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontext)))); + ctx, ctx.getIdentifier(getBuiltinName( + BuiltinValueKind::AutoDiffAllocateSubcontextWithType)))); builtinAutoDiffAllocateSubcontextGenericSignature = decl->getGenericSignature(); assert(builtinAutoDiffAllocateSubcontextGenericSignature); @@ -1085,8 +1085,8 @@ EnumInst *VJPCloner::Implementation::buildPredecessorEnumValue( auto rawBufferValue = builder.createBuiltin( loc, - getASTContext().getIdentifier( - getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontext)), + getASTContext().getIdentifier(getBuiltinName( + BuiltinValueKind::AutoDiffAllocateSubcontextWithType)), rawPtrType, SubstitutionMap(), {borrowedPullbackContextValue, pbTupleMetatype}); diff --git a/lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp b/lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp index 01069fe3e2774..20359097fcf3d 100644 --- a/lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp +++ b/lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp @@ -146,7 +146,7 @@ static bool isBarrier(SILInstruction *inst) { case BuiltinValueKind::COWBufferForReading: case BuiltinValueKind::GetCurrentAsyncTask: case BuiltinValueKind::GetCurrentExecutor: - case BuiltinValueKind::AutoDiffCreateLinearMapContext: + case BuiltinValueKind::AutoDiffCreateLinearMapContextWithType: case BuiltinValueKind::EndAsyncLet: case BuiltinValueKind::EndAsyncLetLifetime: case BuiltinValueKind::CreateTaskGroup: @@ -199,7 +199,7 @@ static bool isBarrier(SILInstruction *inst) { case BuiltinValueKind::ResumeThrowingContinuationReturning: case BuiltinValueKind::ResumeThrowingContinuationThrowing: case BuiltinValueKind::AutoDiffProjectTopLevelSubcontext: - case BuiltinValueKind::AutoDiffAllocateSubcontext: + case BuiltinValueKind::AutoDiffAllocateSubcontextWithType: case BuiltinValueKind::AddressOfBorrowOpaque: case BuiltinValueKind::UnprotectedAddressOfBorrowOpaque: return true; diff --git a/stdlib/public/runtime/AutoDiffSupport.cpp b/stdlib/public/runtime/AutoDiffSupport.cpp index 1f81861d9501d..372bcdac9447f 100644 --- a/stdlib/public/runtime/AutoDiffSupport.cpp +++ b/stdlib/public/runtime/AutoDiffSupport.cpp @@ -43,6 +43,10 @@ static FullMetadata linearMapContextHeapMetadata = { } }; +AutoDiffLinearMapContext::AutoDiffLinearMapContext() + : HeapObject(&linearMapContextHeapMetadata) { +} + AutoDiffLinearMapContext::AutoDiffLinearMapContext( const Metadata *topLevelLinearMapContextMetadata) : HeapObject(&linearMapContextHeapMetadata) { @@ -51,10 +55,14 @@ AutoDiffLinearMapContext::AutoDiffLinearMapContext( } void *AutoDiffLinearMapContext::projectTopLevelSubcontext() const { - auto offset = alignTo(sizeof(AutoDiffLinearMapContext), - alignof(AutoDiffLinearMapContext)); - return const_cast(reinterpret_cast(this) + - offset); + auto offset = alignTo( + sizeof(AutoDiffLinearMapContext), alignof(AutoDiffLinearMapContext)); + return const_cast( + reinterpret_cast(this) + offset); +} + +void *AutoDiffLinearMapContext::allocate(size_t size) { + return allocator.Allocate(size, alignof(AutoDiffLinearMapContext)); } void *AutoDiffLinearMapContext::allocateSubcontext( @@ -68,7 +76,27 @@ void *AutoDiffLinearMapContext::allocateSubcontext( } AutoDiffLinearMapContext *swift::swift_autoDiffCreateLinearMapContext( + size_t topLevelLinearMapStructSize) { + auto allocationSize = alignTo( + sizeof(AutoDiffLinearMapContext), alignof(AutoDiffLinearMapContext)) + + topLevelLinearMapStructSize; + auto *buffer = (AutoDiffLinearMapContext *)malloc(allocationSize); + return ::new (buffer) AutoDiffLinearMapContext; +} + +void *swift::swift_autoDiffProjectTopLevelSubcontext( + AutoDiffLinearMapContext *linearMapContext) { + return static_cast(linearMapContext->projectTopLevelSubcontext()); +} + +void *swift::swift_autoDiffAllocateSubcontext( + AutoDiffLinearMapContext *allocator, size_t size) { + return allocator->allocate(size); +} + +AutoDiffLinearMapContext *swift::swift_autoDiffCreateLinearMapContextWithType( const Metadata *topLevelLinearMapContextMetadata) { + assert(topLevelLinearMapContextMetadata->getValueWitnesses() != nullptr); auto topLevelLinearMapContextSize = topLevelLinearMapContextMetadata->vw_size(); auto allocationSize = alignTo(sizeof(AutoDiffLinearMapContext), @@ -79,13 +107,9 @@ AutoDiffLinearMapContext *swift::swift_autoDiffCreateLinearMapContext( AutoDiffLinearMapContext(topLevelLinearMapContextMetadata); } -void *swift::swift_autoDiffProjectTopLevelSubcontext( - AutoDiffLinearMapContext *linearMapContext) { - return static_cast(linearMapContext->projectTopLevelSubcontext()); -} - -void *swift::swift_autoDiffAllocateSubcontext( +void *swift::swift_autoDiffAllocateSubcontextWithType( AutoDiffLinearMapContext *linearMapContext, const Metadata *linearMapSubcontextMetadata) { + assert(linearMapSubcontextMetadata->getValueWitnesses() != nullptr); return linearMapContext->allocateSubcontext(linearMapSubcontextMetadata); } diff --git a/stdlib/public/runtime/AutoDiffSupport.h b/stdlib/public/runtime/AutoDiffSupport.h index ede01af04337c..b3135ef794628 100644 --- a/stdlib/public/runtime/AutoDiffSupport.h +++ b/stdlib/public/runtime/AutoDiffSupport.h @@ -13,9 +13,9 @@ #ifndef SWIFT_RUNTIME_AUTODIFF_SUPPORT_H #define SWIFT_RUNTIME_AUTODIFF_SUPPORT_H +#include "swift/Runtime/HeapObject.h" #include "swift/ABI/Metadata.h" #include "swift/Runtime/Config.h" -#include "swift/Runtime/HeapObject.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/Allocator.h" @@ -48,7 +48,19 @@ class AutoDiffLinearMapContext : public HeapObject { static_cast(contextObjectPtr)) {} ~AllocatedContextObjectRecord() { - contextObjectMetadata->vw_destroy(contextObjectPtr); + if (contextObjectMetadata != nullptr && contextObjectPtr != nullptr) { + contextObjectMetadata->vw_destroy(contextObjectPtr); + } + } + + AllocatedContextObjectRecord(const AllocatedContextObjectRecord &) = delete; + + AllocatedContextObjectRecord( + AllocatedContextObjectRecord &&other) noexcept { + this->contextObjectMetadata = other.contextObjectMetadata; + this->contextObjectPtr = other.contextObjectPtr; + other.contextObjectMetadata = nullptr; + other.contextObjectPtr = nullptr; } size_t size() const { return contextObjectMetadata->vw_size(); } @@ -67,27 +79,54 @@ class AutoDiffLinearMapContext : public HeapObject { llvm::SmallVector allocatedContextObjects; public: - /// Creates a linear map context. + /// DEPRECATED - Use overloaded constructor taking a `const Metadata *` + /// parameter instead. This constructor might be removed as it leads to memory + /// leaks. + AutoDiffLinearMapContext(); + AutoDiffLinearMapContext(const Metadata *topLevelLinearMapContextMetadata); /// Returns the address of the tail-allocated top-level subcontext. void *projectTopLevelSubcontext() const; + /// Allocates memory for a new subcontext. + /// + /// DEPRECATED - Use `allocateSubcontext` instead. This + /// method might be removed as it leads to memory leaks. + void *allocate(size_t size); + /// Allocates memory for a new subcontext. void *allocateSubcontext(const Metadata *contextObjectMetadata); }; /// Creates a linear map context with a tail-allocated top-level subcontext. +/// +/// DEPRECATED - Use `swift_autoDiffCreateLinearMapContextWithType` instead. +/// This builtin might be removed as it leads to memory leaks. SWIFT_RUNTIME_EXPORT SWIFT_CC(swift) - AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext( - const Metadata *topLevelLinearMapContextMetadata); +AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext( + size_t topLevelSubcontextSize); /// Returns the address of the tail-allocated top-level subcontext. SWIFT_RUNTIME_EXPORT SWIFT_CC(swift) void *swift_autoDiffProjectTopLevelSubcontext(AutoDiffLinearMapContext *); /// Allocates memory for a new subcontext. -SWIFT_RUNTIME_EXPORT SWIFT_CC(swift) void *swift_autoDiffAllocateSubcontext( - AutoDiffLinearMapContext *, const Metadata *linearMapSubcontextMetadata); +/// +/// DEPRECATED - Use `swift_autoDiffAllocateSubcontextWithType` instead. This +/// builtin might be removed as it leads to memory leaks. +SWIFT_RUNTIME_EXPORT SWIFT_CC(swift) +void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, size_t size); + +/// Creates a linear map context with a tail-allocated top-level subcontext. +SWIFT_RUNTIME_EXPORT SWIFT_CC(swift) + AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContextWithType( + const Metadata *topLevelLinearMapContextMetadata); + +/// Allocates memory for a new subcontext. +SWIFT_RUNTIME_EXPORT + SWIFT_CC(swift) void *swift_autoDiffAllocateSubcontextWithType( + AutoDiffLinearMapContext *, + const Metadata *linearMapSubcontextMetadata); } // namespace swift #endif /* SWIFT_RUNTIME_AUTODIFF_SUPPORT_H */ diff --git a/test/AutoDiff/IRGen/runtime.swift b/test/AutoDiff/IRGen/runtime.swift index a4687d4108e89..8a753356212a8 100644 --- a/test/AutoDiff/IRGen/runtime.swift +++ b/test/AutoDiff/IRGen/runtime.swift @@ -1,18 +1,19 @@ // RUN: %target-swift-frontend -parse-stdlib %s -emit-ir | %FileCheck %s import Swift +import _Differentiation -@_silgen_name("test_context_builtins") -func test_context_builtins(t: T) { - let context = Builtin.autoDiffCreateLinearMapContext(T.self) +@_silgen_name("test_context_builtins_with_type") +func test_context_builtins_with_type(t: T) { + let context = Builtin.autoDiffCreateLinearMapContextWithType(T.self) let topLevelSubctxAddr = Builtin.autoDiffProjectTopLevelSubcontext(context) UnsafeMutableRawPointer(topLevelSubctxAddr).storeBytes(of: t, as: T.self) - let newBuffer = Builtin.autoDiffAllocateSubcontext(context, T.self) + let newBuffer = Builtin.autoDiffAllocateSubcontextWithType(context, T.self) UnsafeMutableRawPointer(newBuffer).storeBytes(of: t, as: T.self) } -// CHECK-LABEL: define{{.*}}@test_context_builtins(ptr noalias nocapture %0, ptr %T) +// CHECK-LABEL: define{{.*}}@test_context_builtins_with_type(ptr noalias nocapture %0, ptr %T) // CHECK: entry: -// CHECK: [[CTX:%.*]] = call swiftcc ptr @swift_autoDiffCreateLinearMapContext(ptr %T) +// CHECK: [[CTX:%.*]] = call swiftcc ptr @swift_autoDiffCreateLinearMapContextWithType(ptr %T) // CHECK: call swiftcc ptr @swift_autoDiffProjectTopLevelSubcontext(ptr [[CTX]]) -// CHECK: [[BUF:%.*]] = call swiftcc ptr @swift_autoDiffAllocateSubcontext(ptr [[CTX]], ptr %T) +// CHECK: [[BUF:%.*]] = call swiftcc ptr @swift_autoDiffAllocateSubcontextWithType(ptr [[CTX]], ptr %T) diff --git a/test/AutoDiff/SILGen/autodiff_builtins.swift b/test/AutoDiff/SILGen/autodiff_builtins.swift index 7f21e1f8dacce..b369deec7809c 100644 --- a/test/AutoDiff/SILGen/autodiff_builtins.swift +++ b/test/AutoDiff/SILGen/autodiff_builtins.swift @@ -84,19 +84,20 @@ func applyDerivative_f1_vjp(t0: T) -> (T // CHECK: copy_addr [take] [[D_RESULT_BUFFER_0_FOR_LOAD]] to [init] [[ORIG_RESULT_OUT_PARAM]] // CHECK: return [[PULLBACK]] -@_silgen_name("test_context_builtins") -func test_context_builtins(t: T) { - let context = Builtin.autoDiffCreateLinearMapContext(T.self) + +@_silgen_name("test_context_builtins_with_type") +func test_context_builtins_with_type(t: T) { + let context = Builtin.autoDiffCreateLinearMapContextWithType(T.self) let topLevelSubctxAddr = Builtin.autoDiffProjectTopLevelSubcontext(context) UnsafeMutableRawPointer(topLevelSubctxAddr).storeBytes(of: t, as: T.self) - let newBuffer = Builtin.autoDiffAllocateSubcontext(context, T.self) + let newBuffer = Builtin.autoDiffAllocateSubcontextWithType(context, T.self) UnsafeMutableRawPointer(newBuffer).storeBytes(of: t, as: T.self) } -// CHECK-LABEL: sil{{.*}}@test_context_builtins : $@convention(thin) (@in_guaranteed T) -> () { +// CHECK-LABEL: sil{{.*}}@test_context_builtins_with_type : $@convention(thin) (@in_guaranteed T) -> () { // CHECK: bb0({{%.*}} : $*T): -// CHECK: [[CTX:%.*]] = builtin "autoDiffCreateLinearMapContext"({{%.*}} : $@thick T.Type) : $Builtin.NativeObject // users: {{.*}} +// CHECK: [[CTX:%.*]] = builtin "autoDiffCreateLinearMapContextWithType"({{%.*}} : $@thick T.Type) : $Builtin.NativeObject // users: {{.*}} // CHECK: [[BORROWED_CTX:%.*]] = begin_borrow [lexical] [[CTX]] : $Builtin.NativeObject // users: {{.*}} // CHECK: [[BUF:%.*]] = builtin "autoDiffProjectTopLevelSubcontext"([[BORROWED_CTX]] : $Builtin.NativeObject) : $Builtin.RawPointer // users: {{.*}} -// CHECK: [[BUF:%.*]] = builtin "autoDiffAllocateSubcontext"([[BORROWED_CTX]] : $Builtin.NativeObject, {{.*}} : $@thick T.Type) : $Builtin.RawPointer // users: {{.*}} +// CHECK: [[BUF:%.*]] = builtin "autoDiffAllocateSubcontextWithType"([[BORROWED_CTX]] : $Builtin.NativeObject, {{.*}} : $@thick T.Type) : $Builtin.RawPointer // users: {{.*}} // CHECK: destroy_value [[CTX]] diff --git a/test/AutoDiff/SILGen/autodiff_builtins_deprecated_abi_fails.swift b/test/AutoDiff/SILGen/autodiff_builtins_deprecated_abi_fails.swift deleted file mode 100644 index 5d6b756a7eb91..0000000000000 --- a/test/AutoDiff/SILGen/autodiff_builtins_deprecated_abi_fails.swift +++ /dev/null @@ -1,20 +0,0 @@ -// RUN: %target-swift-frontend -emit-silgen -parse-stdlib %s -// XFAIL: * - -// SIL code with older, deprecated autodiff ABI for linear map context -// allocation should fail, complaining about receiving a 'Builtin.Word' when -// a 'T.Type' is expected. - -struct ExamplePullbackStruct { - var pb0: (T.TangentVector) -> T.TangentVector -} - -@_silgen_name("test_context_builtins_bad") -func test_context_builtins_bad() { - let pbStruct = ExamplePullbackStruct(pb0: { $0 }) - let context = Builtin.autoDiffCreateLinearMapContext(Builtin.sizeof(type(of: pbStruct))) - let topLevelSubctxAddr = Builtin.autoDiffProjectTopLevelSubcontext(context) - UnsafeMutableRawPointer(topLevelSubctxAddr).storeBytes(of: pbStruct, as: type(of: pbStruct)) - let newBuffer = Builtin.autoDiffAllocateSubcontext(context, Builtin.sizeof(type(of: pbStruct))) - UnsafeMutableRawPointer(newBuffer).storeBytes(of: pbStruct, as: type(of: pbStruct)) -}