[llvm] [llvm-rtdyld] Preallocate Memory Slab (PR #71409)

[Sjoerd Meijer via llvm-commits]

https://github.com/sjoerdmeijer updated https://github.com/llvm/llvm-project/pull/71409

>From 73a33c23e737d1ea26266536ae4f40dcb2e9ff67 Mon Sep 17 00:00:00 2001
From: Sjoerd Meijer <smeijer at nvidia.com>
Date: Mon, 6 Nov 2023 15:53:42 +0000
Subject: [PATCH] [llvm-rtdyld] Preallocate Memory Slab (#71409)

There's an option --preallocate that allocates memory upfront, rather than
on-demand, but it didn't seem to work at all, nor were there any tests.

This fixes the range check `CurrentSlabOffset + Size > SlabSize` that will
always result in:

  "Can't allocate enough memory. Tune --preallocate"

because `CurrentSlabOffset` is an address, typically a big number, and
`SlabSize` doesn't account for the start-address of the memory buffer.

This is a first step towards implementing the idea presented here:

https://discourse.llvm.org/t/llvm-rtdyld-aarch64-abi-relocation-restrictions/74616

The next step is to calculate the required memory rather than setting it on the
command-line.
---
 .../AArch64/ELF_ARM64_preallocate.s           | 21 +++++++++++++++++++
 llvm/tools/llvm-rtdyld/llvm-rtdyld.cpp        |  4 +++-
 2 files changed, 24 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/ExecutionEngine/RuntimeDyld/AArch64/ELF_ARM64_preallocate.s

diff --git a/llvm/test/ExecutionEngine/RuntimeDyld/AArch64/ELF_ARM64_preallocate.s b/llvm/test/ExecutionEngine/RuntimeDyld/AArch64/ELF_ARM64_preallocate.s
new file mode 100644
index 000000000000000..b6dbe6c3fc6cd4d
--- /dev/null
+++ b/llvm/test/ExecutionEngine/RuntimeDyld/AArch64/ELF_ARM64_preallocate.s
@@ -0,0 +1,21 @@
+# RUN: llvm-mc -triple=aarch64-none-linux-gnu -filetype=obj -o %t %s
+# RUN: llvm-rtdyld -triple=aarch64-none-linux-gnu -execute --entry=foo --preallocate=0 -check=%s %t
+# RUN: llvm-rtdyld -triple=aarch64-none-linux-gnu -execute --entry=foo --preallocate=1024 -check=%s %t
+
+
+       .text
+       .globl  foo
+       .p2align        2
+       .type   foo, at function
+foo:
+       adrp    x1, .L.str
+       add     x1, x1, :lo12:.L.str
+        mov     w0, wzr
+        ret
+.Lfunc_end0:
+       .size   foo, .Lfunc_end0-foo
+       .type   .L.str, at object
+       .section        .rodata.str1.1,"aMS", at progbits,1
+.L.str:
+       .asciz  "foo"
+       .size   .L.str, 4
diff --git a/llvm/tools/llvm-rtdyld/llvm-rtdyld.cpp b/llvm/tools/llvm-rtdyld/llvm-rtdyld.cpp
index 107b555a99faa40..7567838ef624655 100644
--- a/llvm/tools/llvm-rtdyld/llvm-rtdyld.cpp
+++ b/llvm/tools/llvm-rtdyld/llvm-rtdyld.cpp
@@ -260,12 +260,13 @@ class TrivialMemoryManager : public RTDyldMemoryManager {
     PreallocSlab = MB;
     UsePreallocation = true;
     SlabSize = Size;
+    SlabStart = CurrentSlabOffset = (uintptr_t)MB.base();
   }
 
   uint8_t *allocateFromSlab(uintptr_t Size, unsigned Alignment, bool isCode,
                             StringRef SectionName, unsigned SectionID) {
     Size = alignTo(Size, Alignment);
-    if (CurrentSlabOffset + Size > SlabSize)
+    if (CurrentSlabOffset + Size > SlabStart + SlabSize)
       report_fatal_error("Can't allocate enough memory. Tune --preallocate");
 
     uintptr_t OldSlabOffset = CurrentSlabOffset;
@@ -283,6 +284,7 @@ class TrivialMemoryManager : public RTDyldMemoryManager {
   sys::MemoryBlock PreallocSlab;
   bool UsePreallocation = false;
   uintptr_t SlabSize = 0;
+  uintptr_t SlabStart = 0;
   uintptr_t CurrentSlabOffset = 0;
   SectionIDMap *SecIDMap = nullptr;
 #if defined(__x86_64__) && defined(__ELF__) && defined(__linux__)