[TPDE] Avoid ValueAssignment leaks

Avoid leaking large ValueAssignments when compilation fails by
allocating all assignments through the bump allocator, which will take
care of freeing the allocation in all cases. To avoid excessive memory
usage, all ValueAssignment allocations are now recycled in free lists.

Additionally, modify AddressSanitizer options to abort on errors to
avoid hiding the leak behind "not" in the compilation failure case.

Author: aengelke

tpde-llvm/test/diagnostics/incompatible.ll

@@ -126,3 +126,11 @@ define void @f_v4i64_2(<4 x i64>) {
 define void @f_v4i64_3([2 x <4 x i64>]) {
   ret void
 }
+
+; CHECK: type with incompatible layout at function/call: [2 x <4 x i25>]
+; CHECK-NEXT: unsupported type: [2 x <4 x i25>]
+; CHECK-NEXT: Failed to compile function f_call_v4i25_2
+define void @f_call_v4i25_2(ptr %f) {
+  call void %f([2 x <4 x i25>] zeroinitializer)
+  ret void
+}

tpde-llvm/test/diagnostics/large-types.ll

@@ -28,3 +28,13 @@ define void @f_sa40000i8a40000i8(ptr %p) {
   store {[40000 x i8], [40000 x i8]} %l, ptr %p
   ret void
 }
+
+; CHECK: Failed to compile instruction   store atomic [40000 x i8] %l
+; CHECK: Failed to compile function f_atomic_a40000i8
+; Note: check that large value assignment doesn't leak.
+define void @f_atomic_a40000i8(ptr %p) {
+  %l = load [40000 x i8], ptr %p
+  store atomic [40000 x i8] %l, ptr %p monotonic, align 1
+  store atomic [40000 x i8] %l, ptr %p monotonic, align 1
+  ret void
+}

tpde-llvm/test/lit.cfg.py

@@ -18,6 +18,8 @@
 # Tweak the PATH to include the tools dir and TPDE binaries.
 llvm_config.with_environment('PATH', config.llvm_tools_dir, append_path=True)
 llvm_config.with_environment('PATH', config.tpde_llvm_bin_dir, append_path=True)
+# Abort on ASan errors so that also tests running with "not" fail.
+llvm_config.with_environment('ASAN_OPTIONS', "abort_on_error=1", append_path=True)
 config.substitutions.append(('tpde-llc', 'tpde-llc --regular-exit'))
 config.substitutions.append(('%objdump', 'llvm-objdump -d -r --no-show-raw-insn --symbolize-operands --no-addresses --x86-asm-syntax=intel -'))
 

tpde/include/tpde/ValueAssignment.hpp

@@ -93,7 +93,11 @@ class AssignmentAllocator {
   // Free list 2 holds twice as much as free list 1. Etc.
   static constexpr size_t SlabSize = 16 * 1024;
 
-  static constexpr u32 NumFreeLists = 2;
+  // TODO: define a sensible limit on the maximum number of parts and reduce the
+  // number of free lists. Currently the limit is a little under 2**31; limiting
+  // the number to e.g. 2**16 would allow halving the number of free lists,
+  // where most are unused in the vast majority of cases.
+  static constexpr u32 NumFreeLists = 29;
   static constexpr u32 FirstPartOff = offsetof(ValueAssignment, parts);
 
 public:

tpde/src/ValueAssignment.cpp

@@ -3,6 +3,8 @@
 
 #include "tpde/ValueAssignment.hpp"
 
+#include <limits>
+
 namespace tpde {
 
 namespace {
@@ -23,6 +25,7 @@ AssignmentAllocInfo::AssignmentAllocInfo(u32 part_count) noexcept {
   alloc_size = VASize;
   free_list_idx = 0;
   if (part_count > AssignmentAllocator::NumPartsIncluded) {
+    assert(part_count < (std::numeric_limits<u32>::max() - VASize) / PartSize);
     size += (part_count - AssignmentAllocator::NumPartsIncluded) * PartSize;
     // Round size to next power of two.
     static_assert((VASize & (VASize - 1)) == 0,
@@ -40,11 +43,6 @@ ValueAssignment *
                                        bool skip_free_list) noexcept {
   AssignmentAllocInfo aai(part_count);
 
-  if (aai.free_list_idx >= fixed_free_lists.size()) [[unlikely]] {
-    auto *alloc = new std::byte[aai.size];
-    return new (reinterpret_cast<ValueAssignment *>(alloc)) ValueAssignment{};
-  }
-
   if (!skip_free_list) {
     auto &free_list = fixed_free_lists[aai.free_list_idx];
     if (auto *assignment = free_list) {
@@ -54,22 +52,16 @@ ValueAssignment *
     }
   }
 
-  assert(aai.alloc_size < SlabSize);
   auto *buf = alloc.allocate(aai.alloc_size, alignof(ValueAssignment));
   return new (reinterpret_cast<ValueAssignment *>(buf)) ValueAssignment{};
 }
 
 void AssignmentAllocator::deallocate_slow(
     ValueAssignment *assignment) noexcept {
   AssignmentAllocInfo aai(assignment->part_count);
-  if (aai.free_list_idx < fixed_free_lists.size()) [[likely]] {
-    assignment->next_free_list_entry = fixed_free_lists[aai.free_list_idx];
-    fixed_free_lists[aai.free_list_idx] = assignment;
-    util::poison_memory_region(assignment, aai.size);
-  } else {
-    assignment->~ValueAssignment();
-    delete[] reinterpret_cast<std::byte *>(assignment);
-  }
+  assignment->next_free_list_entry = fixed_free_lists[aai.free_list_idx];
+  fixed_free_lists[aai.free_list_idx] = assignment;
+  util::poison_memory_region(assignment, aai.size);
 }
 
 } // namespace tpde

tpde/test/filetest/lit.cfg.py

@@ -18,5 +18,7 @@
 # Tweak the PATH to include the tools dir.
 llvm_config.with_environment('PATH', config.llvm_tools_dir, append_path=True)
 llvm_config.with_environment('PATH', config.tpde_obj_root, append_path=True)
+# Abort on ASan errors so that also tests running with "not" fail.
+llvm_config.with_environment('ASAN_OPTIONS', "abort_on_error=1", append_path=True)
 config.substitutions.append(("%tpde_test", "tpde_test"))
 config.substitutions.append(('%objdump', 'llvm-objdump -d -r --no-show-raw-insn --symbolize-operands --no-addresses --x86-asm-syntax=intel -'))