[llvm] ab492f6 - [Orc] Take offset inside slab into account in SharedMemoryMapper

[Anubhab Ghosh via llvm-commits]

Author: Anubhab Ghosh
Date: 2022-08-27T07:49:48+05:30
New Revision: ab492f628216fbca2ae2a24c8fd61ca3f9f70a04

URL: https://github.com/llvm/llvm-project/commit/ab492f628216fbca2ae2a24c8fd61ca3f9f70a04
DIFF: https://github.com/llvm/llvm-project/commit/ab492f628216fbca2ae2a24c8fd61ca3f9f70a04.diff

LOG: [Orc] Take offset inside slab into account in SharedMemoryMapper

SharedMemoryMapper assumed each reservation will have its corresponding
allocations starting from the beginning. However with the introduction
of the slab allocator, there can be a possible offset from the start
from where the initialization is being performed.

This commit also simplifies the logic for finding the parent reservation
and makes the assert messages consistent.

Added: 
    

Modified: 
    llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp b/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp
index 773e6d26eea5a..115289d1c056c 100644
--- a/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp
@@ -292,7 +292,7 @@ void SharedMemoryMapper::reserve(size_t NumBytes,
 
 char *SharedMemoryMapper::prepare(ExecutorAddr Addr, size_t ContentSize) {
   auto R = Reservations.upper_bound(Addr);
-  assert(R != Reservations.begin() && "Attempt to prepare unknown range");
+  assert(R != Reservations.begin() && "Attempt to prepare unreserved range");
   R--;
 
   ExecutorAddrDiff Offset = Addr - R->first;
@@ -302,11 +302,11 @@ char *SharedMemoryMapper::prepare(ExecutorAddr Addr, size_t ContentSize) {
 
 void SharedMemoryMapper::initialize(MemoryMapper::AllocInfo &AI,
                                     OnInitializedFunction OnInitialized) {
-  auto Reservation = std::lower_bound(
-      Reservations.rbegin(), Reservations.rend(), AI.MappingBase,
-      [](const auto &A, const auto &B) { return A.first > B; });
-  assert(Reservation != Reservations.rend() &&
-         "Attempt to initialize unreserved range");
+  auto Reservation = Reservations.upper_bound(AI.MappingBase);
+  assert(Reservation != Reservations.begin() && "Attempt to initialize unreserved range");
+  Reservation--;
+
+  auto AllocationOffset = AI.MappingBase - Reservation->first;
 
   tpctypes::SharedMemoryFinalizeRequest FR;
 
@@ -315,8 +315,8 @@ void SharedMemoryMapper::initialize(MemoryMapper::AllocInfo &AI,
   FR.Segments.reserve(AI.Segments.size());
 
   for (auto Segment : AI.Segments) {
-    char *Base =
-        static_cast<char *>(Reservation->second.LocalAddr) + Segment.Offset;
+    char *Base = static_cast<char *>(Reservation->second.LocalAddr) +
+                 AllocationOffset + Segment.Offset;
     std::memset(Base + Segment.ContentSize, 0, Segment.ZeroFillSize);
 
     tpctypes::SharedMemorySegFinalizeRequest SegReq;