Reintegrates BumpPtrAllocator in AutoDiffLinearMapContext

This commit builds on top of a previous commit that fixed memory leak
issues in the autodiff loop context allocation builtins. The previous
commit allocated loop contexts in a box, which is known to have memory
scaling issues. This commit modifies AutoDiffLinearMapContext to perform
memory allocation using the BumpPtrAllocator, like it used to, while
introducing new bookkeeping logic that enables AutoDiffLinearMapContext
to prevent memory leaks, unlike before.

Author: jkshtj

stdlib/public/runtime/AutoDiffSupport.cpp

@@ -21,13 +21,7 @@ using namespace llvm;
 
 SWIFT_CC(swift)
 static void destroyLinearMapContext(SWIFT_CONTEXT HeapObject *obj) {
-  auto *linearMapContext = static_cast<AutoDiffLinearMapContext *>(obj);
-
-  for (auto *heapObjectPtr : linearMapContext->getAllocatedHeapObjects()) {
-    swift_release(heapObjectPtr);
-  }
-
-  linearMapContext->~AutoDiffLinearMapContext();
+  static_cast<AutoDiffLinearMapContext *>(obj)->~AutoDiffLinearMapContext();
   free(obj);
 }
 
@@ -50,53 +44,48 @@ static FullMetadata<HeapMetadata> linearMapContextHeapMetadata = {
 };
 
 AutoDiffLinearMapContext::AutoDiffLinearMapContext(
-    OpaqueValue *const topLevelLinearMapContextProjection)
+    const Metadata *topLevelLinearMapContextMetadata)
     : HeapObject(&linearMapContextHeapMetadata) {
-  this->topLevelLinearMapContextProjection = topLevelLinearMapContextProjection;
+  allocatedContextObjects.push_back(AllocatedContextObjectRecord{
+      topLevelLinearMapContextMetadata, projectTopLevelSubcontext()});
 }
 
-AutoDiffLinearMapContext *swift::swift_autoDiffCreateLinearMapContext(
-    const Metadata *topLevelLinearMapContextMetadata) {
-  // Linear map context metadata must have non-null value witnesses
-  assert(topLevelLinearMapContextMetadata->getValueWitnesses());
-
-  // Allocate a box for the top-level linear map context
-  auto [topLevelContextHeapObjectPtr, toplevelContextProjection] =
-      swift_allocBox(topLevelLinearMapContextMetadata);
-
-  // Create a linear map context object that stores the projection
-  // for the top level context
-  auto linearMapContext =
-      new AutoDiffLinearMapContext(toplevelContextProjection);
+void *AutoDiffLinearMapContext::projectTopLevelSubcontext() const {
+  auto offset = alignTo(sizeof(AutoDiffLinearMapContext),
+                        alignof(AutoDiffLinearMapContext));
+  return const_cast<uint8_t *>(reinterpret_cast<const uint8_t *>(this) +
+                               offset);
+}
 
-  // Stash away the `HeapObject` pointer for the allocated context
-  // for proper "release" during clean up.
-  linearMapContext->storeAllocatedHeapObjectPtr(topLevelContextHeapObjectPtr);
+void *AutoDiffLinearMapContext::allocateSubcontext(
+    const Metadata *contextObjectMetadata) {
+  auto size = contextObjectMetadata->vw_size();
+  auto align = contextObjectMetadata->vw_alignment();
+  auto *contextObjectPtr = allocator.Allocate(size, align);
+  allocatedContextObjects.push_back(
+      AllocatedContextObjectRecord{contextObjectMetadata, contextObjectPtr});
+  return contextObjectPtr;
+}
 
-  // Return the newly created linear map context object
-  return linearMapContext;
+AutoDiffLinearMapContext *swift::swift_autoDiffCreateLinearMapContext(
+    const Metadata *topLevelLinearMapContextMetadata) {
+  auto topLevelLinearMapContextSize =
+      topLevelLinearMapContextMetadata->vw_size();
+  auto allocationSize = alignTo(sizeof(AutoDiffLinearMapContext),
+                                alignof(AutoDiffLinearMapContext)) +
+                        topLevelLinearMapContextSize;
+  auto *buffer = (AutoDiffLinearMapContext *)malloc(allocationSize);
+  return ::new (buffer)
+      AutoDiffLinearMapContext(topLevelLinearMapContextMetadata);
 }
 
 void *swift::swift_autoDiffProjectTopLevelSubcontext(
     AutoDiffLinearMapContext *linearMapContext) {
-  return static_cast<void *>(
-      linearMapContext->getTopLevelLinearMapContextProjection());
+  return static_cast<void *>(linearMapContext->projectTopLevelSubcontext());
 }
 
 void *swift::swift_autoDiffAllocateSubcontext(
     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);
+  return linearMapContext->allocateSubcontext(linearMapSubcontextMetadata);
 }

stdlib/public/runtime/AutoDiffSupport.h

@@ -13,58 +13,68 @@
 #ifndef SWIFT_RUNTIME_AUTODIFF_SUPPORT_H
 #define SWIFT_RUNTIME_AUTODIFF_SUPPORT_H
 
+#include "swift/ABI/Metadata.h"
 #include "swift/Runtime/Config.h"
 #include "swift/Runtime/HeapObject.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Allocator.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 {
+  /// A simple wrapper around a context object allocated by the
+  /// `AutoDiffLinearMapContext` type. This type knows all the "physical"
+  /// properties and behavior of the allocated context object by way of
+  /// storing the allocated type's `TypeMetadata`. It uses this information
+  /// to ensure that the allocated context object is destroyed/deinitialized
+  /// properly, upon its own destruction.
+  class [[nodiscard]] AllocatedContextObjectRecord final {
+    const Metadata *contextObjectMetadata;
+    OpaqueValue *contextObjectPtr;
+
+  public:
+    AllocatedContextObjectRecord(const Metadata *contextObjectMetadata,
+                                 OpaqueValue *contextObjectPtr)
+        : contextObjectMetadata(contextObjectMetadata),
+          contextObjectPtr(contextObjectPtr) {}
+
+    AllocatedContextObjectRecord(const Metadata *contextObjectMetadata,
+                                 void *contextObjectPtr)
+        : AllocatedContextObjectRecord(
+              contextObjectMetadata,
+              static_cast<OpaqueValue *>(contextObjectPtr)) {}
+
+    ~AllocatedContextObjectRecord() {
+      contextObjectMetadata->vw_destroy(contextObjectPtr);
+    }
+
+    size_t size() const { return contextObjectMetadata->vw_size(); }
+
+    size_t align() const { return contextObjectMetadata->vw_alignment(); }
+  };
+
 private:
-  // 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;
+  /// The underlying allocator.
+  // TODO: Use a custom allocator so that the initial slab can be
+  // tail-allocated.
+  llvm::BumpPtrAllocator allocator;
+
+  /// Storage for `AllocatedContextObjectRecord`s, corresponding to the
+  /// subcontext allocations performed by the type.
+  llvm::SmallVector<AllocatedContextObjectRecord, 4> allocatedContextObjects;
 
 public:
   /// Creates a linear map context.
-  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);
-  }
+  AutoDiffLinearMapContext(const Metadata *topLevelLinearMapContextMetadata);
+
+  /// Returns the address of the tail-allocated top-level subcontext.
+  void *projectTopLevelSubcontext() const;
+
+  /// Allocates memory for a new subcontext.
+  void *allocateSubcontext(const Metadata *contextObjectMetadata);
 };
 
 /// Creates a linear map context with a tail-allocated top-level subcontext.