Creates new autodiff loop context allocation builtins and preserves old ones

jkshtj · jkshtj · commit aa9aab8410f3 · 2023-08-23T21:44:44.000-06:00
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.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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<ManagedValue> 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<ManagedValue> 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
@@ -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<FuncDecl>(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
@@ -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
@@ -43,6 +43,10 @@ static FullMetadata<HeapMetadata> 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<uint8_t *>(reinterpret_cast<const uint8_t *>(this) +
-                               offset);
+  auto offset = alignTo(
+      sizeof(AutoDiffLinearMapContext), alignof(AutoDiffLinearMapContext));
+  return const_cast<uint8_t *>(
+      reinterpret_cast<const uint8_t *>(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<void *>(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<void *>(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
diff --git a/test/AutoDiff/IRGen/runtime.swift b/test/AutoDiff/IRGen/runtime.swift
diff --git a/test/AutoDiff/SILGen/autodiff_builtins.swift b/test/AutoDiff/SILGen/autodiff_builtins.swift
diff --git a/test/AutoDiff/SILGen/autodiff_builtins_deprecated_abi_fails.swift b/test/AutoDiff/SILGen/autodiff_builtins_deprecated_abi_fails.swift
