[CIR] Upstream EHScopeStack memory allocator

When the cleanup handling code was initially upstreamed, a SmallVector
was used to simplify the handling of the stack of cleanup objects. However,
that mechanism won't scale well enough for the rate at which cleanup
handlers are going to be pushed and popped while compiling a large program.
This change introduces the custom memory allocator which is used in classic
codegen and the CIR incubator.

Thiis does not otherwise change the cleanup handling implementation and
many parts of the infrastructure are still missing.

This is not intended to have any observable effect on the generated CIR,
but it does change the internal implementation significantly, so it's not
exactly an NFC change. The functionality is covered by existing tests.

Author: andykaylor

clang/include/clang/CIR/MissingFeatures.h

@@ -199,6 +199,7 @@ struct MissingFeatures {
   static bool dataLayoutTypeAllocSize() { return false; }
   static bool deferredCXXGlobalInit() { return false; }
   static bool ehCleanupFlags() { return false; }
+  static bool ehCleanupScope() { return false; }
   static bool ehstackBranches() { return false; }
   static bool emitCheckedInBoundsGEP() { return false; }
   static bool emitCondLikelihoodViaExpectIntrinsic() { return false; }

clang/lib/CIR/CodeGen/CIRGenCleanup.cpp

@@ -16,6 +16,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "CIRGenCleanup.h"
 #include "CIRGenFunction.h"
 
 #include "clang/CIR/MissingFeatures.h"
@@ -33,6 +34,54 @@ using namespace clang::CIRGen;
 
 void EHScopeStack::Cleanup::anchor() {}
 
+/// Push an entry of the given size onto this protected-scope stack.
+char *EHScopeStack::allocate(size_t size) {
+  size = llvm::alignTo(size, ScopeStackAlignment);
+  if (!startOfBuffer) {
+    unsigned capacity = 1024;
+    while (capacity < size)
+      capacity *= 2;
+    startOfBuffer = new char[capacity];
+    startOfData = endOfBuffer = startOfBuffer + capacity;
+  } else if (static_cast<size_t>(startOfData - startOfBuffer) < size) {
+    unsigned currentCapacity = endOfBuffer - startOfBuffer;
+    unsigned usedCapacity = currentCapacity - (startOfData - startOfBuffer);
+
+    unsigned newCapacity = currentCapacity;
+    do {
+      newCapacity *= 2;
+    } while (newCapacity < usedCapacity + size);
+
+    char *newStartOfBuffer = new char[newCapacity];
+    char *newEndOfBuffer = newStartOfBuffer + newCapacity;
+    char *newStartOfData = newEndOfBuffer - usedCapacity;
+    memcpy(newStartOfData, startOfData, usedCapacity);
+    delete[] startOfBuffer;
+    startOfBuffer = newStartOfBuffer;
+    endOfBuffer = newEndOfBuffer;
+    startOfData = newStartOfData;
+  }
+
+  assert(startOfBuffer + size <= startOfData);
+  startOfData -= size;
+  return startOfData;
+}
+
+void EHScopeStack::deallocate(size_t size) {
+  startOfData += llvm::alignTo(size, ScopeStackAlignment);
+}
+
+void *EHScopeStack::pushCleanup(CleanupKind kind, size_t size) {
+  char *buffer = allocate(size);
+
+  // When the full implementation is upstreamed, this will allocate
+  // extra memory for and construct a wrapper object that is used to
+  // manage the cleanup generation.
+  assert(!cir::MissingFeatures::ehCleanupScope());
+
+  return buffer;
+}
+
 static mlir::Block *getCurCleanupBlock(CIRGenFunction &cgf) {
   mlir::OpBuilder::InsertionGuard guard(cgf.getBuilder());
   mlir::Block *cleanup =
@@ -44,26 +93,34 @@ static mlir::Block *getCurCleanupBlock(CIRGenFunction &cgf) {
 /// current insertion point is threaded through the cleanup, as are
 /// any branch fixups on the cleanup.
 void CIRGenFunction::popCleanupBlock() {
-  assert(!ehStack.cleanupStack.empty() && "cleanup stack is empty!");
+  assert(!ehStack.empty() && "cleanup stack is empty!");
+
+  // The memory for the cleanup continues to be owned by the EHScopeStack
+  // allocator, so we just destroy the object rather than attempting to
+  // free it.
+  EHScopeStack::Cleanup &cleanup = *ehStack.begin();
+
+  // The eventual implementation here will use the EHCleanupScope helper class.
+  assert(!cir::MissingFeatures::ehCleanupScope());
+
   mlir::OpBuilder::InsertionGuard guard(builder);
-  std::unique_ptr<EHScopeStack::Cleanup> cleanup =
-      ehStack.cleanupStack.pop_back_val();
 
   assert(!cir::MissingFeatures::ehCleanupFlags());
   mlir::Block *cleanupEntry = getCurCleanupBlock(*this);
   builder.setInsertionPointToEnd(cleanupEntry);
-  cleanup->emit(*this);
+  cleanup.emit(*this);
+
+  ehStack.deallocate(cleanup.getSize());
 }
 
 /// Pops cleanup blocks until the given savepoint is reached.
-void CIRGenFunction::popCleanupBlocks(size_t oldCleanupStackDepth) {
+void CIRGenFunction::popCleanupBlocks(
+    EHScopeStack::stable_iterator oldCleanupStackDepth) {
   assert(!cir::MissingFeatures::ehstackBranches());
 
-  assert(ehStack.getStackDepth() >= oldCleanupStackDepth);
-
   // Pop cleanup blocks until we reach the base stack depth for the
   // current scope.
-  while (ehStack.getStackDepth() > oldCleanupStackDepth) {
+  while (ehStack.stable_begin() != oldCleanupStackDepth) {
     popCleanupBlock();
   }
 }

clang/lib/CIR/CodeGen/CIRGenCleanup.h

@@ -0,0 +1,43 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// These classes support the generation of CIR for cleanups, initially based
+// on LLVM IR cleanup handling, but ought to change as CIR evolves.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef CLANG_LIB_CIR_CODEGEN_CGCLEANUP_H
+#define CLANG_LIB_CIR_CODEGEN_CGCLEANUP_H
+
+#include "EHScopeStack.h"
+
+namespace clang::CIRGen {
+
+/// A non-stable pointer into the scope stack.
+class EHScopeStack::iterator {
+  char *ptr;
+
+  friend class EHScopeStack;
+  explicit iterator(char *ptr) : ptr(ptr) {}
+
+public:
+  iterator() : ptr(nullptr) {}
+
+  EHScopeStack::Cleanup *get() const {
+    return reinterpret_cast<EHScopeStack::Cleanup *>(ptr);
+  }
+
+  EHScopeStack::Cleanup &operator*() const { return *get(); }
+};
+
+inline EHScopeStack::iterator EHScopeStack::begin() const {
+  return iterator(startOfData);
+}
+
+} // namespace clang::CIRGen
+#endif // CLANG_LIB_CIR_CODEGEN_CGCLEANUP_H

clang/lib/CIR/CodeGen/CIRGenDecl.cpp

@@ -667,6 +667,12 @@ struct DestroyObject final : EHScopeStack::Cleanup {
   void emit(CIRGenFunction &cgf) override {
     cgf.emitDestroy(addr, type, destroyer);
   }
+
+  // This is a placeholder until EHCleanupScope is implemented.
+  size_t getSize() const override {
+    assert(!cir::MissingFeatures::ehCleanupScope());
+    return sizeof(DestroyObject);
+  }
 };
 } // namespace
 

clang/lib/CIR/CodeGen/CIRGenFunction.cpp

@@ -28,8 +28,6 @@ CIRGenFunction::CIRGenFunction(CIRGenModule &cgm, CIRGenBuilderTy &builder,
                                bool suppressNewContext)
     : CIRGenTypeCache(cgm), cgm{cgm}, builder(builder) {
   ehStack.setCGF(this);
-  currentCleanupStackDepth = 0;
-  assert(ehStack.getStackDepth() == 0);
 }
 
 CIRGenFunction::~CIRGenFunction() {}
@@ -409,6 +407,8 @@ void CIRGenFunction::startFunction(GlobalDecl gd, QualType returnType,
   const auto *fd = dyn_cast_or_null<FunctionDecl>(d);
   curFuncDecl = d->getNonClosureContext();
 
+  prologueCleanupDepth = ehStack.stable_begin();
+
   mlir::Block *entryBB = &fn.getBlocks().front();
   builder.setInsertionPointToStart(entryBB);
 
@@ -475,11 +475,11 @@ void CIRGenFunction::finishFunction(SourceLocation endLoc) {
   // important to do this before we enter the return block or return
   // edges will be *really* confused.
   // TODO(cir): Use prologueCleanupDepth here.
-  bool hasCleanups = ehStack.getStackDepth() != currentCleanupStackDepth;
+  bool hasCleanups = ehStack.stable_begin() != prologueCleanupDepth;
   if (hasCleanups) {
     assert(!cir::MissingFeatures::generateDebugInfo());
     // FIXME(cir): should we clearInsertionPoint? breaks many testcases
-    popCleanupBlocks(currentCleanupStackDepth);
+    popCleanupBlocks(prologueCleanupDepth);
   }
 }
 

clang/lib/CIR/CodeGen/CIRGenFunction.h

@@ -601,9 +601,13 @@ class CIRGenFunction : public CIRGenTypeCache {
                      FunctionArgList args, clang::SourceLocation loc,
                      clang::SourceLocation startLoc);
 
+  /// The cleanup depth enclosing all the cleanups associated with the
+  /// parameters.
+  EHScopeStack::stable_iterator prologueCleanupDepth;
+
   /// Takes the old cleanup stack size and emits the cleanup blocks
   /// that have been added.
-  void popCleanupBlocks(size_t oldCleanupStackDepth);
+  void popCleanupBlocks(EHScopeStack::stable_iterator oldCleanupStackDepth);
   void popCleanupBlock();
 
   /// Push a cleanup to be run at the end of the current full-expression.  Safe
@@ -622,7 +626,7 @@ class CIRGenFunction : public CIRGenTypeCache {
   /// Enters a new scope for capturing cleanups, all of which
   /// will be executed once the scope is exited.
   class RunCleanupsScope {
-    size_t cleanupStackDepth, oldCleanupStackDepth;
+    EHScopeStack::stable_iterator cleanupStackDepth, oldCleanupStackDepth;
 
   protected:
     bool performCleanup;
@@ -638,7 +642,7 @@ class CIRGenFunction : public CIRGenTypeCache {
     /// Enter a new cleanup scope.
     explicit RunCleanupsScope(CIRGenFunction &cgf)
         : performCleanup(true), cgf(cgf) {
-      cleanupStackDepth = cgf.ehStack.getStackDepth();
+      cleanupStackDepth = cgf.ehStack.stable_begin();
       oldCleanupStackDepth = cgf.currentCleanupStackDepth;
       cgf.currentCleanupStackDepth = cleanupStackDepth;
     }
@@ -663,7 +667,7 @@ class CIRGenFunction : public CIRGenTypeCache {
   };
 
   // Cleanup stack depth of the RunCleanupsScope that was pushed most recently.
-  size_t currentCleanupStackDepth;
+  EHScopeStack::stable_iterator currentCleanupStackDepth = ehStack.stable_end();
 
 public:
   /// Represents a scope, including function bodies, compound statements, and

clang/lib/CIR/CodeGen/CIRGenStmt.cpp

@@ -412,7 +412,7 @@ mlir::LogicalResult CIRGenFunction::emitReturnStmt(const ReturnStmt &s) {
   auto *retBlock = curLexScope->getOrCreateRetBlock(*this, loc);
   // This should emit a branch through the cleanup block if one exists.
   builder.create<cir::BrOp>(loc, retBlock);
-  if (ehStack.getStackDepth() != currentCleanupStackDepth)
+  if (ehStack.stable_begin() != currentCleanupStackDepth)
     cgm.errorNYI(s.getSourceRange(), "return with cleanup stack");
   builder.createBlock(builder.getBlock()->getParent());
 

clang/lib/CIR/CodeGen/EHScopeStack.h

@@ -42,7 +42,47 @@ enum CleanupKind : unsigned {
 /// A stack of scopes which respond to exceptions, including cleanups
 /// and catch blocks.
 class EHScopeStack {
+  friend class CIRGenFunction;
+
 public:
+  // TODO(ogcg): Switch to alignof(uint64_t) instead of 8
+  enum { ScopeStackAlignment = 8 };
+
+  /// A saved depth on the scope stack.  This is necessary because
+  /// pushing scopes onto the stack invalidates iterators.
+  class stable_iterator {
+    friend class EHScopeStack;
+
+    /// Offset from startOfData to endOfBuffer.
+    ptrdiff_t size;
+
+    stable_iterator(ptrdiff_t size) : size(size) {}
+
+  public:
+    static stable_iterator invalid() { return stable_iterator(-1); }
+    stable_iterator() : size(-1) {}
+
+    bool isValid() const { return size >= 0; }
+
+    /// Returns true if this scope encloses I.
+    /// Returns false if I is invalid.
+    /// This scope must be valid.
+    bool encloses(stable_iterator other) const { return size <= other.size; }
+
+    /// Returns true if this scope strictly encloses I: that is,
+    /// if it encloses I and is not I.
+    /// Returns false is I is invalid.
+    /// This scope must be valid.
+    bool strictlyEncloses(stable_iterator I) const { return size < I.size; }
+
+    friend bool operator==(stable_iterator A, stable_iterator B) {
+      return A.size == B.size;
+    }
+    friend bool operator!=(stable_iterator A, stable_iterator B) {
+      return A.size != B.size;
+    }
+  };
+
   /// Information for lazily generating a cleanup.  Subclasses must be
   /// POD-like: cleanups will not be destructed, and they will be
   /// allocated on the cleanup stack and freely copied and moved
@@ -68,30 +108,75 @@ class EHScopeStack {
     ///
     // \param flags cleanup kind.
     virtual void emit(CIRGenFunction &cgf) = 0;
-  };
 
-  // Classic codegen has a finely tuned custom allocator and a complex stack
-  // management scheme. We'll probably eventually want to find a way to share
-  // that implementation. For now, we will use a very simplified implementation
-  // to get cleanups working.
-  llvm::SmallVector<std::unique_ptr<Cleanup>, 8> cleanupStack;
+    // This is a placeholder until EHScope is implemented.
+    virtual size_t getSize() const = 0;
+  };
 
 private:
+  // The implementation for this class is in CIRGenCleanup.h and
+  // CIRGenCleanup.cpp; the definition is here because it's used as a
+  // member of CIRGenFunction.
+
+  /// The start of the scope-stack buffer, i.e. the allocated pointer
+  /// for the buffer.  All of these pointers are either simultaneously
+  /// null or simultaneously valid.
+  char *startOfBuffer = nullptr;
+
+  /// The end of the buffer.
+  char *endOfBuffer = nullptr;
+
+  /// The first valid entry in the buffer.
+  char *startOfData = nullptr;
+
   /// The CGF this Stack belong to
   CIRGenFunction *cgf = nullptr;
 
+  // This class uses a custom allocator for maximum efficiency because cleanups
+  // are allocated and freed very frequently. It's basically a bump pointer
+  // allocator, but we can't use LLVM's BumpPtrAllocator because we use offsets
+  // into the buffer as stable iterators.
+  char *allocate(size_t size);
+  void deallocate(size_t size);
+
+  void *pushCleanup(CleanupKind kind, size_t dataSize);
+
 public:
   EHScopeStack() = default;
-  ~EHScopeStack() = default;
+  ~EHScopeStack() { delete[] startOfBuffer; }
 
   /// Push a lazily-created cleanup on the stack.
   template <class T, class... As> void pushCleanup(CleanupKind kind, As... a) {
-    cleanupStack.push_back(std::make_unique<T>(a...));
+    static_assert(alignof(T) <= ScopeStackAlignment,
+                  "Cleanup's alignment is too large.");
+    void *buffer = pushCleanup(kind, sizeof(T));
+    [[maybe_unused]] Cleanup *obj = new (buffer) T(a...);
   }
 
   void setCGF(CIRGenFunction *inCGF) { cgf = inCGF; }
 
-  size_t getStackDepth() const { return cleanupStack.size(); }
+  /// Pops a cleanup scope off the stack.  This is private to CIRGenCleanup.cpp.
+  void popCleanup();
+
+  /// Determines whether the exception-scopes stack is empty.
+  bool empty() const { return startOfData == endOfBuffer; }
+
+  /// An unstable reference to a scope-stack depth.  Invalidated by
+  /// pushes but not pops.
+  class iterator;
+
+  /// Returns an iterator pointing to the innermost EH scope.
+  iterator begin() const;
+
+  /// Create a stable reference to the top of the EH stack.  The
+  /// returned reference is valid until that scope is popped off the
+  /// stack.
+  stable_iterator stable_begin() const {
+    return stable_iterator(endOfBuffer - startOfData);
+  }
+
+  /// Create a stable reference to the bottom of the EH stack.
+  static stable_iterator stable_end() { return stable_iterator(0); }
 };
 
 } // namespace clang::CIRGen