[AutoDiff] Fixes memory leaks in autodiff linear map context allocation builtins #67944

jkshtj · 2023-08-15T18:14:24Z

When the differentiating a function containing loops, we allocate a linear map 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 PR fixes the autodiff linear map context allocation builtins to destroy and deallocate any such objects.

Fixes #67323

jkshtj · 2023-08-15T18:15:30Z

@asl @BradLarson Could I get a review on this?

BradLarson · 2023-08-15T18:33:08Z

@swift-ci Please test.

asl · 2023-08-15T19:55:28Z

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 *);


So, there will be no link errors for things compiled with old compiler and linked with new runtime and vice versa?

Ah, it slipped my mind and I meant to ask this. I wasn't able to figure out how to make this a link time error.

The easiest way: rename functions :) Hard way: invent some backward compatibility scheme. @rxwei any suggestions? I think we do not care about ABI stability for now, isn't it?

Haven't made a decision on this in the latest revision. Waiting on @rxwei's input. If breaking ABI is going to be a concern we could deprecate existing builtins instead of completely removing them.

Just to be safe, we should add a new runtime function without modifying the signature of the old one.

Should we deprecate the old ones as well? They do lead to memory leaks.

What do you mean by deprecate? We just don't call it

I mean, this is not an API. There is not a real way to deprecate it.

@jkshtj I believe what @rxwei means is as follows:

Add new functions, say swift_autoDiffCreateLinearMapContext2 etc. with the desired functionality

Keep old functions as-is

As a result: old code would still be able to run with the new runtime. The behavior will be the same, well, it will leak. The new code will call that new runtime functions and will be leak-less.

test/AutoDiff/SILGen/autodiff_builtins.swift

test/AutoDiff/stdlib/callee_differential_not_leaked_in_func_with_loops.swift

jkshtj · 2023-08-16T18:29:04Z

@asl as per offline discussion w/ @BradLarson, since we're already aware of the memory scaling issues with allocating loop context objects directly on the heap I have created a new revision to reintegrate BumpPtrAllocator in AutodiffLinearMapContext while still preventing memory leaks.

asl · 2023-08-16T19:40:45Z

since we're already aware of the #34886 with allocating loop context objects directly on the heap

@jkshtj Well, this was quite different case, as the chain of boxes were explicitly created ;) Anyway, as we already discussed, as soon as you had the initial implementation on top of the boxes, it was possible to change the underlying implementation to something more lightweight, as we certainly did not want to reference count the pullback tuples.

asl · 2023-08-16T19:41:10Z

@swift-ci please test

Fixes swiftlang#67323

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.

BradLarson · 2023-08-18T17:20:45Z

@swift-ci Please test

asl · 2023-08-23T22:42:55Z

@swift-ci please test windows

asl · 2023-08-23T22:43:07Z

include/swift/AST/Builtins.def

@@ -994,6 +994,12 @@ BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffProjectTopLevelSubcontext, "autoDiffP
 // autoDiffAllocateSubcontext: (Builtin.NativeObject, Builtin.Word) -> Builtin.RawPointer
 BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffAllocateSubcontext, "autoDiffAllocateSubcontext", "", Special)

+// autoDiffCreateLinearMapContext2: (T.Type) -> Builtin.NativeObject


Some bike-shedding: maybe rename these to autoDiffCreateLinearMapContextWithType etc.?

lib/AST/Builtins.cpp

lib/IRGen/GenCall.cpp

asl · 2023-08-23T22:45:27Z

lib/IRGen/GenCall.cpp

@@ -5500,8 +5501,9 @@ 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,


Do we need to keep these functions after all? We are not going to call these builtins from the compiler.

asl · 2023-08-23T22:47:20Z

lib/AST/Builtins.cpp

@@ -1625,6 +1625,20 @@ static ValueDecl *getAutoDiffAllocateSubcontext(ASTContext &ctx,
      ctx.TheRawPointerType);
 }

+static ValueDecl *getAutoDiffCreateLinearMapContext2(ASTContext &ctx,


As the functions above with old builtins are never called from the compiler, I'd steal their names for the helpers that call new builtins

I think it'll be better to keep the helpers symmetric with the names of the builtins. Otherwise autoDiffProjectTopLevelSubcontext will be the odd one out in different places.

…ld 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.

BradLarson · 2023-08-24T15:15:48Z

@swift-ci Please test

asl

Thanks!

Resolve conflicts introduced in #67944 as follows: ``` diff --git a/include/swift/Runtime/RuntimeFunctions.def b/include/swift/Runtime/RuntimeFunctions.def index 3c973b5884b..44cde707d17 100644 --- a/include/swift/Runtime/RuntimeFunctions.def +++ b/include/swift/Runtime/RuntimeFunctions.def @@ -2537,16 +2537,10 @@ FUNCTION(AutoDiffCreateLinearMapContextWithType, swift_autoDiffCreateLinearMapContextWithType, SwiftCC, DifferentiationAvailability, RETURNS(RefCountedPtrTy), -<<<<<<< HEAD - ARGS(SizeTy), + ARGS(TypeMetadataPtrTy), ATTRS(NoUnwind), EFFECT(AutoDiff), MEMEFFECTS(ArgMemOnly)) -======= - ARGS(TypeMetadataPtrTy), - ATTRS(NoUnwind, ArgMemOnly), - EFFECT(AutoDiff)) ->>>>>>> public-github/main // void *swift_autoDiffProjectTopLevelSubcontext(AutoDiffLinearMapContext *); FUNCTION(AutoDiffProjectTopLevelSubcontext, @@ -2563,16 +2557,10 @@ FUNCTION(AutoDiffAllocateSubcontextWithType, swift_autoDiffAllocateSubcontextWithType, SwiftCC, DifferentiationAvailability, RETURNS(Int8PtrTy), -<<<<<<< HEAD - ARGS(RefCountedPtrTy, SizeTy), + ARGS(RefCountedPtrTy, TypeMetadataPtrTy), ATTRS(NoUnwind), EFFECT(AutoDiff), MEMEFFECTS(ArgMemOnly)) -======= - ARGS(RefCountedPtrTy, TypeMetadataPtrTy), - ATTRS(NoUnwind, ArgMemOnly), - EFFECT(AutoDiff)) ->>>>>>> public-github/main ```

jkshtj requested review from hborla, slavapestov and xedin as code owners August 15, 2023 18:14

jkshtj changed the title ~~Fixes memory leaks in autodiff linear map context allocation builtins~~ [Autodiff] Fixes memory leaks in autodiff linear map context allocation builtins Aug 15, 2023

asl reviewed Aug 15, 2023

View reviewed changes

test/AutoDiff/SILGen/autodiff_builtins.swift Outdated Show resolved Hide resolved

asl reviewed Aug 15, 2023

View reviewed changes

test/AutoDiff/stdlib/callee_differential_not_leaked_in_func_with_loops.swift Outdated Show resolved Hide resolved

asl reviewed Aug 15, 2023

View reviewed changes

test/AutoDiff/stdlib/callee_differential_not_leaked_in_func_with_loops.swift Show resolved Hide resolved

jkshtj force-pushed the leaks branch from a12aeba to 83f13d4 Compare August 15, 2023 22:26

jkshtj requested a review from asl August 16, 2023 18:24

jkshtj added 2 commits August 17, 2023 13:15

Fixes memory leaks in autodiff linear map context allocation builtins

2ccc0b3

Fixes swiftlang#67323

jkshtj force-pushed the leaks branch from 3956749 to dd0dce4 Compare August 18, 2023 06:48

jkshtj requested a review from rxwei August 18, 2023 06:49

jkshtj changed the title ~~[Autodiff] Fixes memory leaks in autodiff linear map context allocation builtins~~ [AutoDiff] Fixes memory leaks in autodiff linear map context allocation builtins Aug 18, 2023

xedin removed their request for review August 23, 2023 16:30

asl reviewed Aug 23, 2023

View reviewed changes

lib/AST/Builtins.cpp Show resolved Hide resolved

asl reviewed Aug 23, 2023

View reviewed changes

lib/IRGen/GenCall.cpp Show resolved Hide resolved

asl reviewed Aug 23, 2023

View reviewed changes

lib/IRGen/GenCall.cpp Show resolved Hide resolved

asl reviewed Aug 23, 2023

View reviewed changes

asl added the AutoDiff label Aug 23, 2023

jkshtj force-pushed the leaks branch from dd0dce4 to aa9aab8 Compare August 24, 2023 15:03

jkshtj force-pushed the leaks branch from aa9aab8 to 8102e6a Compare August 24, 2023 15:07

jkshtj requested a review from asl August 24, 2023 15:10

asl approved these changes Aug 24, 2023

View reviewed changes

asl merged commit d5a3e2e into swiftlang:main Aug 24, 2023

nate-chandler mentioned this pull request Aug 24, 2023

🌲: Merge branch 'main' into rebranch. #68133

Merged

jkshtj mentioned this pull request Sep 8, 2023

[AutoDiff] Pullbacks w/ loops can segfault #68392

Closed

jkshtj mentioned this pull request Nov 8, 2023

[Autodiff] EXC_BAD_ACCESS when running a @differentiable function in the context of an XCTest. #69238

Open

jkshtj deleted the leaks branch March 13, 2024 23:23

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[AutoDiff] Fixes memory leaks in autodiff linear map context allocation builtins #67944

[AutoDiff] Fixes memory leaks in autodiff linear map context allocation builtins #67944

jkshtj commented Aug 15, 2023

jkshtj commented Aug 15, 2023

BradLarson commented Aug 15, 2023

asl Aug 15, 2023

jkshtj Aug 15, 2023

asl Aug 15, 2023

jkshtj Aug 16, 2023 •

edited

Loading

rxwei Aug 16, 2023 •

edited

Loading

jkshtj Aug 16, 2023

rxwei Aug 16, 2023

rxwei Aug 16, 2023

asl Aug 16, 2023

jkshtj commented Aug 16, 2023

asl commented Aug 16, 2023 •

edited

Loading

asl commented Aug 16, 2023

BradLarson commented Aug 18, 2023

asl commented Aug 23, 2023

asl Aug 23, 2023

asl Aug 23, 2023

jkshtj Aug 24, 2023

asl Aug 23, 2023

jkshtj Aug 24, 2023

BradLarson commented Aug 24, 2023

asl left a comment

[AutoDiff] Fixes memory leaks in autodiff linear map context allocation builtins #67944

[AutoDiff] Fixes memory leaks in autodiff linear map context allocation builtins #67944

Conversation

jkshtj commented Aug 15, 2023

jkshtj commented Aug 15, 2023

BradLarson commented Aug 15, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jkshtj Aug 16, 2023 • edited Loading

Choose a reason for hiding this comment

rxwei Aug 16, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jkshtj commented Aug 16, 2023

asl commented Aug 16, 2023 • edited Loading

asl commented Aug 16, 2023

BradLarson commented Aug 18, 2023

asl commented Aug 23, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

BradLarson commented Aug 24, 2023

asl left a comment

Choose a reason for hiding this comment

jkshtj Aug 16, 2023 •

edited

Loading

rxwei Aug 16, 2023 •

edited

Loading

asl commented Aug 16, 2023 •

edited

Loading