fix(rocjitsu): Fix local memory fallback for VMID misses

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/events.cpp

@@ -49,13 +49,25 @@ void EventState::adopt_page(void *ptr, size_t size) {
   }
 }
 
-bool EventState::release_page(void *ptr) {
+EventState::PageReleaseResult EventState::release_page(void *ptr, size_t size) {
   std::lock_guard<std::mutex> lock(mutex_);
-  if (page != ptr)
-    return false;
+  if (!page)
+    return PageReleaseResult::kNotMapped;
+
+  const uintptr_t request_begin = reinterpret_cast<uintptr_t>(ptr);
+  const uintptr_t page_begin = reinterpret_cast<uintptr_t>(page);
+  if (size > UINTPTR_MAX - request_begin || page_size > UINTPTR_MAX - page_begin)
+    return PageReleaseResult::kOverlap;
+  const uintptr_t request_end = request_begin + size;
+  const uintptr_t page_end = page_begin + page_size;
+  if (request_begin >= page_end || request_end <= page_begin)
+    return PageReleaseResult::kNotMapped;
+  if (request_begin != page_begin || size != page_size)
+    return PageReleaseResult::kOverlap;
+
   page = nullptr;
   page_size = 0;
-  return true;
+  return PageReleaseResult::kReleased;
 }
 
 /// @brief Signal event(s) from the CP's interrupt callback.

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/events.h

@@ -34,6 +34,8 @@ namespace rocjitsu {
 /// for mmap/munmap operations.
 class EventState {
 public:
+  enum class PageReleaseResult { kNotMapped, kReleased, kOverlap };
+
   ~EventState();
 
   /// @brief Handle KFD CREATE_EVENT ioctl.
@@ -74,11 +76,11 @@ class EventState {
   /// @details Idempotent, subsequent calls after the first are no-ops.
   void adopt_page(void *ptr, size_t size);
 
-  /// @brief If @p ptr is the adopted event page, clear page/page_size and return true.
-  /// @details Clears the fields under mutex_ so the CP interrupt thread (which
-  /// reads page/page_size under the same lock in signal_interrupt) can never race
-  /// a concurrent munmap tearing the mapping down.
-  [[nodiscard]] bool release_page(void *ptr);
+  /// @brief Release an exact event-page mapping or report a protected overlap.
+  /// @details Clears page/page_size only when the requested range exactly matches
+  /// the adopted page. Classification and release happen under mutex_, so the CP
+  /// interrupt thread cannot retain a pointer to an unmapped event page.
+  [[nodiscard]] PageReleaseResult release_page(void *ptr, size_t size);
 
   /// @brief Signal event(s) from the CP's interrupt callback.
   /// @details When event_id is non-zero, signals that specific event. When

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/interposer.cpp

@@ -2267,20 +2267,13 @@ RJ_INTERPOSER_EXPORT void *mmap(void *addr, size_t length, int prot, int flags,
       if (lookup == RemoteDriver::MemfdLookup::kFound) {
         // memfd_out is a caller-owned dup; close it once we are done. Its
         // validity is independent of a concurrent RemoteDriver teardown/close.
-        auto total = static_cast<off_t>(length) + memfd_offset;
-        [[maybe_unused]] auto ft_rc = ftruncate(memfd_out, total);
-        fallocate(memfd_out, 0, memfd_offset, static_cast<off_t>(length));
         auto *raw = InterposerContext::real().mmap(
             addr, length, prot, (flags & ~MAP_ANONYMOUS) | MAP_SHARED, memfd_out, memfd_offset);
         // Preserve the mmap errno across close() (which may set its own).
         int mmap_errno = errno;
         InterposerContext::real().close(memfd_out);
-        if (raw != MAP_FAILED) {
-#ifdef MADV_POPULATE_WRITE
-          InterposerContext::real().madvise(raw, length, MADV_POPULATE_WRITE);
-#endif
+        if (raw != MAP_FAILED)
           return raw;
-        }
         // A daemon-shared range matched but the shared mapping failed. Fail
         // closed with the real errno rather than falling through to an anonymous
         // MAP_FIXED mapping, which would silently detach this GPUVM address from

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/kfd_process.h

@@ -16,13 +16,17 @@
 #include "rocjitsu/vm/amdgpu/mtype.h"
 #include "util/unique_handle.h"
 
+#include <algorithm>
 #include <atomic>
 #include <cassert>
 #include <cstdint>
+#include <map>
 #include <mutex>
+#include <optional>
 #include <shared_mutex>
 #include <unordered_map>
 #include <unordered_set>
+#include <utility>
 #include <vector>
 
 #include <sys/types.h> // pid_t
@@ -79,21 +83,65 @@ class KfdProcess {
   struct GpuAllocation {
     uint64_t gpu_va = 0;
     uint64_t size = 0;
+    uint64_t backing_size = 0;
     void *host_ptr = nullptr;
     uint32_t flags = 0;
+    amdgpu::Mtype backing_mtype = amdgpu::Mtype::RW;
     uint64_t handle = 0;
     int memfd = -1;
     uint32_t gpu_id = 0;
     bool user_va = false;
     bool imported = false;
     int dmabuf_fd = -1;
+    bool permanent_backing = false;
+    bool mapped_to_gpu = false;
+    bool backing_authoritative = false;
+    // Page-aligned allocation-relative intervals already consumed from an
+    // existing MAP_FIXED client reservation. Keys are inclusive starts and
+    // values are exclusive ends.
+    std::map<uint64_t, uint64_t> imported_client_extents;
     // True when the driver created host_ptr (mmap it itself) and must munmap it
     // on teardown. False for caller-owned pages (e.g. reused MAP_FIXED pages
     // from the thunk) that the driver must never unmap, since unmapping them
     // races with the owning process still accessing the memory.
     bool host_ptr_owned = false;
   };
 
+  /// @brief A stable lookup result for allocation-owned GPU backing.
+  /// @details An inaccessible result identifies a managed address that must not
+  /// fall through to identity-pointer translation. The guard keeps accessible
+  /// backing alive until the caller finishes the memory operation.
+  struct ResolvedBacking {
+    uint8_t *address = nullptr;
+    uint8_t *page_base = nullptr;
+    uint8_t *range_base = nullptr;
+    uint64_t range_size = 0;
+    amdgpu::Mtype mtype = amdgpu::Mtype::RW;
+    bool gpu_accessible = false;
+    std::shared_lock<std::shared_mutex> guard;
+  };
+
+  [[nodiscard]] static bool is_valid_gpu_range(uint64_t gpu_va, uint64_t size) {
+    return size != 0 && gpu_va <= UINT64_MAX - (size - 1);
+  }
+
+  /// @brief Whether an allocation already covers any page in a proposed range.
+  /// @details The caller must hold alloc_mutex_. GPU mappings are page-granular,
+  /// so byte-disjoint ranges sharing a page are treated as overlapping.
+  [[nodiscard]] bool overlaps_allocation_locked(uint64_t gpu_va, uint64_t size) const {
+    assert(is_valid_gpu_range(gpu_va, size));
+    const uint64_t first_page = gpu_va >> kPageShift;
+    const uint64_t end_page = ((gpu_va + size - 1) >> kPageShift) + 1;
+    for (const auto &[handle, alloc] : allocations_) {
+      (void)handle;
+      const uint64_t alloc_first_page = alloc.gpu_va >> kPageShift;
+      const uint64_t alloc_end_page = ((alloc.gpu_va + alloc.size - 1) >> kPageShift) + 1;
+      if (first_page < alloc_end_page && alloc_first_page < end_page)
+        return true;
+    }
+    return false;
+  }
+
   /// @brief Memory policy descriptor.
   struct MemoryPolicy {
     uint64_t alternate_base = 0;
@@ -254,6 +302,41 @@ class KfdProcess {
   /// @brief Return the mutation counter used by GpuMemory translation caches.
   const uint64_t *page_table_generation() const { return &page_table_generation_; }
 
+  /// @brief Resolve a managed GPU range to allocation-owned host backing.
+  /// @details Unmapped allocations still return a result with
+  /// gpu_accessible=false, preventing unsafe identity-pointer fallback. The
+  /// returned guard prevents teardown while accessible backing is in use.
+  [[nodiscard]] std::optional<ResolvedBacking> resolve_backing(uint64_t gpu_va, size_t size) const {
+    if (size == 0 || gpu_va > UINT64_MAX - (static_cast<uint64_t>(size) - 1))
+      return std::nullopt;
+
+    std::shared_lock lock(backing_ranges_mutex_);
+    auto it = std::upper_bound(
+        backing_ranges_.begin(), backing_ranges_.end(), gpu_va,
+        [](uint64_t addr, const BackingRange &range) { return addr < range.first; });
+    if (it == backing_ranges_.begin())
+      return std::nullopt;
+    --it;
+
+    const uint64_t offset = gpu_va - it->first;
+    if (offset >= it->size || static_cast<uint64_t>(size) > it->size - offset)
+      return std::nullopt;
+
+    const uintptr_t address_value = reinterpret_cast<uintptr_t>(it->host_base) + offset;
+    return ResolvedBacking{reinterpret_cast<uint8_t *>(address_value),
+                           reinterpret_cast<uint8_t *>(address_value - (gpu_va & (kPageSize - 1))),
+                           it->host_base,
+                           it->size,
+                           it->mtype,
+                           it->gpu_accessible,
+                           std::move(lock)};
+  }
+
+  /// @brief Republish allocation-backed address ranges after a state change.
+  /// @details The caller must hold alloc_mutex_. Acquiring the unique range lock
+  /// waits for in-flight GPU accesses before backing can be unmapped.
+  void refresh_backing_ranges_locked() { publish_backing_ranges_locked(); }
+
   mutable std::shared_mutex page_table_mutex_;
   PageTable page_table_;
 
@@ -271,6 +354,9 @@ class KfdProcess {
   std::mutex op_mutex_;
   mutable std::mutex alloc_mutex_;
   std::unordered_map<uint64_t, GpuAllocation> allocations_;
+  // Page-aligned CPU alias intervals. Keys are inclusive host addresses and
+  // values are exclusive ends.
+  std::map<uintptr_t, uintptr_t> cpu_mappings_;
   uint64_t next_handle_ = 1;
   uint64_t next_gpu_va_;
 
@@ -309,6 +395,42 @@ class KfdProcess {
   ///          this generation; all reads and writes occur while holding
   ///          page_table_mutex_, so the counter itself does not need atomics.
   uint64_t page_table_generation_{1};
+
+  struct BackingRange {
+    uint64_t first = 0;
+    uint64_t size = 0;
+    uint8_t *host_base = nullptr;
+    amdgpu::Mtype mtype = amdgpu::Mtype::RW;
+    bool gpu_accessible = false;
+  };
+
+  void publish_backing_ranges_locked() {
+    std::vector<BackingRange> ranges;
+    ranges.reserve(allocations_.size());
+    for (const auto &[handle, alloc] : allocations_) {
+      (void)handle;
+      if (!alloc.permanent_backing || !alloc.host_ptr)
+        continue;
+      assert(is_valid_gpu_range(alloc.gpu_va, alloc.size));
+      ranges.push_back({alloc.gpu_va, alloc.size, static_cast<uint8_t *>(alloc.host_ptr),
+                        alloc.backing_mtype, alloc.mapped_to_gpu});
+    }
+
+    std::sort(ranges.begin(), ranges.end(), [](const BackingRange &lhs, const BackingRange &rhs) {
+      return lhs.first < rhs.first;
+    });
+    for (size_t i = 1; i < ranges.size(); ++i) {
+      assert(ranges[i - 1].first <= ranges[i].first);
+      assert(ranges[i - 1].size <= ranges[i].first - ranges[i - 1].first &&
+             "overlapping GPU allocation backings");
+    }
+
+    std::unique_lock lock(backing_ranges_mutex_);
+    backing_ranges_ = std::move(ranges);
+  }
+
+  mutable std::shared_mutex backing_ranges_mutex_;
+  std::vector<BackingRange> backing_ranges_;
 };
 
 } // namespace rocjitsu