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.

Author: jkshtj

include/swift/AST/Builtins.def

@@ -985,15 +985,21 @@ 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
+// autoDiffCreateLinearMapContext: (Builtin.Word) -> 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, T.Type) -> Builtin.RawPointer
+// autoDiffAllocateSubcontext: (Builtin.NativeObject, Builtin.Word) -> Builtin.RawPointer
 BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffAllocateSubcontext, "autoDiffAllocateSubcontext", "", Special)
 
+// autoDiffCreateLinearMapContext2: (T.Type) -> Builtin.NativeObject
+BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffCreateLinearMapContext2, "autoDiffCreateLinearMapContext2", "", Special)
+
+// autoDiffAllocateSubcontext2: (Builtin.NativeObject, T.Type) -> Builtin.RawPointer
+BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffAllocateSubcontext2, "autoDiffAllocateSubcontext2", "", Special)
+
 /// Build a Builtin.Executor value from an "ordinary" serial executor
 /// reference.
 BUILTIN_MISC_OPERATION_WITH_SILGEN(BuildOrdinarySerialExecutorRef,

include/swift/Runtime/RuntimeFunctions.def

@@ -2273,12 +2273,12 @@ FUNCTION(TaskGroupDestroy,
          ATTRS(NoUnwind),
          EFFECT(Concurrency))
 
-// AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext(const Metadata *);
+// AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext(size_t);
 FUNCTION(AutoDiffCreateLinearMapContext,
          swift_autoDiffCreateLinearMapContext, SwiftCC,
          DifferentiationAvailability,
          RETURNS(RefCountedPtrTy),
-         ARGS(TypeMetadataPtrTy),
+         ARGS(SizeTy),
          ATTRS(NoUnwind, ArgMemOnly),
          EFFECT(AutoDiff))
 
@@ -2291,11 +2291,29 @@ FUNCTION(AutoDiffProjectTopLevelSubcontext,
          ATTRS(NoUnwind, ArgMemOnly),
          EFFECT(AutoDiff))
 
-// void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, const Metadata *);
+// void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, size_t);
 FUNCTION(AutoDiffAllocateSubcontext,
          swift_autoDiffAllocateSubcontext, SwiftCC,
          DifferentiationAvailability,
          RETURNS(Int8PtrTy),
+         ARGS(RefCountedPtrTy, SizeTy),
+         ATTRS(NoUnwind, ArgMemOnly),
+         EFFECT(AutoDiff))
+
+// AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext2(const Metadata *);
+FUNCTION(AutoDiffCreateLinearMapContext2,
+         swift_autoDiffCreateLinearMapContext2, SwiftCC,
+         DifferentiationAvailability,
+         RETURNS(RefCountedPtrTy),
+         ARGS(TypeMetadataPtrTy),
+         ATTRS(NoUnwind, ArgMemOnly),
+         EFFECT(AutoDiff))
+
+// void *swift_autoDiffAllocateSubcontext2(AutoDiffLinearMapContext *, const Metadata *);
+FUNCTION(AutoDiffAllocateSubcontext2,
+         swift_autoDiffAllocateSubcontext2, SwiftCC,
+         DifferentiationAvailability,
+         RETURNS(Int8PtrTy),
          ARGS(RefCountedPtrTy, TypeMetadataPtrTy),
          ATTRS(NoUnwind, ArgMemOnly),
          EFFECT(AutoDiff))

lib/AST/Builtins.cpp

@@ -1608,9 +1608,8 @@ 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(id, {BuiltinIntegerType::getWordType(ctx)},
+                            ctx.TheNativeObjectType);
 }
 
 static ValueDecl *getAutoDiffProjectTopLevelSubcontext(ASTContext &ctx,
@@ -1621,6 +1620,20 @@ static ValueDecl *getAutoDiffProjectTopLevelSubcontext(ASTContext &ctx,
 
 static ValueDecl *getAutoDiffAllocateSubcontext(ASTContext &ctx,
                                                 Identifier id) {
+  return getBuiltinFunction(
+      id, {ctx.TheNativeObjectType, BuiltinIntegerType::getWordType(ctx)},
+      ctx.TheRawPointerType);
+}
+
+static ValueDecl *getAutoDiffCreateLinearMapContext2(ASTContext &ctx,
+                                                     Identifier id) {
+  return getBuiltinFunction(ctx, id, _thin, _generics(_unrestricted),
+                            _parameters(_metatype(_typeparam(0))),
+                            _nativeObject);
+}
+
+static ValueDecl *getAutoDiffAllocateSubcontext2(ASTContext &ctx,
+                                                 Identifier id) {
   return getBuiltinFunction(
       ctx, id, _thin, _generics(_unrestricted),
       _parameters(_nativeObject, _metatype(_typeparam(0))), _rawPointer);
@@ -2975,6 +2988,12 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
 
   case BuiltinValueKind::AutoDiffAllocateSubcontext:
     return getAutoDiffAllocateSubcontext(Context, Id);
+
+  case BuiltinValueKind::AutoDiffCreateLinearMapContext2:
+    return getAutoDiffCreateLinearMapContext2(Context, Id);
+
+  case BuiltinValueKind::AutoDiffAllocateSubcontext2:
+    return getAutoDiffAllocateSubcontext2(Context, Id);
   }
 
   llvm_unreachable("bad builtin value!");

lib/IRGen/GenBuiltin.cpp

@@ -1308,8 +1308,8 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
   }
 
   if (Builtin.ID == BuiltinValueKind::AutoDiffCreateLinearMapContext) {
-    auto topLevelSubcontextMetaType = args.claimNext();
-    out.add(emitAutoDiffCreateLinearMapContext(IGF, topLevelSubcontextMetaType)
+    auto topLevelSubcontextSize = args.claimNext();
+    out.add(emitAutoDiffCreateLinearMapContext(IGF, topLevelSubcontextSize)
                 .getAddress());
     return;
   }
@@ -1323,11 +1323,27 @@ 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();
+    out.add(
+        emitAutoDiffAllocateSubcontext(IGF, allocatorAddr, size).getAddress());
+    return;
+  }
+
+  if (Builtin.ID == BuiltinValueKind::AutoDiffCreateLinearMapContext2) {
+    auto topLevelSubcontextMetaType = args.claimNext();
+    out.add(emitAutoDiffCreateLinearMapContext2(IGF, topLevelSubcontextMetaType)
+                .getAddress());
+    return;
+  }
+
+  if (Builtin.ID == BuiltinValueKind::AutoDiffAllocateSubcontext2) {
     Address allocatorAddr(args.claimNext(), IGF.IGM.RefCountedStructTy,
                           IGF.IGM.getPointerAlignment());
     auto subcontextMetatype = args.claimNext();
     out.add(
-        emitAutoDiffAllocateSubcontext(IGF, allocatorAddr, subcontextMetatype)
+        emitAutoDiffAllocateSubcontext2(IGF, allocatorAddr, subcontextMetatype)
             .getAddress());
     return;
   }

lib/IRGen/GenCall.cpp

@@ -5479,13 +5479,12 @@ IRGenFunction::getFunctionPointerForResumeIntrinsic(llvm::Value *resume) {
   return fnPtr;
 }
 
-Address irgen::emitAutoDiffCreateLinearMapContext(
-    IRGenFunction &IGF, llvm::Value *topLevelSubcontextMetatype) {
-  topLevelSubcontextMetatype = IGF.Builder.CreateBitCast(
-      topLevelSubcontextMetatype, IGF.IGM.TypeMetadataPtrTy);
+Address
+irgen::emitAutoDiffCreateLinearMapContext(IRGenFunction &IGF,
+                                          llvm::Value *topLevelSubcontextSize) {
   auto *call = IGF.Builder.CreateCall(
       IGF.IGM.getAutoDiffCreateLinearMapContextFunctionPointer(),
-      {topLevelSubcontextMetatype});
+      {topLevelSubcontextSize});
   call->setDoesNotThrow();
   call->setCallingConv(IGF.IGM.SwiftCC);
   return Address(call, IGF.IGM.RefCountedStructTy,
@@ -5504,11 +5503,35 @@ Address irgen::emitAutoDiffProjectTopLevelSubcontext(
 
 Address irgen::emitAutoDiffAllocateSubcontext(IRGenFunction &IGF,
                                               Address context,
-                                              llvm::Value *subcontextMetatype) {
+                                              llvm::Value *size) {
+  auto *call = IGF.Builder.CreateCall(
+      IGF.IGM.getAutoDiffAllocateSubcontextFunctionPointer(),
+      {context.getAddress(), size});
+  call->setDoesNotThrow();
+  call->setCallingConv(IGF.IGM.SwiftCC);
+  return Address(call, IGF.IGM.Int8Ty, IGF.IGM.getPointerAlignment());
+}
+
+Address irgen::emitAutoDiffCreateLinearMapContext2(
+    IRGenFunction &IGF, llvm::Value *topLevelSubcontextMetatype) {
+  topLevelSubcontextMetatype = IGF.Builder.CreateBitCast(
+      topLevelSubcontextMetatype, IGF.IGM.TypeMetadataPtrTy);
+  auto *call = IGF.Builder.CreateCall(
+      IGF.IGM.getAutoDiffCreateLinearMapContext2FunctionPointer(),
+      {topLevelSubcontextMetatype});
+  call->setDoesNotThrow();
+  call->setCallingConv(IGF.IGM.SwiftCC);
+  return Address(call, IGF.IGM.RefCountedStructTy,
+                 IGF.IGM.getPointerAlignment());
+}
+
+Address
+irgen::emitAutoDiffAllocateSubcontext2(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.getAutoDiffAllocateSubcontext2FunctionPointer(),
       {context.getAddress(), subcontextMetatype});
   call->setDoesNotThrow();
   call->setCallingConv(IGF.IGM.SwiftCC);

lib/IRGen/GenCall.h

@@ -263,11 +263,17 @@ namespace irgen {
 
   Address
   emitAutoDiffCreateLinearMapContext(IRGenFunction &IGF,
-                                     llvm::Value *topLevelSubcontextMetatype);
+                                     llvm::Value *topLevelSubcontextSize);
   Address emitAutoDiffProjectTopLevelSubcontext(
       IRGenFunction &IGF, Address context);
   Address emitAutoDiffAllocateSubcontext(IRGenFunction &IGF, Address context,
-                                         llvm::Value *subcontextMetatype);
+                                         llvm::Value *size);
+
+  Address
+  emitAutoDiffCreateLinearMapContext2(IRGenFunction &IGF,
+                                      llvm::Value *topLevelSubcontextMetatype);
+  Address emitAutoDiffAllocateSubcontext2(IRGenFunction &IGF, Address context,
+                                          llvm::Value *subcontextMetatype);
 
   FunctionPointer getFunctionPointerForDispatchCall(IRGenModule &IGM,
                                                     const FunctionPointer &fn);

lib/SIL/IR/OperandOwnership.cpp

@@ -957,6 +957,8 @@ BUILTIN_OPERAND_OWNERSHIP(BitwiseEscape, BuildMainActorExecutorRef)
 
 BUILTIN_OPERAND_OWNERSHIP(TrivialUse, AutoDiffCreateLinearMapContext)
 
+BUILTIN_OPERAND_OWNERSHIP(TrivialUse, AutoDiffCreateLinearMapContext2)
+BUILTIN_OPERAND_OWNERSHIP(PointerEscape, AutoDiffAllocateSubcontext2)
 #undef BUILTIN_OPERAND_OWNERSHIP
 
 #define SHOULD_NEVER_VISIT_BUILTIN(ID)                                         \

lib/SIL/IR/ValueOwnership.cpp

@@ -574,6 +574,8 @@ CONSTANT_OWNERSHIP_BUILTIN(None, InitializeNonDefaultDistributedActor)
 CONSTANT_OWNERSHIP_BUILTIN(Owned, AutoDiffCreateLinearMapContext)
 CONSTANT_OWNERSHIP_BUILTIN(None, AutoDiffProjectTopLevelSubcontext)
 CONSTANT_OWNERSHIP_BUILTIN(None, AutoDiffAllocateSubcontext)
+CONSTANT_OWNERSHIP_BUILTIN(Owned, AutoDiffCreateLinearMapContext2)
+CONSTANT_OWNERSHIP_BUILTIN(None, AutoDiffAllocateSubcontext2)
 CONSTANT_OWNERSHIP_BUILTIN(None, GetCurrentExecutor)
 CONSTANT_OWNERSHIP_BUILTIN(None, ResumeNonThrowingContinuationReturning)
 CONSTANT_OWNERSHIP_BUILTIN(None, ResumeThrowingContinuationReturning)

lib/SIL/Utils/MemAccessUtils.cpp

@@ -2569,6 +2569,8 @@ static void visitBuiltinAddress(BuiltinInst *builtin,
     case BuiltinValueKind::CreateAsyncTaskInGroup:
     case BuiltinValueKind::AutoDiffCreateLinearMapContext:
     case BuiltinValueKind::AutoDiffAllocateSubcontext:
+    case BuiltinValueKind::AutoDiffCreateLinearMapContext2:
+    case BuiltinValueKind::AutoDiffAllocateSubcontext2:
     case BuiltinValueKind::InitializeDefaultActor:
     case BuiltinValueKind::InitializeDistributedRemoteActor:
     case BuiltinValueKind::InitializeNonDefaultDistributedActor:

lib/SILGen/SILGenBuiltin.cpp

@@ -1750,6 +1750,32 @@ static ManagedValue emitBuiltinAutoDiffAllocateSubcontext(
   return ManagedValue::forObjectRValueWithoutOwnership(builtinApply);
 }
 
+static ManagedValue emitBuiltinAutoDiffCreateLinearMapContext2(
+    SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
+    ArrayRef<ManagedValue> args, SGFContext C) {
+  ASTContext &ctx = SGF.getASTContext();
+  auto *builtinApply = SGF.B.createBuiltin(
+      loc,
+      ctx.getIdentifier(
+          getBuiltinName(BuiltinValueKind::AutoDiffCreateLinearMapContext2)),
+      SILType::getNativeObjectType(ctx), subs,
+      /*args*/ {args[0].getValue()});
+  return SGF.emitManagedRValueWithCleanup(builtinApply);
+}
+
+static ManagedValue emitBuiltinAutoDiffAllocateSubcontext2(
+    SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
+    ArrayRef<ManagedValue> args, SGFContext C) {
+  ASTContext &ctx = SGF.getASTContext();
+  auto *builtinApply = SGF.B.createBuiltin(
+      loc,
+      ctx.getIdentifier(
+          getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontext2)),
+      SILType::getRawPointerType(ctx), subs,
+      /*args*/ {args[0].borrow(SGF, loc).getValue(), args[1].getValue()});
+  return ManagedValue::forObjectRValueWithoutOwnership(builtinApply);
+}
+
 // The only reason these need special attention is that we want these to
 // be borrowed arguments, but the default emitter doesn't handle borrowed
 // arguments correctly.

lib/SILOptimizer/Differentiation/VJPCloner.cpp

@@ -132,7 +132,7 @@ class VJPCloner::Implementation final
     pullbackContextValue = Builder.createBuiltin(
         original->getLocation(),
         getASTContext().getIdentifier(
-            getBuiltinName(BuiltinValueKind::AutoDiffCreateLinearMapContext)),
+            getBuiltinName(BuiltinValueKind::AutoDiffCreateLinearMapContext2)),
         SILType::getNativeObjectType(getASTContext()), SubstitutionMap(),
         {pbTupleMetatype});
     borrowedPullbackContextValue = Builder.createBeginBorrow(
@@ -156,8 +156,8 @@ class VJPCloner::Implementation final
       return builtinAutoDiffAllocateSubcontextGenericSignature;
     auto &ctx = getASTContext();
     auto *decl = cast<FuncDecl>(getBuiltinValueDecl(
-        ctx, ctx.getIdentifier(
-            getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontext))));
+        ctx, ctx.getIdentifier(getBuiltinName(
+                 BuiltinValueKind::AutoDiffAllocateSubcontext2))));
     builtinAutoDiffAllocateSubcontextGenericSignature =
         decl->getGenericSignature();
     assert(builtinAutoDiffAllocateSubcontextGenericSignature);
@@ -1086,7 +1086,7 @@ EnumInst *VJPCloner::Implementation::buildPredecessorEnumValue(
     auto rawBufferValue = builder.createBuiltin(
         loc,
         getASTContext().getIdentifier(
-            getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontext)),
+            getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontext2)),
         rawPtrType, SubstitutionMap(),
         {borrowedPullbackContextValue, pbTupleMetatype});
 

lib/SILOptimizer/Transforms/AccessEnforcementReleaseSinking.cpp

@@ -147,6 +147,7 @@ static bool isBarrier(SILInstruction *inst) {
     case BuiltinValueKind::GetCurrentAsyncTask:
     case BuiltinValueKind::GetCurrentExecutor:
     case BuiltinValueKind::AutoDiffCreateLinearMapContext:
+    case BuiltinValueKind::AutoDiffCreateLinearMapContext2:
     case BuiltinValueKind::EndAsyncLet:
     case BuiltinValueKind::EndAsyncLetLifetime:
     case BuiltinValueKind::CreateTaskGroup:
@@ -200,6 +201,7 @@ static bool isBarrier(SILInstruction *inst) {
     case BuiltinValueKind::ResumeThrowingContinuationThrowing:
     case BuiltinValueKind::AutoDiffProjectTopLevelSubcontext:
     case BuiltinValueKind::AutoDiffAllocateSubcontext:
+    case BuiltinValueKind::AutoDiffAllocateSubcontext2:
     case BuiltinValueKind::AddressOfBorrowOpaque:
     case BuiltinValueKind::UnprotectedAddressOfBorrowOpaque:
       return true;