Skip to content

Commit d5a3e2e

Browse files
jkshtjjkshtj
authored
[AutoDiff] Fixes memory leaks in autodiff linear map context allocation builtins (#67944)
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 released explicitly. Fix the autodiff linear map context allocation builtins to correctly release such objects and not just free the memory they occupy.
1 parent 80f4588 commit d5a3e2e

File tree

17 files changed

+289
-103
lines changed

17 files changed

+289
-103
lines changed

‎include/swift/AST/Builtins.def‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -985,14 +985,14 @@ BUILTIN_MISC_OPERATION_WITH_SILGEN(CreateAsyncTaskInGroup,
985985
/// is a pure value and therefore we can consider it as readnone).
986986
BUILTIN_MISC_OPERATION_WITH_SILGEN(GlobalStringTablePointer, "globalStringTablePointer", "n", Special)
987987

988-
// autoDiffCreateLinearMapContext: (Builtin.Word) -> Builtin.NativeObject
989-
BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffCreateLinearMapContext, "autoDiffCreateLinearMapContext", "", Special)
988+
// autoDiffCreateLinearMapContextWithType: (T.Type) -> Builtin.NativeObject
989+
BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffCreateLinearMapContextWithType, "autoDiffCreateLinearMapContextWithType", "", Special)
990990

991991
// autoDiffProjectTopLevelSubcontext: (Builtin.NativeObject) -> Builtin.RawPointer
992992
BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffProjectTopLevelSubcontext, "autoDiffProjectTopLevelSubcontext", "n", Special)
993993

994-
// autoDiffAllocateSubcontext: (Builtin.NativeObject, Builtin.Word) -> Builtin.RawPointer
995-
BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffAllocateSubcontext, "autoDiffAllocateSubcontext", "", Special)
994+
// autoDiffAllocateSubcontextWithType: (Builtin.NativeObject, T.Type) -> Builtin.RawPointer
995+
BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffAllocateSubcontextWithType, "autoDiffAllocateSubcontextWithType", "", Special)
996996

997997
/// Build a Builtin.Executor value from an "ordinary" serial executor
998998
/// reference.

‎include/swift/Runtime/RuntimeFunctions.def‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2273,12 +2273,12 @@ FUNCTION(TaskGroupDestroy,
22732273
ATTRS(NoUnwind),
22742274
EFFECT(Concurrency))
22752275

2276-
// AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext(size_t);
2277-
FUNCTION(AutoDiffCreateLinearMapContext,
2278-
swift_autoDiffCreateLinearMapContext, SwiftCC,
2276+
// AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContextWithType(const Metadata *);
2277+
FUNCTION(AutoDiffCreateLinearMapContextWithType,
2278+
swift_autoDiffCreateLinearMapContextWithType, SwiftCC,
22792279
DifferentiationAvailability,
22802280
RETURNS(RefCountedPtrTy),
2281-
ARGS(SizeTy),
2281+
ARGS(TypeMetadataPtrTy),
22822282
ATTRS(NoUnwind, ArgMemOnly),
22832283
EFFECT(AutoDiff))
22842284

@@ -2291,12 +2291,12 @@ FUNCTION(AutoDiffProjectTopLevelSubcontext,
22912291
ATTRS(NoUnwind, ArgMemOnly),
22922292
EFFECT(AutoDiff))
22932293

2294-
// void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, size_t);
2295-
FUNCTION(AutoDiffAllocateSubcontext,
2296-
swift_autoDiffAllocateSubcontext, SwiftCC,
2294+
// void *swift_autoDiffAllocateSubcontextWithType(AutoDiffLinearMapContext *, const Metadata *);
2295+
FUNCTION(AutoDiffAllocateSubcontextWithType,
2296+
swift_autoDiffAllocateSubcontextWithType, SwiftCC,
22972297
DifferentiationAvailability,
22982298
RETURNS(Int8PtrTy),
2299-
ARGS(RefCountedPtrTy, SizeTy),
2299+
ARGS(RefCountedPtrTy, TypeMetadataPtrTy),
23002300
ATTRS(NoUnwind, ArgMemOnly),
23012301
EFFECT(AutoDiff))
23022302

‎lib/AST/Builtins.cpp‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1609,7 +1609,8 @@ static ValueDecl *getBuildComplexEqualitySerialExecutorRef(ASTContext &ctx,
16091609
static ValueDecl *getAutoDiffCreateLinearMapContext(ASTContext &ctx,
16101610
Identifier id) {
16111611
return getBuiltinFunction(
1612-
id, {BuiltinIntegerType::getWordType(ctx)}, ctx.TheNativeObjectType);
1612+
ctx, id, _thin, _generics(_unrestricted),
1613+
_parameters(_metatype(_typeparam(0))), _nativeObject);
16131614
}
16141615

16151616
static ValueDecl *getAutoDiffProjectTopLevelSubcontext(ASTContext &ctx,
@@ -1621,8 +1622,8 @@ static ValueDecl *getAutoDiffProjectTopLevelSubcontext(ASTContext &ctx,
16211622
static ValueDecl *getAutoDiffAllocateSubcontext(ASTContext &ctx,
16221623
Identifier id) {
16231624
return getBuiltinFunction(
1624-
id, {ctx.TheNativeObjectType, BuiltinIntegerType::getWordType(ctx)},
1625-
ctx.TheRawPointerType);
1625+
ctx, id, _thin, _generics(_unrestricted),
1626+
_parameters(_nativeObject, _metatype(_typeparam(0))), _rawPointer);
16261627
}
16271628

16281629
static ValueDecl *getPoundAssert(ASTContext &Context, Identifier Id) {
@@ -2966,13 +2967,13 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
29662967
case BuiltinValueKind::HopToActor:
29672968
return getHopToActor(Context, Id);
29682969

2969-
case BuiltinValueKind::AutoDiffCreateLinearMapContext:
2970+
case BuiltinValueKind::AutoDiffCreateLinearMapContextWithType:
29702971
return getAutoDiffCreateLinearMapContext(Context, Id);
29712972

29722973
case BuiltinValueKind::AutoDiffProjectTopLevelSubcontext:
29732974
return getAutoDiffProjectTopLevelSubcontext(Context, Id);
29742975

2975-
case BuiltinValueKind::AutoDiffAllocateSubcontext:
2976+
case BuiltinValueKind::AutoDiffAllocateSubcontextWithType:
29762977
return getAutoDiffAllocateSubcontext(Context, Id);
29772978
}
29782979

‎lib/IRGen/GenBuiltin.cpp‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1307,9 +1307,10 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
13071307
return;
13081308
}
13091309

1310-
if (Builtin.ID == BuiltinValueKind::AutoDiffCreateLinearMapContext) {
1311-
auto topLevelSubcontextSize = args.claimNext();
1312-
out.add(emitAutoDiffCreateLinearMapContext(IGF, topLevelSubcontextSize)
1310+
if (Builtin.ID == BuiltinValueKind::AutoDiffCreateLinearMapContextWithType) {
1311+
auto topLevelSubcontextMetaType = args.claimNext();
1312+
out.add(emitAutoDiffCreateLinearMapContextWithType(
1313+
IGF, topLevelSubcontextMetaType)
13131314
.getAddress());
13141315
return;
13151316
}
@@ -1322,12 +1323,13 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
13221323
return;
13231324
}
13241325

1325-
if (Builtin.ID == BuiltinValueKind::AutoDiffAllocateSubcontext) {
1326+
if (Builtin.ID == BuiltinValueKind::AutoDiffAllocateSubcontextWithType) {
13261327
Address allocatorAddr(args.claimNext(), IGF.IGM.RefCountedStructTy,
13271328
IGF.IGM.getPointerAlignment());
1328-
auto size = args.claimNext();
1329-
out.add(
1330-
emitAutoDiffAllocateSubcontext(IGF, allocatorAddr, size).getAddress());
1329+
auto subcontextMetatype = args.claimNext();
1330+
out.add(emitAutoDiffAllocateSubcontextWithType(IGF, allocatorAddr,
1331+
subcontextMetatype)
1332+
.getAddress());
13311333
return;
13321334
}
13331335

‎lib/IRGen/GenCall.cpp‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5479,11 +5479,13 @@ IRGenFunction::getFunctionPointerForResumeIntrinsic(llvm::Value *resume) {
54795479
return fnPtr;
54805480
}
54815481

5482-
Address irgen::emitAutoDiffCreateLinearMapContext(
5483-
IRGenFunction &IGF, llvm::Value *topLevelSubcontextSize) {
5482+
Address irgen::emitAutoDiffCreateLinearMapContextWithType(
5483+
IRGenFunction &IGF, llvm::Value *topLevelSubcontextMetatype) {
5484+
topLevelSubcontextMetatype = IGF.Builder.CreateBitCast(
5485+
topLevelSubcontextMetatype, IGF.IGM.TypeMetadataPtrTy);
54845486
auto *call = IGF.Builder.CreateCall(
5485-
IGF.IGM.getAutoDiffCreateLinearMapContextFunctionPointer(),
5486-
{topLevelSubcontextSize});
5487+
IGF.IGM.getAutoDiffCreateLinearMapContextWithTypeFunctionPointer(),
5488+
{topLevelSubcontextMetatype});
54875489
call->setDoesNotThrow();
54885490
call->setCallingConv(IGF.IGM.SwiftCC);
54895491
return Address(call, IGF.IGM.RefCountedStructTy,
@@ -5500,11 +5502,13 @@ Address irgen::emitAutoDiffProjectTopLevelSubcontext(
55005502
return Address(call, IGF.IGM.Int8Ty, IGF.IGM.getPointerAlignment());
55015503
}
55025504

5503-
Address irgen::emitAutoDiffAllocateSubcontext(
5504-
IRGenFunction &IGF, Address context, llvm::Value *size) {
5505+
Address irgen::emitAutoDiffAllocateSubcontextWithType(
5506+
IRGenFunction &IGF, Address context, llvm::Value *subcontextMetatype) {
5507+
subcontextMetatype =
5508+
IGF.Builder.CreateBitCast(subcontextMetatype, IGF.IGM.TypeMetadataPtrTy);
55055509
auto *call = IGF.Builder.CreateCall(
5506-
IGF.IGM.getAutoDiffAllocateSubcontextFunctionPointer(),
5507-
{context.getAddress(), size});
5510+
IGF.IGM.getAutoDiffAllocateSubcontextWithTypeFunctionPointer(),
5511+
{context.getAddress(), subcontextMetatype});
55085512
call->setDoesNotThrow();
55095513
call->setCallingConv(IGF.IGM.SwiftCC);
55105514
return Address(call, IGF.IGM.Int8Ty, IGF.IGM.getPointerAlignment());

‎lib/IRGen/GenCall.h‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -261,12 +261,15 @@ namespace irgen {
261261
CanSILFunctionType fnType, Explosion &result,
262262
Explosion &error);
263263

264-
Address emitAutoDiffCreateLinearMapContext(
265-
IRGenFunction &IGF, llvm::Value *topLevelSubcontextSize);
264+
Address emitAutoDiffCreateLinearMapContextWithType(
265+
IRGenFunction &IGF, llvm::Value *topLevelSubcontextMetatype);
266+
266267
Address emitAutoDiffProjectTopLevelSubcontext(
267268
IRGenFunction &IGF, Address context);
268-
Address emitAutoDiffAllocateSubcontext(
269-
IRGenFunction &IGF, Address context, llvm::Value *size);
269+
270+
Address
271+
emitAutoDiffAllocateSubcontextWithType(IRGenFunction &IGF, Address context,
272+
llvm::Value *subcontextMetatype);
270273

271274
FunctionPointer getFunctionPointerForDispatchCall(IRGenModule &IGM,
272275
const FunctionPointer &fn);

‎lib/SIL/IR/OperandOwnership.cpp‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -944,7 +944,7 @@ BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse, InitializeDistributedRemoteActor)
944944
BUILTIN_OPERAND_OWNERSHIP(InstantaneousUse,
945945
InitializeNonDefaultDistributedActor)
946946

947-
BUILTIN_OPERAND_OWNERSHIP(PointerEscape, AutoDiffAllocateSubcontext)
947+
BUILTIN_OPERAND_OWNERSHIP(PointerEscape, AutoDiffAllocateSubcontextWithType)
948948
BUILTIN_OPERAND_OWNERSHIP(PointerEscape, AutoDiffProjectTopLevelSubcontext)
949949

950950
// FIXME: ConvertTaskToJob is documented as taking NativePointer. It's operand's
@@ -956,8 +956,7 @@ BUILTIN_OPERAND_OWNERSHIP(BitwiseEscape, BuildComplexEqualitySerialExecutorRef)
956956
BUILTIN_OPERAND_OWNERSHIP(BitwiseEscape, BuildDefaultActorExecutorRef)
957957
BUILTIN_OPERAND_OWNERSHIP(BitwiseEscape, BuildMainActorExecutorRef)
958958

959-
BUILTIN_OPERAND_OWNERSHIP(TrivialUse, AutoDiffCreateLinearMapContext)
960-
959+
BUILTIN_OPERAND_OWNERSHIP(TrivialUse, AutoDiffCreateLinearMapContextWithType)
961960
#undef BUILTIN_OPERAND_OWNERSHIP
962961

963962
#define SHOULD_NEVER_VISIT_BUILTIN(ID) \

‎lib/SIL/IR/ValueOwnership.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -572,9 +572,9 @@ CONSTANT_OWNERSHIP_BUILTIN(None, InitializeDefaultActor)
572572
CONSTANT_OWNERSHIP_BUILTIN(None, DestroyDefaultActor)
573573
CONSTANT_OWNERSHIP_BUILTIN(None, InitializeDistributedRemoteActor)
574574
CONSTANT_OWNERSHIP_BUILTIN(None, InitializeNonDefaultDistributedActor)
575-
CONSTANT_OWNERSHIP_BUILTIN(Owned, AutoDiffCreateLinearMapContext)
575+
CONSTANT_OWNERSHIP_BUILTIN(Owned, AutoDiffCreateLinearMapContextWithType)
576576
CONSTANT_OWNERSHIP_BUILTIN(None, AutoDiffProjectTopLevelSubcontext)
577-
CONSTANT_OWNERSHIP_BUILTIN(None, AutoDiffAllocateSubcontext)
577+
CONSTANT_OWNERSHIP_BUILTIN(None, AutoDiffAllocateSubcontextWithType)
578578
CONSTANT_OWNERSHIP_BUILTIN(None, GetCurrentExecutor)
579579
CONSTANT_OWNERSHIP_BUILTIN(None, ResumeNonThrowingContinuationReturning)
580580
CONSTANT_OWNERSHIP_BUILTIN(None, ResumeThrowingContinuationReturning)

‎lib/SIL/Utils/MemAccessUtils.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2567,8 +2567,8 @@ static void visitBuiltinAddress(BuiltinInst *builtin,
25672567
case BuiltinValueKind::CancelAsyncTask:
25682568
case BuiltinValueKind::CreateAsyncTask:
25692569
case BuiltinValueKind::CreateAsyncTaskInGroup:
2570-
case BuiltinValueKind::AutoDiffCreateLinearMapContext:
2571-
case BuiltinValueKind::AutoDiffAllocateSubcontext:
2570+
case BuiltinValueKind::AutoDiffCreateLinearMapContextWithType:
2571+
case BuiltinValueKind::AutoDiffAllocateSubcontextWithType:
25722572
case BuiltinValueKind::InitializeDefaultActor:
25732573
case BuiltinValueKind::InitializeDistributedRemoteActor:
25742574
case BuiltinValueKind::InitializeNonDefaultDistributedActor:

‎lib/SILGen/SILGenBuiltin.cpp‎

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1708,16 +1708,15 @@ static ManagedValue emitBuiltinHopToActor(SILGenFunction &SGF, SILLocation loc,
17081708
return ManagedValue::forObjectRValueWithoutOwnership(SGF.emitEmptyTuple(loc));
17091709
}
17101710

1711-
static ManagedValue emitBuiltinAutoDiffCreateLinearMapContext(
1711+
static ManagedValue emitBuiltinAutoDiffCreateLinearMapContextWithType(
17121712
SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
17131713
ArrayRef<ManagedValue> args, SGFContext C) {
17141714
ASTContext &ctx = SGF.getASTContext();
17151715
auto *builtinApply = SGF.B.createBuiltin(
17161716
loc,
1717-
ctx.getIdentifier(
1718-
getBuiltinName(BuiltinValueKind::AutoDiffCreateLinearMapContext)),
1719-
SILType::getNativeObjectType(ctx),
1720-
subs,
1717+
ctx.getIdentifier(getBuiltinName(
1718+
BuiltinValueKind::AutoDiffCreateLinearMapContextWithType)),
1719+
SILType::getNativeObjectType(ctx), subs,
17211720
/*args*/ {args[0].getValue()});
17221721
return SGF.emitManagedRValueWithCleanup(builtinApply);
17231722
}
@@ -1736,16 +1735,15 @@ static ManagedValue emitBuiltinAutoDiffProjectTopLevelSubcontext(
17361735
return ManagedValue::forObjectRValueWithoutOwnership(builtinApply);
17371736
}
17381737

1739-
static ManagedValue emitBuiltinAutoDiffAllocateSubcontext(
1738+
static ManagedValue emitBuiltinAutoDiffAllocateSubcontextWithType(
17401739
SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
17411740
ArrayRef<ManagedValue> args, SGFContext C) {
17421741
ASTContext &ctx = SGF.getASTContext();
17431742
auto *builtinApply = SGF.B.createBuiltin(
17441743
loc,
17451744
ctx.getIdentifier(
1746-
getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontext)),
1747-
SILType::getRawPointerType(ctx),
1748-
subs,
1745+
getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontextWithType)),
1746+
SILType::getRawPointerType(ctx), subs,
17491747
/*args*/ {args[0].borrow(SGF, loc).getValue(), args[1].getValue()});
17501748
return ManagedValue::forObjectRValueWithoutOwnership(builtinApply);
17511749
}

0 commit comments

Comments
 (0)