Fixes memory leaks in autodiff linear map context allocation builtins

jkshtj · jkshtj · commit 2ccc0b31efb3 · 2023-08-17T13:15:33.000-07:00
Fixes #67323
diff --git a/include/swift/AST/Builtins.def b/include/swift/AST/Builtins.def
@@ -985,13 +985,13 @@ 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: (Builtin.Word) -> Builtin.NativeObject
+// autoDiffCreateLinearMapContext: (T.Type) -> 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, Builtin.Word) -> Builtin.RawPointer
+// autoDiffAllocateSubcontext: (Builtin.NativeObject, T.Type) -> Builtin.RawPointer
 BUILTIN_MISC_OPERATION_WITH_SILGEN(AutoDiffAllocateSubcontext, "autoDiffAllocateSubcontext", "", Special)
 
 /// Build a Builtin.Executor value from an "ordinary" serial executor
diff --git a/include/swift/Runtime/RuntimeFunctions.def b/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 *);
 FUNCTION(AutoDiffCreateLinearMapContext,
          swift_autoDiffCreateLinearMapContext, SwiftCC,
          DifferentiationAvailability,
          RETURNS(RefCountedPtrTy),
-         ARGS(SizeTy),
+         ARGS(TypeMetadataPtrTy),
          ATTRS(NoUnwind, ArgMemOnly),
          EFFECT(AutoDiff))
 
@@ -2291,12 +2291,12 @@ FUNCTION(AutoDiffProjectTopLevelSubcontext,
          ATTRS(NoUnwind, ArgMemOnly),
          EFFECT(AutoDiff))
 
-// void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, size_t);
+// void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, const Metadata *);
 FUNCTION(AutoDiffAllocateSubcontext,
          swift_autoDiffAllocateSubcontext, SwiftCC,
          DifferentiationAvailability,
          RETURNS(Int8PtrTy),
-         ARGS(RefCountedPtrTy, SizeTy),
+         ARGS(RefCountedPtrTy, TypeMetadataPtrTy),
          ATTRS(NoUnwind, ArgMemOnly),
          EFFECT(AutoDiff))
 
diff --git a/lib/AST/Builtins.cpp b/lib/AST/Builtins.cpp
@@ -1608,8 +1608,9 @@ static ValueDecl *getBuildComplexEqualitySerialExecutorRef(ASTContext &ctx,
 
 static ValueDecl *getAutoDiffCreateLinearMapContext(ASTContext &ctx,
                                                     Identifier id) {
-  return getBuiltinFunction(
-      id, {BuiltinIntegerType::getWordType(ctx)}, ctx.TheNativeObjectType);
+  return getBuiltinFunction(ctx, id, _thin, _generics(_unrestricted),
+                            _parameters(_metatype(_typeparam(0))),
+                            _nativeObject);
 }
 
 static ValueDecl *getAutoDiffProjectTopLevelSubcontext(ASTContext &ctx,
@@ -1621,8 +1622,8 @@ static ValueDecl *getAutoDiffProjectTopLevelSubcontext(ASTContext &ctx,
 static ValueDecl *getAutoDiffAllocateSubcontext(ASTContext &ctx,
                                                 Identifier id) {
   return getBuiltinFunction(
-      id, {ctx.TheNativeObjectType, BuiltinIntegerType::getWordType(ctx)},
-      ctx.TheRawPointerType);
+      ctx, id, _thin, _generics(_unrestricted),
+      _parameters(_nativeObject, _metatype(_typeparam(0))), _rawPointer);
 }
 
 static ValueDecl *getPoundAssert(ASTContext &Context, Identifier Id) {
diff --git a/lib/IRGen/GenBuiltin.cpp b/lib/IRGen/GenBuiltin.cpp
@@ -1308,8 +1308,8 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
   }
 
   if (Builtin.ID == BuiltinValueKind::AutoDiffCreateLinearMapContext) {
-    auto topLevelSubcontextSize = args.claimNext();
-    out.add(emitAutoDiffCreateLinearMapContext(IGF, topLevelSubcontextSize)
+    auto topLevelSubcontextMetaType = args.claimNext();
+    out.add(emitAutoDiffCreateLinearMapContext(IGF, topLevelSubcontextMetaType)
                 .getAddress());
     return;
   }
@@ -1325,9 +1325,10 @@ 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();
+    auto subcontextMetatype = args.claimNext();
     out.add(
-        emitAutoDiffAllocateSubcontext(IGF, allocatorAddr, size).getAddress());
+        emitAutoDiffAllocateSubcontext(IGF, allocatorAddr, subcontextMetatype)
+            .getAddress());
     return;
   }
 
diff --git a/lib/IRGen/GenCall.cpp b/lib/IRGen/GenCall.cpp
@@ -5480,10 +5480,12 @@ IRGenFunction::getFunctionPointerForResumeIntrinsic(llvm::Value *resume) {
 }
 
 Address irgen::emitAutoDiffCreateLinearMapContext(
-    IRGenFunction &IGF, llvm::Value *topLevelSubcontextSize) {
+    IRGenFunction &IGF, llvm::Value *topLevelSubcontextMetatype) {
+  topLevelSubcontextMetatype = IGF.Builder.CreateBitCast(
+      topLevelSubcontextMetatype, IGF.IGM.TypeMetadataPtrTy);
   auto *call = IGF.Builder.CreateCall(
       IGF.IGM.getAutoDiffCreateLinearMapContextFunctionPointer(),
-      {topLevelSubcontextSize});
+      {topLevelSubcontextMetatype});
   call->setDoesNotThrow();
   call->setCallingConv(IGF.IGM.SwiftCC);
   return Address(call, IGF.IGM.RefCountedStructTy,
@@ -5500,11 +5502,14 @@ 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,
+                                              Address context,
+                                              llvm::Value *subcontextMetatype) {
+  subcontextMetatype =
+      IGF.Builder.CreateBitCast(subcontextMetatype, IGF.IGM.TypeMetadataPtrTy);
   auto *call = IGF.Builder.CreateCall(
       IGF.IGM.getAutoDiffAllocateSubcontextFunctionPointer(),
-      {context.getAddress(), size});
+      {context.getAddress(), subcontextMetatype});
   call->setDoesNotThrow();
   call->setCallingConv(IGF.IGM.SwiftCC);
   return Address(call, IGF.IGM.Int8Ty, IGF.IGM.getPointerAlignment());
diff --git a/lib/IRGen/GenCall.h b/lib/IRGen/GenCall.h
@@ -261,12 +261,13 @@ namespace irgen {
                        CanSILFunctionType fnType, Explosion &result,
                        Explosion &error);
 
-  Address emitAutoDiffCreateLinearMapContext(
-      IRGenFunction &IGF, llvm::Value *topLevelSubcontextSize);
+  Address
+  emitAutoDiffCreateLinearMapContext(IRGenFunction &IGF,
+                                     llvm::Value *topLevelSubcontextMetatype);
   Address emitAutoDiffProjectTopLevelSubcontext(
       IRGenFunction &IGF, Address context);
-  Address emitAutoDiffAllocateSubcontext(
-      IRGenFunction &IGF, Address context, llvm::Value *size);
+  Address emitAutoDiffAllocateSubcontext(IRGenFunction &IGF, Address context,
+                                         llvm::Value *subcontextMetatype);
 
   FunctionPointer getFunctionPointerForDispatchCall(IRGenModule &IGM,
                                                     const FunctionPointer &fn);
diff --git a/lib/SILOptimizer/Differentiation/VJPCloner.cpp b/lib/SILOptimizer/Differentiation/VJPCloner.cpp
@@ -17,6 +17,8 @@
 
 #define DEBUG_TYPE "differentiation"
 
+#include "swift/AST/Types.h"
+
 #include "swift/SILOptimizer/Differentiation/VJPCloner.h"
 #include "swift/SILOptimizer/Analysis/DifferentiableActivityAnalysis.h"
 #include "swift/SILOptimizer/Differentiation/ADContext.h"
@@ -118,15 +120,21 @@ class VJPCloner::Implementation final
     auto pullbackTupleType =
       remapASTType(pullbackInfo.getLinearMapTupleType(returnBB)->getCanonicalType());
     Builder.setInsertionPoint(vjp->getEntryBlock());
-    auto topLevelSubcontextSize = emitMemoryLayoutSize(
-        Builder, original->getLocation(), pullbackTupleType);
+
+    auto pbTupleMetatypeType =
+        CanMetatypeType::get(pullbackTupleType, MetatypeRepresentation::Thick);
+    auto pbTupleMetatypeSILType =
+        SILType::getPrimitiveObjectType(pbTupleMetatypeType);
+    auto pbTupleMetatype =
+        Builder.createMetatype(original->getLocation(), pbTupleMetatypeSILType);
+
     // Create an context.
     pullbackContextValue = Builder.createBuiltin(
         original->getLocation(),
         getASTContext().getIdentifier(
             getBuiltinName(BuiltinValueKind::AutoDiffCreateLinearMapContext)),
-        SILType::getNativeObjectType(getASTContext()),
-        SubstitutionMap(), {topLevelSubcontextSize});
+        SILType::getNativeObjectType(getASTContext()), SubstitutionMap(),
+        {pbTupleMetatype});
     borrowedPullbackContextValue = Builder.createBeginBorrow(
         original->getLocation(), pullbackContextValue);
     LLVM_DEBUG(getADDebugStream()
@@ -1067,14 +1075,21 @@ EnumInst *VJPCloner::Implementation::buildPredecessorEnumValue(
     assert(enumEltType == rawPtrType);
     auto pbTupleType =
       remapASTType(pullbackInfo.getLinearMapTupleType(predBB)->getCanonicalType());
-    SILValue pbTupleSize =
-        emitMemoryLayoutSize(Builder, loc, pbTupleType);
+
+    auto pbTupleMetatypeType =
+        CanMetatypeType::get(pbTupleType, MetatypeRepresentation::Thick);
+    auto pbTupleMetatypeSILType =
+        SILType::getPrimitiveObjectType(pbTupleMetatypeType);
+    auto pbTupleMetatype =
+        Builder.createMetatype(original->getLocation(), pbTupleMetatypeSILType);
+
     auto rawBufferValue = builder.createBuiltin(
         loc,
         getASTContext().getIdentifier(
             getBuiltinName(BuiltinValueKind::AutoDiffAllocateSubcontext)),
         rawPtrType, SubstitutionMap(),
-        {borrowedPullbackContextValue, pbTupleSize});
+        {borrowedPullbackContextValue, pbTupleMetatype});
+
     auto typedBufferValue =
       builder.createPointerToAddress(
         loc, rawBufferValue, pbTupleVal->getType().getAddressType(),
diff --git a/stdlib/public/runtime/AutoDiffSupport.cpp b/stdlib/public/runtime/AutoDiffSupport.cpp
@@ -13,15 +13,21 @@
 #include "AutoDiffSupport.h"
 #include "swift/ABI/Metadata.h"
 #include "swift/Runtime/HeapObject.h"
-
+#include "llvm/ADT/SmallVector.h"
 #include <new>
 
 using namespace swift;
 using namespace llvm;
 
 SWIFT_CC(swift)
 static void destroyLinearMapContext(SWIFT_CONTEXT HeapObject *obj) {
-  static_cast<AutoDiffLinearMapContext *>(obj)->~AutoDiffLinearMapContext();
+  auto *linearMapContext = static_cast<AutoDiffLinearMapContext *>(obj);
+
+  for (auto *heapObjectPtr : linearMapContext->getAllocatedHeapObjects()) {
+    swift_release(heapObjectPtr);
+  }
+
+  linearMapContext->~AutoDiffLinearMapContext();
   free(obj);
 }
 
@@ -43,36 +49,54 @@ static FullMetadata<HeapMetadata> linearMapContextHeapMetadata = {
   }
 };
 
-AutoDiffLinearMapContext::AutoDiffLinearMapContext()
+AutoDiffLinearMapContext::AutoDiffLinearMapContext(
+    OpaqueValue *const topLevelLinearMapContextProjection)
     : HeapObject(&linearMapContextHeapMetadata) {
+  this->topLevelLinearMapContextProjection = topLevelLinearMapContextProjection;
 }
 
-void *AutoDiffLinearMapContext::projectTopLevelSubcontext() const {
-  auto offset = alignTo(
-      sizeof(AutoDiffLinearMapContext), alignof(AutoDiffLinearMapContext));
-  return const_cast<uint8_t *>(
-      reinterpret_cast<const uint8_t *>(this) + offset);
-}
+AutoDiffLinearMapContext *swift::swift_autoDiffCreateLinearMapContext(
+    const Metadata *topLevelLinearMapContextMetadata) {
+  // Linear map context metadata must have non-null value witnesses
+  assert(topLevelLinearMapContextMetadata->getValueWitnesses());
 
-void *AutoDiffLinearMapContext::allocate(size_t size) {
-  return allocator.Allocate(size, alignof(AutoDiffLinearMapContext));
-}
+  // Allocate a box for the top-level linear map context
+  auto [topLevelContextHeapObjectPtr, toplevelContextProjection] =
+      swift_allocBox(topLevelLinearMapContextMetadata);
 
-AutoDiffLinearMapContext *swift::swift_autoDiffCreateLinearMapContext(
-    size_t topLevelLinearMapStructSize) {
-  auto allocationSize = alignTo(
-      sizeof(AutoDiffLinearMapContext), alignof(AutoDiffLinearMapContext))
-      + topLevelLinearMapStructSize;
-  auto *buffer = (AutoDiffLinearMapContext *)malloc(allocationSize);
-  return ::new (buffer) AutoDiffLinearMapContext;
+  // Create a linear map context object that stores the projection
+  // for the top level context
+  auto linearMapContext =
+      new AutoDiffLinearMapContext(toplevelContextProjection);
+
+  // Stash away the `HeapObject` pointer for the allocated context
+  // for proper "release" during clean up.
+  linearMapContext->storeAllocatedHeapObjectPtr(topLevelContextHeapObjectPtr);
+
+  // Return the newly created linear map context object
+  return linearMapContext;
 }
 
 void *swift::swift_autoDiffProjectTopLevelSubcontext(
-    AutoDiffLinearMapContext *allocator) {
-  return allocator->projectTopLevelSubcontext();
+    AutoDiffLinearMapContext *linearMapContext) {
+  return static_cast<void *>(
+      linearMapContext->getTopLevelLinearMapContextProjection());
 }
 
 void *swift::swift_autoDiffAllocateSubcontext(
-    AutoDiffLinearMapContext *allocator, size_t size) {
-  return allocator->allocate(size);
+    AutoDiffLinearMapContext *linearMapContext,
+    const Metadata *linearMapSubcontextMetadata) {
+  // Linear map context metadata must have non-null value witnesses
+  assert(linearMapSubcontextMetadata->getValueWitnesses());
+
+  // Allocate a box for the linear map subcontext
+  auto [subcontextHeapObjectPtr, subcontextProjection] =
+      swift_allocBox(linearMapSubcontextMetadata);
+
+  // Stash away the `HeapObject` pointer for the allocated context
+  // for proper "release" during clean up.
+  linearMapContext->storeAllocatedHeapObjectPtr(subcontextHeapObjectPtr);
+
+  // Return the subcontext projection
+  return static_cast<void *>(subcontextProjection);
 }
diff --git a/stdlib/public/runtime/AutoDiffSupport.h b/stdlib/public/runtime/AutoDiffSupport.h
@@ -13,44 +13,71 @@
 #ifndef SWIFT_RUNTIME_AUTODIFF_SUPPORT_H
 #define SWIFT_RUNTIME_AUTODIFF_SUPPORT_H
 
-#include "swift/Runtime/HeapObject.h"
 #include "swift/Runtime/Config.h"
-#include "llvm/Support/Allocator.h"
+#include "swift/Runtime/HeapObject.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
 
 namespace swift {
-
 /// A data structure responsible for efficiently allocating closure contexts for
 /// linear maps such as pullbacks, including recursive branching trace enum
 /// case payloads.
 class AutoDiffLinearMapContext : public HeapObject {
 private:
-  /// The underlying allocator.
-  // TODO: Use a custom allocator so that the initial slab can be
-  // tail-allocated.
-  llvm::BumpPtrAllocator allocator;
+  // TODO: Commenting out BumpPtrAllocator temporarily
+  // until we move away from the interim solution of allocating
+  // boxes for linear map contexts/subcontexts.
+  //
+  // /// The underlying allocator.
+  // // TODO: Use a custom allocator so that the initial slab can be
+  // // tail-allocated.
+  // llvm::BumpPtrAllocator allocator;
+
+  /// A projection/pointer to the memory storing the
+  /// top-level linear map context object.
+  OpaqueValue *topLevelLinearMapContextProjection;
+
+  /// Storage for `HeapObject` pointers to the linear map
+  /// context and subcontexts allocated for derivatives with
+  /// loops.
+  llvm::SmallVector<HeapObject *, 4> allocatedHeapObjects;
 
 public:
   /// Creates a linear map context.
-  AutoDiffLinearMapContext();
-  /// Returns the address of the tail-allocated top-level subcontext.
-  void *projectTopLevelSubcontext() const;
-  /// Allocates memory for a new subcontext.
-  void *allocate(size_t size);
+  AutoDiffLinearMapContext(OpaqueValue *const);
+
+  // TODO: Commenting out BumpPtrAllocator temporarily
+  // until we move away from the interim solution of allocating
+  // boxes for linear map contexts/subcontexts.
+  //
+  // llvm::BumpPtrAllocator& getAllocator() const {
+  //   return const_cast<llvm::BumpPtrAllocator&>(this->allocator);
+  // }
+
+  OpaqueValue *getTopLevelLinearMapContextProjection() const {
+    return this->topLevelLinearMapContextProjection;
+  }
+
+  llvm::ArrayRef<HeapObject *> getAllocatedHeapObjects() const {
+    return this->allocatedHeapObjects;
+  }
+
+  void storeAllocatedHeapObjectPtr(HeapObject *allocatedHeapObjectPtr) {
+    this->allocatedHeapObjects.push_back(allocatedHeapObjectPtr);
+  }
 };
 
 /// Creates a linear map context with a tail-allocated top-level subcontext.
 SWIFT_RUNTIME_EXPORT SWIFT_CC(swift)
-AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext(
-    size_t topLevelSubcontextSize);
+    AutoDiffLinearMapContext *swift_autoDiffCreateLinearMapContext(
+        const Metadata *topLevelLinearMapContextMetadata);
 
 /// Returns the address of the tail-allocated top-level subcontext.
 SWIFT_RUNTIME_EXPORT SWIFT_CC(swift)
 void *swift_autoDiffProjectTopLevelSubcontext(AutoDiffLinearMapContext *);
 
 /// Allocates memory for a new subcontext.
-SWIFT_RUNTIME_EXPORT SWIFT_CC(swift)
-void *swift_autoDiffAllocateSubcontext(AutoDiffLinearMapContext *, size_t size);
-
-}
-
+SWIFT_RUNTIME_EXPORT SWIFT_CC(swift) void *swift_autoDiffAllocateSubcontext(
+    AutoDiffLinearMapContext *, const Metadata *linearMapSubcontextMetadata);
+} // namespace swift
 #endif /* SWIFT_RUNTIME_AUTODIFF_SUPPORT_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
diff --git a/test/AutoDiff/stdlib/callee_differential_not_leaked_in_func_with_loops.swift b/test/AutoDiff/stdlib/callee_differential_not_leaked_in_func_with_loops.swift
