Skip to content

Commit dd0dce4

Browse files
jkshtjjkshtj
committed
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.
1 parent 91bafd4 commit dd0dce4

File tree

16 files changed

+274
-48
lines changed

16 files changed

+274
-48
lines changed

‎include/swift/AST/Builtins.def

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -985,15 +985,21 @@ 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: (T.Type) -> Builtin.NativeObject
988+
// autoDiffCreateLinearMapContext: (Builtin.Word) -> Builtin.NativeObject
989989
BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffCreateLinearMapContext, "autoDiffCreateLinearMapContext", "", Special)
990990

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

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

997+
// autoDiffCreateLinearMapContext2: (T.Type) -> Builtin.NativeObject
998+
BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffCreateLinearMapContext2, "autoDiffCreateLinearMapContext2", "", Special)
999+
1000+
// autoDiffAllocateSubcontext2: (Builtin.NativeObject, T.Type) -> Builtin.RawPointer
1001+
BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffAllocateSubcontext2, "autoDiffAllocateSubcontext2", "", Special)
1002+
9971003
/// Build a Builtin.Executor value from an "ordinary" serial executor
9981004
/// reference.
9991005
BUILTIN_MISC_OPERATION_WITH_SILGEN(BuildOrdinarySerialExecutorRef,

‎include/swift/Runtime/RuntimeFunctions.def

Lines changed: 21 additions & 3 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(const Metadata *);
2276+
// AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext(size_t);
22772277
FUNCTION(AutoDiffCreateLinearMapContext,
22782278
swift_autoDiffCreateLinearMapContext, SwiftCC,
22792279
DifferentiationAvailability,
22802280
RETURNS(RefCountedPtrTy),
2281-
ARGS(TypeMetadataPtrTy),
2281+
ARGS(SizeTy),
22822282
ATTRS(NoUnwind, ArgMemOnly),
22832283
EFFECT(AutoDiff))
22842284

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

2294-
// void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, const Metadata *);
2294+
// void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, size_t);
22952295
FUNCTION(AutoDiffAllocateSubcontext,
22962296
swift_autoDiffAllocateSubcontext, SwiftCC,
22972297
DifferentiationAvailability,
22982298
RETURNS(Int8PtrTy),
2299+
ARGS(RefCountedPtrTy, SizeTy),
2300+
ATTRS(NoUnwind, ArgMemOnly),
2301+
EFFECT(AutoDiff))
2302+
2303+
// AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext2(const Metadata *);
2304+
FUNCTION(AutoDiffCreateLinearMapContext2,
2305+
swift_autoDiffCreateLinearMapContext2, SwiftCC,
2306+
DifferentiationAvailability,
2307+
RETURNS(RefCountedPtrTy),
2308+
ARGS(TypeMetadataPtrTy),
2309+
ATTRS(NoUnwind, ArgMemOnly),
2310+
EFFECT(AutoDiff))
2311+
2312+
// void *swift_autoDiffAllocateSubcontext2(AutoDiffLinearMapContext *, const Metadata *);
2313+
FUNCTION(AutoDiffAllocateSubcontext2,
2314+
swift_autoDiffAllocateSubcontext2, SwiftCC,
2315+
DifferentiationAvailability,
2316+
RETURNS(Int8PtrTy),
22992317
ARGS(RefCountedPtrTy, TypeMetadataPtrTy),
23002318
ATTRS(NoUnwind, ArgMemOnly),
23012319
EFFECT(AutoDiff))

‎lib/AST/Builtins.cpp

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1608,9 +1608,8 @@ static ValueDecl *getBuildComplexEqualitySerialExecutorRef(ASTContext &ctx,
16081608

16091609
static ValueDecl *getAutoDiffCreateLinearMapContext(ASTContext &ctx,
16101610
Identifier id) {
1611-
return getBuiltinFunction(ctx, id, _thin, _generics(_unrestricted),
1612-
_parameters(_metatype(_typeparam(0))),
1613-
_nativeObject);
1611+
return getBuiltinFunction(id, {BuiltinIntegerType::getWordType(ctx)},
1612+
ctx.TheNativeObjectType);
16141613
}
16151614

16161615
static ValueDecl *getAutoDiffProjectTopLevelSubcontext(ASTContext &ctx,
@@ -1621,6 +1620,20 @@ static ValueDecl *getAutoDiffProjectTopLevelSubcontext(ASTContext &ctx,
16211620

16221621
static ValueDecl *getAutoDiffAllocateSubcontext(ASTContext &ctx,
16231622
Identifier id) {
1623+
return getBuiltinFunction(
1624+
id, {ctx.TheNativeObjectType, BuiltinIntegerType::getWordType(ctx)},
1625+
ctx.TheRawPointerType);
1626+
}
1627+
1628+
static ValueDecl *getAutoDiffCreateLinearMapContext2(ASTContext &ctx,
1629+
Identifier id) {
1630+
return getBuiltinFunction(ctx, id, _thin, _generics(_unrestricted),
1631+
_parameters(_metatype(_typeparam(0))),
1632+
_nativeObject);
1633+
}
1634+
1635+
static ValueDecl *getAutoDiffAllocateSubcontext2(ASTContext &ctx,
1636+
Identifier id) {
16241637
return getBuiltinFunction(
16251638
ctx, id, _thin, _generics(_unrestricted),
16261639
_parameters(_nativeObject, _metatype(_typeparam(0))), _rawPointer);
@@ -2975,6 +2988,12 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
29752988

29762989
case BuiltinValueKind::AutoDiffAllocateSubcontext:
29772990
return getAutoDiffAllocateSubcontext(Context, Id);
2991+
2992+
case BuiltinValueKind::AutoDiffCreateLinearMapContext2:
2993+
return getAutoDiffCreateLinearMapContext2(Context, Id);
2994+
2995+
case BuiltinValueKind::AutoDiffAllocateSubcontext2:
2996+
return getAutoDiffAllocateSubcontext2(Context, Id);
29782997
}
29792998

29802999
llvm_unreachable("bad builtin value!");

‎lib/IRGen/GenBuiltin.cpp

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1308,8 +1308,8 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
13081308
}
13091309

13101310
if (Builtin.ID == BuiltinValueKind::AutoDiffCreateLinearMapContext) {
1311-
auto topLevelSubcontextMetaType = args.claimNext();
1312-
out.add(emitAutoDiffCreateLinearMapContext(IGF, topLevelSubcontextMetaType)
1311+
auto topLevelSubcontextSize = args.claimNext();
1312+
out.add(emitAutoDiffCreateLinearMapContext(IGF, topLevelSubcontextSize)
13131313
.getAddress());
13141314
return;
13151315
}
@@ -1323,11 +1323,27 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
13231323
}
13241324

13251325
if (Builtin.ID == BuiltinValueKind::AutoDiffAllocateSubcontext) {
1326+
Address allocatorAddr(args.claimNext(), IGF.IGM.RefCountedStructTy,
1327+
IGF.IGM.getPointerAlignment());
1328+
auto size = args.claimNext();
1329+
out.add(
1330+
emitAutoDiffAllocateSubcontext(IGF, allocatorAddr, size).getAddress());
1331+
return;
1332+
}
1333+
1334+
if (Builtin.ID == BuiltinValueKind::AutoDiffCreateLinearMapContext2) {
1335+
auto topLevelSubcontextMetaType = args.claimNext();
1336+
out.add(emitAutoDiffCreateLinearMapContext2(IGF, topLevelSubcontextMetaType)
1337+
.getAddress());
1338+
return;
1339+
}
1340+
1341+
if (Builtin.ID == BuiltinValueKind::AutoDiffAllocateSubcontext2) {
13261342
Address allocatorAddr(args.claimNext(), IGF.IGM.RefCountedStructTy,
13271343
IGF.IGM.getPointerAlignment());
13281344
auto subcontextMetatype = args.claimNext();
13291345
out.add(
1330-
emitAutoDiffAllocateSubcontext(IGF, allocatorAddr, subcontextMetatype)
1346+
emitAutoDiffAllocateSubcontext2(IGF, allocatorAddr, subcontextMetatype)
13311347
.getAddress());
13321348
return;
13331349
}

‎lib/IRGen/GenCall.cpp

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

5482-
Address irgen::emitAutoDiffCreateLinearMapContext(
5483-
IRGenFunction &IGF, llvm::Value *topLevelSubcontextMetatype) {
5484-
topLevelSubcontextMetatype = IGF.Builder.CreateBitCast(
5485-
topLevelSubcontextMetatype, IGF.IGM.TypeMetadataPtrTy);
5482+
Address
5483+
irgen::emitAutoDiffCreateLinearMapContext(IRGenFunction &IGF,
5484+
llvm::Value *topLevelSubcontextSize) {
54865485
auto *call = IGF.Builder.CreateCall(
54875486
IGF.IGM.getAutoDiffCreateLinearMapContextFunctionPointer(),
5488-
{topLevelSubcontextMetatype});
5487+
{topLevelSubcontextSize});
54895488
call->setDoesNotThrow();
54905489
call->setCallingConv(IGF.IGM.SwiftCC);
54915490
return Address(call, IGF.IGM.RefCountedStructTy,
@@ -5504,11 +5503,35 @@ Address irgen::emitAutoDiffProjectTopLevelSubcontext(
55045503

55055504
Address irgen::emitAutoDiffAllocateSubcontext(IRGenFunction &IGF,
55065505
Address context,
5507-
llvm::Value *subcontextMetatype) {
5506+
llvm::Value *size) {
5507+
auto *call = IGF.Builder.CreateCall(
5508+
IGF.IGM.getAutoDiffAllocateSubcontextFunctionPointer(),
5509+
{context.getAddress(), size});
5510+
call->setDoesNotThrow();
5511+
call->setCallingConv(IGF.IGM.SwiftCC);
5512+
return Address(call, IGF.IGM.Int8Ty, IGF.IGM.getPointerAlignment());
5513+
}
5514+
5515+
Address irgen::emitAutoDiffCreateLinearMapContext2(
5516+
IRGenFunction &IGF, llvm::Value *topLevelSubcontextMetatype) {
5517+
topLevelSubcontextMetatype = IGF.Builder.CreateBitCast(
5518+
topLevelSubcontextMetatype, IGF.IGM.TypeMetadataPtrTy);
5519+
auto *call = IGF.Builder.CreateCall(
5520+
IGF.IGM.getAutoDiffCreateLinearMapContext2FunctionPointer(),
5521+
{topLevelSubcontextMetatype});
5522+
call->setDoesNotThrow();
5523+
call->setCallingConv(IGF.IGM.SwiftCC);
5524+
return Address(call, IGF.IGM.RefCountedStructTy,
5525+
IGF.IGM.getPointerAlignment());
5526+
}
5527+
5528+
Address
5529+
irgen::emitAutoDiffAllocateSubcontext2(IRGenFunction &IGF, Address context,
5530+
llvm::Value *subcontextMetatype) {
55085531
subcontextMetatype =
55095532
IGF.Builder.CreateBitCast(subcontextMetatype, IGF.IGM.TypeMetadataPtrTy);
55105533
auto *call = IGF.Builder.CreateCall(
5511-
IGF.IGM.getAutoDiffAllocateSubcontextFunctionPointer(),
5534+
IGF.IGM.getAutoDiffAllocateSubcontext2FunctionPointer(),
55125535
{context.getAddress(), subcontextMetatype});
55135536
call->setDoesNotThrow();
55145537
call->setCallingConv(IGF.IGM.SwiftCC);

‎lib/IRGen/GenCall.h

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -263,11 +263,17 @@ namespace irgen {
263263

264264
Address
265265
emitAutoDiffCreateLinearMapContext(IRGenFunction &IGF,
266-
llvm::Value *topLevelSubcontextMetatype);
266+
llvm::Value *topLevelSubcontextSize);
267267
Address emitAutoDiffProjectTopLevelSubcontext(
268268
IRGenFunction &IGF, Address context);
269269
Address emitAutoDiffAllocateSubcontext(IRGenFunction &IGF, Address context,
270-
llvm::Value *subcontextMetatype);
270+
llvm::Value *size);
271+
272+
Address
273+
emitAutoDiffCreateLinearMapContext2(IRGenFunction &IGF,
274+
llvm::Value *topLevelSubcontextMetatype);
275+
Address emitAutoDiffAllocateSubcontext2(IRGenFunction &IGF, Address context,
276+
llvm::Value *subcontextMetatype);
271277

272278
FunctionPointer getFunctionPointerForDispatchCall(IRGenModule &IGM,
273279
const FunctionPointer &fn);

‎lib/SIL/IR/OperandOwnership.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -957,6 +957,8 @@ BUILTIN_OPERAND_OWNERSHIP(BitwiseEscape, BuildMainActorExecutorRef)
957957

958958
BUILTIN_OPERAND_OWNERSHIP(TrivialUse, AutoDiffCreateLinearMapContext)
959959

960+
BUILTIN_OPERAND_OWNERSHIP(TrivialUse, AutoDiffCreateLinearMapContext2)
961+
BUILTIN_OPERAND_OWNERSHIP(PointerEscape, AutoDiffAllocateSubcontext2)
960962
#undef BUILTIN_OPERAND_OWNERSHIP
961963

962964
#define SHOULD_NEVER_VISIT_BUILTIN(ID) \

‎lib/SIL/IR/ValueOwnership.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -574,6 +574,8 @@ CONSTANT_OWNERSHIP_BUILTIN(None, InitializeNonDefaultDistributedActor)
574574
CONSTANT_OWNERSHIP_BUILTIN(Owned, AutoDiffCreateLinearMapContext)
575575
CONSTANT_OWNERSHIP_BUILTIN(None, AutoDiffProjectTopLevelSubcontext)
576576
CONSTANT_OWNERSHIP_BUILTIN(None, AutoDiffAllocateSubcontext)
577+
CONSTANT_OWNERSHIP_BUILTIN(Owned, AutoDiffCreateLinearMapContext2)
578+
CONSTANT_OWNERSHIP_BUILTIN(None, AutoDiffAllocateSubcontext2)
577579
CONSTANT_OWNERSHIP_BUILTIN(None, GetCurrentExecutor)
578580
CONSTANT_OWNERSHIP_BUILTIN(None, ResumeNonThrowingContinuationReturning)
579581
CONSTANT_OWNERSHIP_BUILTIN(None, ResumeThrowingContinuationReturning)

‎lib/SIL/Utils/MemAccessUtils.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2569,6 +2569,8 @@ static void visitBuiltinAddress(BuiltinInst *builtin,
25692569
case BuiltinValueKind::CreateAsyncTaskInGroup:
25702570
case BuiltinValueKind::AutoDiffCreateLinearMapContext:
25712571
case BuiltinValueKind::AutoDiffAllocateSubcontext:
2572+
case BuiltinValueKind::AutoDiffCreateLinearMapContext2:
2573+
case BuiltinValueKind::AutoDiffAllocateSubcontext2:
25722574
case BuiltinValueKind::InitializeDefaultActor:
25732575
case BuiltinValueKind::InitializeDistributedRemoteActor:
25742576
case BuiltinValueKind::InitializeNonDefaultDistributedActor:

‎lib/SILGen/SILGenBuiltin.cpp

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1750,6 +1750,32 @@ static ManagedValue emitBuiltinAutoDiffAllocateSubcontext(
17501750
return ManagedValue::forObjectRValueWithoutOwnership(builtinApply);
17511751
}
17521752

1753+
static ManagedValue emitBuiltinAutoDiffCreateLinearMapContext2(
1754+
SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
1755+
ArrayRef<ManagedValue> args, SGFContext C) {
1756+
ASTContext &ctx = SGF.getASTContext();
1757+
auto *builtinApply = SGF.B.createBuiltin(
1758+
loc,
1759+
ctx.getIdentifier(
1760+
getBuiltinName(BuiltinValueKind::AutoDiffCreateLinearMapContext2)),
1761+
SILType::getNativeObjectType(ctx), subs,
1762+
/*args*/ {args[0].getValue()});
1763+
return SGF.emitManagedRValueWithCleanup(builtinApply);
1764+
}
1765+
1766+
static ManagedValue emitBuiltinAutoDiffAllocateSubcontext2(
1767+
SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
1768+
ArrayRef<ManagedValue> args, SGFContext C) {
1769+
ASTContext &ctx = SGF.getASTContext();
1770+
auto *builtinApply = SGF.B.createBuiltin(
1771+
loc,
1772+
ctx.getIdentifier(
1773+
getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontext2)),
1774+
SILType::getRawPointerType(ctx), subs,
1775+
/*args*/ {args[0].borrow(SGF, loc).getValue(), args[1].getValue()});
1776+
return ManagedValue::forObjectRValueWithoutOwnership(builtinApply);
1777+
}
1778+
17531779
// The only reason these need special attention is that we want these to
17541780
// be borrowed arguments, but the default emitter doesn't handle borrowed
17551781
// arguments correctly.

‎lib/SILOptimizer/Differentiation/VJPCloner.cpp

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ class VJPCloner::Implementation final
132132
pullbackContextValue = Builder.createBuiltin(
133133
original->getLocation(),
134134
getASTContext().getIdentifier(
135-
getBuiltinName(BuiltinValueKind::AutoDiffCreateLinearMapContext)),
135+
getBuiltinName(BuiltinValueKind::AutoDiffCreateLinearMapContext2)),
136136
SILType::getNativeObjectType(getASTContext()), SubstitutionMap(),
137137
{pbTupleMetatype});
138138
borrowedPullbackContextValue = Builder.createBeginBorrow(
@@ -156,8 +156,8 @@ class VJPCloner::Implementation final
156156
return builtinAutoDiffAllocateSubcontextGenericSignature;
157157
auto &ctx = getASTContext();
158158
auto *decl = cast<FuncDecl>(getBuiltinValueDecl(
159-
ctx, ctx.getIdentifier(
160-
getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontext))));
159+
ctx, ctx.getIdentifier(getBuiltinName(
160+
BuiltinValueKind::AutoDiffAllocateSubcontext2))));
161161
builtinAutoDiffAllocateSubcontextGenericSignature =
162162
decl->getGenericSignature();
163163
assert(builtinAutoDiffAllocateSubcontextGenericSignature);
@@ -1086,7 +1086,7 @@ EnumInst *VJPCloner::Implementation::buildPredecessorEnumValue(
10861086
auto rawBufferValue = builder.createBuiltin(
10871087
loc,
10881088
getASTContext().getIdentifier(
1089-
getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontext)),
1089+
getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontext2)),
10901090
rawPtrType, SubstitutionMap(),
10911091
{borrowedPullbackContextValue, pbTupleMetatype});
10921092

‎lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,7 @@ static bool isBarrier(SILInstruction *inst) {
147147
case BuiltinValueKind::GetCurrentAsyncTask:
148148
case BuiltinValueKind::GetCurrentExecutor:
149149
case BuiltinValueKind::AutoDiffCreateLinearMapContext:
150+
case BuiltinValueKind::AutoDiffCreateLinearMapContext2:
150151
case BuiltinValueKind::EndAsyncLet:
151152
case BuiltinValueKind::EndAsyncLetLifetime:
152153
case BuiltinValueKind::CreateTaskGroup:
@@ -200,6 +201,7 @@ static bool isBarrier(SILInstruction *inst) {
200201
case BuiltinValueKind::ResumeThrowingContinuationThrowing:
201202
case BuiltinValueKind::AutoDiffProjectTopLevelSubcontext:
202203
case BuiltinValueKind::AutoDiffAllocateSubcontext:
204+
case BuiltinValueKind::AutoDiffAllocateSubcontext2:
203205
case BuiltinValueKind::AddressOfBorrowOpaque:
204206
case BuiltinValueKind::UnprotectedAddressOfBorrowOpaque:
205207
return true;

0 commit comments

Comments
 (0)