[llvm-branch-commits] [llvm] [Support] Recycler: Enforce minimum allocation size (PR #121425)

[Akshat Oke via llvm-branch-commits]

https://github.com/optimisan updated https://github.com/llvm/llvm-project/pull/121425

>From 686d1cd44e2eea89457e3ad7be1977d97ebf198b Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Thu, 19 Dec 2024 08:33:48 +0000
Subject: [PATCH 1/2] [CodeGen] LiveRegMatrix: Use allocator through a
 unique_ptr

---
 llvm/include/llvm/CodeGen/LiveRegMatrix.h | 10 +++-------
 llvm/lib/CodeGen/LiveRegMatrix.cpp        |  2 +-
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/LiveRegMatrix.h b/llvm/include/llvm/CodeGen/LiveRegMatrix.h
index 486392ca3c49d5..a7c1e2bcfe1c43 100644
--- a/llvm/include/llvm/CodeGen/LiveRegMatrix.h
+++ b/llvm/include/llvm/CodeGen/LiveRegMatrix.h
@@ -48,7 +48,7 @@ class LiveRegMatrix {
   unsigned UserTag = 0;
 
   // The matrix is represented as a LiveIntervalUnion per register unit.
-  LiveIntervalUnion::Allocator LIUAlloc;
+  std::unique_ptr<LiveIntervalUnion::Allocator> LIUAlloc;
   LiveIntervalUnion::Array Matrix;
 
   // Cached queries per register unit.
@@ -59,15 +59,11 @@ class LiveRegMatrix {
   unsigned RegMaskVirtReg = 0;
   BitVector RegMaskUsable;
 
-  LiveRegMatrix() = default;
+  LiveRegMatrix() : LIUAlloc(new LiveIntervalUnion::Allocator()) {};
   void releaseMemory();
 
 public:
-  LiveRegMatrix(LiveRegMatrix &&Other)
-      : TRI(Other.TRI), LIS(Other.LIS), VRM(Other.VRM), UserTag(Other.UserTag),
-        Matrix(std::move(Other.Matrix)), Queries(std::move(Other.Queries)),
-        RegMaskTag(Other.RegMaskTag), RegMaskVirtReg(Other.RegMaskVirtReg),
-        RegMaskUsable(std::move(Other.RegMaskUsable)) {}
+  LiveRegMatrix(LiveRegMatrix &&Other) = default;
 
   void init(MachineFunction &MF, LiveIntervals &LIS, VirtRegMap &VRM);
 
diff --git a/llvm/lib/CodeGen/LiveRegMatrix.cpp b/llvm/lib/CodeGen/LiveRegMatrix.cpp
index 9744c47d5a8510..3367171a15662f 100644
--- a/llvm/lib/CodeGen/LiveRegMatrix.cpp
+++ b/llvm/lib/CodeGen/LiveRegMatrix.cpp
@@ -66,7 +66,7 @@ void LiveRegMatrix::init(MachineFunction &MF, LiveIntervals &pLIS,
   unsigned NumRegUnits = TRI->getNumRegUnits();
   if (NumRegUnits != Matrix.size())
     Queries.reset(new LiveIntervalUnion::Query[NumRegUnits]);
-  Matrix.init(LIUAlloc, NumRegUnits);
+  Matrix.init(*LIUAlloc, NumRegUnits);
 
   // Make sure no stale queries get reused.
   invalidateVirtRegs();

>From 4496358f0a2d81710b0d3e451ccb93962c8978a7 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Wed, 1 Jan 2025 06:33:26 +0000
Subject: [PATCH 2/2] [Support] Recycler: Enforce minimum allocation size

Recycler uses reinterpret_cast to an internal structure of size 8.
Invalid write occurs if Recycler is used for objects with sizes less
than 8.
---
 llvm/include/llvm/Support/Recycler.h    |  1 +
 llvm/unittests/Support/CMakeLists.txt   |  1 +
 llvm/unittests/Support/RecyclerTest.cpp | 46 +++++++++++++++++++++++++
 3 files changed, 48 insertions(+)
 create mode 100644 llvm/unittests/Support/RecyclerTest.cpp

diff --git a/llvm/include/llvm/Support/Recycler.h b/llvm/include/llvm/Support/Recycler.h
index bbd9ae321ae30c..2bf30b0b3d5247 100644
--- a/llvm/include/llvm/Support/Recycler.h
+++ b/llvm/include/llvm/Support/Recycler.h
@@ -85,6 +85,7 @@ class Recycler {
                   "Recycler allocation alignment is less than object align!");
     static_assert(sizeof(SubClass) <= Size,
                   "Recycler allocation size is less than object size!");
+    static_assert(Size >= 8 && "Recycler size must be atleast 8");
     return FreeList ? reinterpret_cast<SubClass *>(pop_val())
                     : static_cast<SubClass *>(Allocator.Allocate(Size, Align));
   }
diff --git a/llvm/unittests/Support/CMakeLists.txt b/llvm/unittests/Support/CMakeLists.txt
index d64f89847aa8e7..6de81658264420 100644
--- a/llvm/unittests/Support/CMakeLists.txt
+++ b/llvm/unittests/Support/CMakeLists.txt
@@ -69,6 +69,7 @@ add_llvm_unittest(SupportTests
   PerThreadBumpPtrAllocatorTest.cpp
   ProcessTest.cpp
   ProgramTest.cpp
+  RecyclerTest.cpp
   RegexTest.cpp
   ReverseIterationTest.cpp
   ReplaceFileTest.cpp
diff --git a/llvm/unittests/Support/RecyclerTest.cpp b/llvm/unittests/Support/RecyclerTest.cpp
new file mode 100644
index 00000000000000..8cd763c0b83f8a
--- /dev/null
+++ b/llvm/unittests/Support/RecyclerTest.cpp
@@ -0,0 +1,46 @@
+//===--- unittest/Support/RecyclerTest.cpp --------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Support/Recycler.h"
+#include "llvm/Support/AllocatorBase.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+struct Object1 {
+  char Data[1];
+};
+
+class DecoratedMallocAllocator : public MallocAllocator {
+public:
+  int DeallocCount = 0;
+
+  template <typename T> void Deallocate(T *Ptr) {
+    DeallocCount++;
+    MallocAllocator::Deallocate(Ptr);
+  }
+};
+
+TEST(RecyclerTest, RecycleAllocation) {
+  DecoratedMallocAllocator Allocator;
+  // Recycler needs size to be atleast 8 bytes.
+  Recycler<Object1, 8, 8> R;
+  Object1 *A1 = R.Allocate(Allocator);
+  Object1 *A2 = R.Allocate(Allocator);
+  R.Deallocate(Allocator, A2);
+  Object1 *A3 = R.Allocate(Allocator);
+  EXPECT_EQ(A2, A3); // reuse the deallocated object.
+  R.Deallocate(Allocator, A1);
+  R.Deallocate(Allocator, A3);
+  R.clear(Allocator); // Should deallocate A1 and A3.
+  EXPECT_EQ(Allocator.DeallocCount, 2);
+}
+
+} // end anonymous namespace