[clang][bytecode] Optimize InterpStack

[tbaederr]

Replace StackChunk::End with StackChunk::Size, mark the allocating code paths as unlikely and move grow() into the header, which allows us to template this for the Size parameter. Since we only push our primitive types on the stack and all the sizes are aligned to pointer size multiples, this only results in a few instantiations.

[tbaederr]

https://llvm-compile-time-tracker.com/compare.php?from=a0162219165b94b950fe10547307378c0a306b74&to=03ef1e9181aaec7f0601736a42e5db842b0d4e4f&stat=instructions:u

[llvmbot]

@llvm/pr-subscribers-clang

Author: Timm Baeder (tbaederr)

Changes

Replace StackChunk::End with StackChunk::Size, mark the allocating code paths as unlikely and move grow() into the header, which allows us to template this for the Size parameter. Since we only push our primitive types on the stack and all the sizes are aligned to pointer size multiples, this only results in a few instantiations.

---

Full diff: https://github.com/llvm/llvm-project/pull/159400.diff

2 Files Affected:

• (modified) clang/lib/AST/ByteCode/InterpStack.cpp (+13-26)
• (modified) clang/lib/AST/ByteCode/InterpStack.h (+30-7)

diff --git a/clang/lib/AST/ByteCode/InterpStack.cpp b/clang/lib/AST/ByteCode/InterpStack.cpp
index 7920378f365f9..992546560eec4 100644
--- a/clang/lib/AST/ByteCode/InterpStack.cpp
+++ b/clang/lib/AST/ByteCode/InterpStack.cpp
@@ -24,9 +24,6 @@ InterpStack::~InterpStack() {
     std::free(Chunk->Next);
   if (Chunk)
     std::free(Chunk);
-  Chunk = nullptr;
-  StackSize = 0;
-  ItemTypes.clear();
 }
 
 // We keep the last chunk around to reuse.
@@ -56,29 +53,12 @@ void InterpStack::clearTo(size_t NewSize) {
   assert(size() == NewSize);
 }
 
-void *InterpStack::grow(size_t Size) {
-  assert(Size < ChunkSize - sizeof(StackChunk) && "Object too large");
-
-  if (!Chunk || sizeof(StackChunk) + Chunk->size() + Size > ChunkSize) {
-    if (Chunk && Chunk->Next) {
-      Chunk = Chunk->Next;
-    } else {
-      StackChunk *Next = new (std::malloc(ChunkSize)) StackChunk(Chunk);
-      if (Chunk)
-        Chunk->Next = Next;
-      Chunk = Next;
-    }
-  }
-
-  auto *Object = reinterpret_cast<void *>(Chunk->End);
-  Chunk->End += Size;
-  StackSize += Size;
-  return Object;
-}
-
 void *InterpStack::peekData(size_t Size) const {
   assert(Chunk && "Stack is empty!");
 
+  if (LLVM_LIKELY(Size <= Chunk->size()))
+    return reinterpret_cast<void *>(Chunk->start() + Chunk->Size - Size);
+
   StackChunk *Ptr = Chunk;
   while (Size > Ptr->size()) {
     Size -= Ptr->size();
@@ -86,24 +66,31 @@ void *InterpStack::peekData(size_t Size) const {
     assert(Ptr && "Offset too large");
   }
 
-  return reinterpret_cast<void *>(Ptr->End - Size);
+  return reinterpret_cast<void *>(Ptr->start() + Ptr->Size - Size);
 }
 
 void InterpStack::shrink(size_t Size) {
   assert(Chunk && "Chunk is empty!");
 
+  // Likely case is that we simply remove something from the current chunk.
+  if (LLVM_LIKELY(Size <= Chunk->size())) {
+    Chunk->Size -= Size;
+    StackSize -= Size;
+    return;
+  }
+
   while (Size > Chunk->size()) {
     Size -= Chunk->size();
     if (Chunk->Next) {
       std::free(Chunk->Next);
       Chunk->Next = nullptr;
     }
-    Chunk->End = Chunk->start();
+    Chunk->Size = 0;
     Chunk = Chunk->Prev;
     assert(Chunk && "Offset too large");
   }
 
-  Chunk->End -= Size;
+  Chunk->Size -= Size;
   StackSize -= Size;
 }
 
diff --git a/clang/lib/AST/ByteCode/InterpStack.h b/clang/lib/AST/ByteCode/InterpStack.h
index b0f9f6e225682..c647dfa6d85ea 100644
--- a/clang/lib/AST/ByteCode/InterpStack.h
+++ b/clang/lib/AST/ByteCode/InterpStack.h
@@ -24,14 +24,14 @@ namespace interp {
 /// Stack frame storing temporaries and parameters.
 class InterpStack final {
 public:
-  InterpStack() {}
+  InterpStack() = default;
 
   /// Destroys the stack, freeing up storage.
   ~InterpStack();
 
   /// Constructs a value in place on the top of the stack.
   template <typename T, typename... Tys> void push(Tys &&...Args) {
-    new (grow(aligned_size<T>())) T(std::forward<Tys>(Args)...);
+    new (grow<aligned_size<T>()>()) T(std::forward<Tys>(Args)...);
     ItemTypes.push_back(toPrimType<T>());
   }
 
@@ -89,7 +89,7 @@ class InterpStack final {
 private:
   /// All stack slots are aligned to the native pointer alignment for storage.
   /// The size of an object is rounded up to a pointer alignment multiple.
-  template <typename T> constexpr size_t aligned_size() const {
+  template <typename T> static constexpr size_t aligned_size() {
     constexpr size_t PtrAlign = alignof(void *);
     return ((sizeof(T) + PtrAlign - 1) / PtrAlign) * PtrAlign;
   }
@@ -100,7 +100,30 @@ class InterpStack final {
   }
 
   /// Grows the stack to accommodate a value and returns a pointer to it.
-  void *grow(size_t Size);
+  template <size_t Size> void *grow() {
+    assert(Size < ChunkSize - sizeof(StackChunk) && "Object too large");
+    static_assert(aligned(Size));
+
+    // Allocate a new stack chunk if necessary.
+    if (LLVM_UNLIKELY(!Chunk)) {
+      Chunk = new (std::malloc(ChunkSize)) StackChunk(Chunk);
+    } else if (LLVM_UNLIKELY(Chunk->size() >
+                             ChunkSize - sizeof(StackChunk) - Size)) {
+      if (Chunk->Next) {
+        Chunk = Chunk->Next;
+      } else {
+        StackChunk *Next = new (std::malloc(ChunkSize)) StackChunk(Chunk);
+        Chunk->Next = Next;
+        Chunk = Next;
+      }
+    }
+
+    auto *Object = reinterpret_cast<void *>(Chunk->start() + Chunk->Size);
+    Chunk->Size += Size;
+    StackSize += Size;
+    return Object;
+  }
+
   /// Returns a pointer from the top of the stack.
   void *peekData(size_t Size) const;
   /// Shrinks the stack.
@@ -118,13 +141,13 @@ class InterpStack final {
   struct StackChunk {
     StackChunk *Next;
     StackChunk *Prev;
-    char *End;
+    uint32_t Size;
 
     StackChunk(StackChunk *Prev = nullptr)
-        : Next(nullptr), Prev(Prev), End(reinterpret_cast<char *>(this + 1)) {}
+        : Next(nullptr), Prev(Prev), Size(0) {}
 
     /// Returns the size of the chunk, minus the header.
-    size_t size() const { return End - start(); }
+    size_t size() const { return Size; }
 
     /// Returns a pointer to the start of the data region.
     char *start() { return reinterpret_cast<char *>(this + 1); }