[rocjitsu] Add intrusive shared ownership utilities and migrate consumers

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/code/analysis/indirect_branch_discovery.cpp

@@ -11,6 +11,7 @@
 #include "rocjitsu/isa/operand.h"
 #include "rocjitsu/isa/register_set.h"
 #include "rocjitsu/isa/target_registry.h"
+#include "util/intrusive_shared_ptr.h"
 
 #include <algorithm>
 #include <array>
@@ -2348,6 +2349,8 @@ class RestoredSgprFacts {
   };
 
   using Storage = std::vector<Entry>;
+  using CountedStorage = util::RefCountedNonAtomic<Storage>;
+  using StoragePtr = util::IntrusiveSharedPtr<CountedStorage>;
 
 public:
   [[nodiscard]] const StashedPcHalf *find(uint16_t sgpr) const {
@@ -2374,64 +2377,66 @@ class RestoredSgprFacts {
   void erase(uint16_t sgpr) {
     if (!values_)
       return;
-    const auto found = std::ranges::lower_bound(*values_, sgpr, {}, &Entry::sgpr);
-    if (found == values_->end() || found->sgpr != sgpr)
+    const auto found = std::ranges::lower_bound(values_->value, sgpr, {}, &Entry::sgpr);
+    if (found == values_->value.end() || found->sgpr != sgpr)
       return;
-    const size_t index = static_cast<size_t>(found - values_->begin());
+    const size_t index = static_cast<size_t>(found - values_->value.begin());
     Storage &updated = writable_values();
     updated.erase(updated.begin() + static_cast<Storage::difference_type>(index));
     reset_if_empty();
   }
 
   template <typename Predicate> void erase_if(Predicate &&predicate) {
-    if (!values_ ||
-        std::ranges::none_of(*values_, [&](const Entry &entry) { return predicate(entry.sgpr); }))
+    if (!values_ || std::ranges::none_of(values_->value,
+                                         [&](const Entry &entry) { return predicate(entry.sgpr); }))
       return;
     std::erase_if(writable_values(), [&](const Entry &entry) { return predicate(entry.sgpr); });
     reset_if_empty();
   }
 
   void intersect_with(const RestoredSgprFacts &other) {
-    if (!values_ || values_ == other.values_)
+    if (!values_ || values_.get() == other.values_.get())
       return;
     const auto conflicts = [&](const Entry &entry) {
       const StashedPcHalf *incoming = other.find(entry.sgpr);
       return incoming == nullptr || *incoming != entry.half;
     };
-    if (std::ranges::none_of(*values_, conflicts))
+    if (std::ranges::none_of(values_->value, conflicts))
       return;
     std::erase_if(writable_values(), conflicts);
     reset_if_empty();
   }
 
-  [[nodiscard]] bool empty() const { return !values_ || values_->empty(); }
-  [[nodiscard]] size_t size() const { return values_ ? values_->size() : 0; }
+  [[nodiscard]] bool empty() const { return !values_ || values_->value.empty(); }
+  [[nodiscard]] size_t size() const { return values_ ? values_->value.size() : 0; }
   void clear() { values_.reset(); }
 
   friend bool operator==(const RestoredSgprFacts &lhs, const RestoredSgprFacts &rhs) {
-    return lhs.values_ == rhs.values_ || lhs.values() == rhs.values();
+    return lhs.values_.get() == rhs.values_.get() || lhs.values() == rhs.values();
   }
 
 private:
   [[nodiscard]] const Storage &values() const {
     static const Storage empty;
-    return values_ ? *values_ : empty;
+    return values_ ? values_->value : empty;
   }
 
   [[nodiscard]] Storage &writable_values() {
     if (!values_)
-      values_ = std::make_shared<Storage>();
-    else if (!values_.unique())
-      values_ = std::make_shared<Storage>(*values_);
-    return *values_;
+      values_ = StoragePtr::make();
+    else if (values_->ref_count() != 1)
+      values_ = StoragePtr::make(values_->value);
+    return values_->value;
   }
 
   void reset_if_empty() {
-    if (values_->empty())
+    if (values_->value.empty())
       values_.reset();
   }
 
-  std::shared_ptr<Storage> values_;
+  // All owners are local flow states in one synchronous analysis invocation.
+  // Concurrent invocations allocate independent backing storage.
+  StoragePtr values_;
 };
 
 struct VectorLaneFlowState {

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

@@ -17,7 +17,9 @@
 #include "rocjitsu/vm/amdgpu/gpu_handles.h"
 #include "rocjitsu/vm/amdgpu/legacy_page_table.h"
 #include "rocjitsu/vm/amdgpu/mtype.h"
+#include "rocjitsu/vm/amdgpu/vm_mutation_epoch.h"
 #include "util/distributed_shared_mutex.h"
+#include "util/intrusive_shared_ptr.h"
 #include "util/unique_handle.h"
 
 #include <algorithm>
@@ -513,7 +515,7 @@ class KfdProcess {
   const uint64_t *page_table_generation() const { return &page_table_generation_; }
 
   /// @brief Return the retained token invalidated before each page-policy mutation.
-  std::shared_ptr<const std::atomic<uint64_t>> page_table_mutation_epoch() const {
+  util::IntrusiveSharedPtr<const amdgpu::VmMutationEpoch> page_table_mutation_epoch() const {
     return page_table_mutation_epoch_;
   }
 
@@ -891,7 +893,7 @@ class KfdProcess {
   // Invalidate copied policies before mutation, including a partial update
   // that throws before publishing the ordinary translation generation.
   void invalidate_page_policies_locked() {
-    page_table_mutation_epoch_->fetch_add(1, std::memory_order_release);
+    page_table_mutation_epoch_->value.fetch_add(1, std::memory_order_release);
     page_table_cache_state_->invalidate_and_wait();
   }
 
@@ -906,15 +908,12 @@ class KfdProcess {
   ///          and page_table_mutex_; readers hold at least one of those locks,
   ///          so the counter itself does not need atomics.
   uint64_t page_table_generation_{1};
-  // Cache hits read the epoch while misses copy and release its shared_ptr.
-  // Keep the epoch off the cache line containing make_shared's reference count.
-  struct alignas(64) PageTableMutationEpoch : std::atomic<uint64_t> {
-    PageTableMutationEpoch() : std::atomic<uint64_t>(1) {}
-  };
+  // Cache hits read the epoch while misses copy and release its owner.
+  // Keep the epoch off the cache line containing the reference count.
   /// @brief Retained atomic token for lockless copied-policy checks.
   /// @details Invalidate before mutation, including partially throwing updates.
-  std::shared_ptr<std::atomic<uint64_t>> page_table_mutation_epoch_ =
-      std::make_shared<PageTableMutationEpoch>();
+  util::IntrusiveSharedPtr<amdgpu::VmMutationEpoch> page_table_mutation_epoch_ =
+      util::IntrusiveSharedPtr<amdgpu::VmMutationEpoch>::make(1);
   std::shared_ptr<amdgpu::LegacyPageTableCacheState> page_table_cache_state_ =
       std::make_shared<amdgpu::LegacyPageTableCacheState>();
 };

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

@@ -137,7 +137,7 @@ class LegacyGpuVmAdapter::Binding final : public AddressSpaceTranslator,
             .begin = address & ~(kLegacyPageSize - 1),
             .size = kLegacyPageSize,
             .mutation_epoch = mutation_epoch_,
-            .captured_epoch = mutation_epoch_->load(std::memory_order_relaxed)};
+            .captured_epoch = mutation_epoch_->value.load(std::memory_order_relaxed)};
   }
 
   [[nodiscard]] VmTranslationResult probe_translation(uint64_t address, std::size_t size,
@@ -302,7 +302,7 @@ class LegacyGpuVmAdapter::Binding final : public AddressSpaceTranslator,
   // Synchronous word reads and stores must not destroy the final request-mutex owner
   // under the caller's VM-state lock, even after raw registry replacement.
   const std::shared_ptr<util::DistributedSharedMutex> original_request_mutex_;
-  std::shared_ptr<const std::atomic<uint64_t>> mutation_epoch_;
+  util::IntrusiveSharedPtr<const amdgpu::VmMutationEpoch> mutation_epoch_;
   std::shared_ptr<void> frontend_lifetime_;
 };
 

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/amdgpu/command_processor.cpp

@@ -573,7 +573,7 @@ VmAccessOutcome CommandProcessor::init_wavefront_regs(ComputeUnitCore *cu, Wavef
       for (uint32_t preload_index = 0; preload_index < preload_length; ++preload_index) {
         const uint64_t address = preload_addr + preload_index * 4;
         const AtomicLoadResult loaded = pkt.execution_access
-                                            ? read_gpu_u32(*pkt.execution_access, address)
+                                            ? read_gpu_u32(pkt.execution_access->value, address)
                                             : read_gpu_u32(pkt.address_space, address);
         if (loaded.outcome != VmAccessOutcome::Complete)
           return loaded.outcome;
@@ -744,7 +744,8 @@ VmAccessOutcome CommandProcessor::init_wavefront_regs(ComputeUnitCore *cu, Wavef
       return VmAccessOutcome::Malformed;
     uint64_t wave_scratch = scratch_pool + scratch_slot * per_wave_size;
     std::optional<GpuVmAccess> fallback_scratch_access;
-    const GpuVmAccess *scratch_access = pkt.execution_access.get();
+    const GpuVmAccess *scratch_access =
+        pkt.execution_access ? &pkt.execution_access->value : nullptr;
     if (scratch_access == nullptr) {
       fallback_scratch_access = snapshot_gpu_access(pkt.address_space);
       scratch_access = fallback_scratch_access ? &*fallback_scratch_access : nullptr;
@@ -2367,7 +2368,7 @@ CommandProcessor::dispatch_workgroups(DispatchEntry &entry) {
   // A dispatch may remain queued after its admission snapshot is invalidated.
   // Fall back to current operation snapshots for subsequent wave setup rather
   // than retrying the permanently revoked snapshot forever.
-  if (entry.execution_access && entry.execution_access->revoked())
+  if (entry.execution_access && entry.execution_access->value.revoked())
     entry.execution_access.reset();
 
   // A peer can publish the shared terminal-fault latch before this CP drains
@@ -3482,7 +3483,7 @@ AqlAdmissionResult CommandProcessor::request_dynamic_scratch(ComputeQueueRecord
     if (queue.queue_desc_va == 0)
       return {.status = AqlAdmissionStatus::Faulted};
     state.phase = QueueScratchRequestPhase::PublishReadPointer;
-    state.access = std::make_shared<GpuVmAccess>(transaction_access);
+    state.access = util::IntrusiveSharedPtr<GpuVmAccessSnapshot>::make(transaction_access);
     state.packet_index = packet_index;
     state.status = status;
     util::Logger::cp([&](auto &os) {
@@ -3503,7 +3504,8 @@ AqlAdmissionResult CommandProcessor::request_dynamic_scratch(ComputeQueueRecord
       // Earlier dispatches have retired; leave the requesting packet itself
       // unconsumed at packet_index.
       if (!queue.read_pointer_journal.publication_pending()) {
-        queue.read_pointer_journal.retire(state.packet_index, state.packet_index, *state.access);
+        queue.read_pointer_journal.retire(state.packet_index, state.packet_index,
+                                          state.access->value);
       }
       const VmAccessOutcome outcome = queue.read_pointer_journal.publish();
       if (outcome != VmAccessOutcome::Complete)
@@ -3514,7 +3516,7 @@ AqlAdmissionResult CommandProcessor::request_dynamic_scratch(ComputeQueueRecord
 
     case QueueScratchRequestPhase::ReadSignalHandle: {
       const VmAccessOutcome outcome =
-          read_gpu_block(*state.access, queue.queue_desc_va + kQueueInactiveSignalOffset,
+          read_gpu_block(state.access->value, queue.queue_desc_va + kQueueInactiveSignalOffset,
                          &state.signal_address, sizeof(state.signal_address));
       if (outcome != VmAccessOutcome::Complete)
         return admission_from_vm_outcome(outcome);
@@ -3525,7 +3527,7 @@ AqlAdmissionResult CommandProcessor::request_dynamic_scratch(ComputeQueueRecord
     }
 
     case QueueScratchRequestPhase::StoreStatus: {
-      const AtomicCompareExchangeResult stored = state.access->compare_exchange(
+      const AtomicCompareExchangeResult stored = state.access->value.compare_exchange(
           state.signal_address + kSignalValueOffset, sizeof(uint64_t), 0, state.status);
       if (stored.outcome != VmAccessOutcome::Complete)
         return admission_from_vm_outcome(stored.outcome);
@@ -3537,7 +3539,7 @@ AqlAdmissionResult CommandProcessor::request_dynamic_scratch(ComputeQueueRecord
 
     case QueueScratchRequestPhase::ReadMailboxPointer: {
       const VmAccessOutcome outcome =
-          read_gpu_block(*state.access, state.signal_address + kMailboxPointerOffset,
+          read_gpu_block(state.access->value, state.signal_address + kMailboxPointerOffset,
                          &state.mailbox_pointer, sizeof(state.mailbox_pointer));
       if (outcome != VmAccessOutcome::Complete)
         return admission_from_vm_outcome(outcome);
@@ -3547,8 +3549,8 @@ AqlAdmissionResult CommandProcessor::request_dynamic_scratch(ComputeQueueRecord
 
     case QueueScratchRequestPhase::ReadEventId: {
       const VmAccessOutcome outcome =
-          read_gpu_block(*state.access, state.signal_address + kEventIdOffset, &state.event_id,
-                         sizeof(state.event_id));
+          read_gpu_block(state.access->value, state.signal_address + kEventIdOffset,
+                         &state.event_id, sizeof(state.event_id));
       if (outcome != VmAccessOutcome::Complete)
         return admission_from_vm_outcome(outcome);
       state.phase = QueueScratchRequestPhase::StoreMailbox;
@@ -3557,7 +3559,7 @@ AqlAdmissionResult CommandProcessor::request_dynamic_scratch(ComputeQueueRecord
 
     case QueueScratchRequestPhase::StoreMailbox:
       if (state.mailbox_pointer != 0) {
-        const VmAccessOutcome outcome = state.access->atomic_store(
+        const VmAccessOutcome outcome = state.access->value.atomic_store(
             state.mailbox_pointer, sizeof(uint64_t), uint64_t(state.event_id));
         if (outcome != VmAccessOutcome::Complete)
           return admission_from_vm_outcome(outcome);
@@ -3614,15 +3616,15 @@ VmAccessOutcome CommandProcessor::advance_scratch_reclaim(ComputeQueueRecord &qu
       if (queue.queue_desc_va == 0 || !entry.address_space)
         return VmAccessOutcome::Faulted;
       if (std::optional<GpuVmAccess> access = snapshot_gpu_access(entry.address_space))
-        state.access = std::make_shared<GpuVmAccess>(std::move(*access));
+        state.access = util::IntrusiveSharedPtr<GpuVmAccessSnapshot>::make(std::move(*access));
       else
         return VmAccessOutcome::Faulted;
       state.phase = QueueScratchReclaimPhase::ReadSignalHandle;
       continue;
 
     case QueueScratchReclaimPhase::ReadSignalHandle: {
       const VmAccessOutcome outcome =
-          read_gpu_block(*state.access, queue.queue_desc_va + kQueueInactiveSignalOffset,
+          read_gpu_block(state.access->value, queue.queue_desc_va + kQueueInactiveSignalOffset,
                          &state.signal_address, sizeof(state.signal_address));
       if (outcome != VmAccessOutcome::Complete)
         return outcome;
@@ -3633,7 +3635,7 @@ VmAccessOutcome CommandProcessor::advance_scratch_reclaim(ComputeQueueRecord &qu
     }
 
     case QueueScratchReclaimPhase::StoreStatus: {
-      const AtomicCompareExchangeResult stored = state.access->compare_exchange(
+      const AtomicCompareExchangeResult stored = state.access->value.compare_exchange(
           state.signal_address + kSignalValueOffset, sizeof(uint64_t), 0, kLargeScratchReclaim);
       if (stored.outcome != VmAccessOutcome::Complete)
         return stored.outcome;
@@ -3645,7 +3647,7 @@ VmAccessOutcome CommandProcessor::advance_scratch_reclaim(ComputeQueueRecord &qu
 
     case QueueScratchReclaimPhase::ReadMailboxPointer: {
       const VmAccessOutcome outcome =
-          read_gpu_block(*state.access, state.signal_address + kMailboxPointerOffset,
+          read_gpu_block(state.access->value, state.signal_address + kMailboxPointerOffset,
                          &state.mailbox_pointer, sizeof(state.mailbox_pointer));
       if (outcome != VmAccessOutcome::Complete)
         return outcome;
@@ -3655,8 +3657,8 @@ VmAccessOutcome CommandProcessor::advance_scratch_reclaim(ComputeQueueRecord &qu
 
     case QueueScratchReclaimPhase::ReadEventId: {
       const VmAccessOutcome outcome =
-          read_gpu_block(*state.access, state.signal_address + kEventIdOffset, &state.event_id,
-                         sizeof(state.event_id));
+          read_gpu_block(state.access->value, state.signal_address + kEventIdOffset,
+                         &state.event_id, sizeof(state.event_id));
       if (outcome != VmAccessOutcome::Complete)
         return outcome;
       state.phase = QueueScratchReclaimPhase::StoreMailbox;
@@ -3665,7 +3667,7 @@ VmAccessOutcome CommandProcessor::advance_scratch_reclaim(ComputeQueueRecord &qu
 
     case QueueScratchReclaimPhase::StoreMailbox:
       if (state.mailbox_pointer != 0) {
-        const VmAccessOutcome outcome = state.access->atomic_store(
+        const VmAccessOutcome outcome = state.access->value.atomic_store(
             state.mailbox_pointer, sizeof(uint64_t), uint64_t(state.event_id));
         if (outcome != VmAccessOutcome::Complete)
           return outcome;
@@ -3683,24 +3685,24 @@ VmAccessOutcome CommandProcessor::advance_scratch_reclaim(ComputeQueueRecord &qu
 
     case QueueScratchReclaimPhase::WaitingForRuntime: {
       const AtomicLoadResult signal_status =
-          read_gpu_u64(*state.access, state.signal_address + kSignalValueOffset);
+          read_gpu_u64(state.access->value, state.signal_address + kSignalValueOffset);
       if (signal_status.outcome != VmAccessOutcome::Complete)
         return signal_status.outcome;
 
       uint32_t queue_properties = 0;
       VmAccessOutcome outcome =
-          read_gpu_block(*state.access, queue.queue_desc_va + kQueuePropertiesOffset,
+          read_gpu_block(state.access->value, queue.queue_desc_va + kQueuePropertiesOffset,
                          &queue_properties, sizeof(queue_properties));
       if (outcome != VmAccessOutcome::Complete)
         return outcome;
 
       const AtomicLoadResult scratch_backing =
-          read_gpu_u64(*state.access, queue.queue_desc_va + kScratchBackingOffset);
+          read_gpu_u64(state.access->value, queue.queue_desc_va + kScratchBackingOffset);
       if (scratch_backing.outcome != VmAccessOutcome::Complete)
         return scratch_backing.outcome;
 
       uint32_t compute_tmpring_size = 0;
-      outcome = read_gpu_block(*state.access, queue.queue_desc_va + kComputeTmpringSizeOffset,
+      outcome = read_gpu_block(state.access->value, queue.queue_desc_va + kComputeTmpringSizeOffset,
                                &compute_tmpring_size, sizeof(compute_tmpring_size));
       if (outcome != VmAccessOutcome::Complete)
         return outcome;
@@ -3756,7 +3758,7 @@ VmAccessOutcome CommandProcessor::retire_aql_packet(ComputeQueueRecord &queue,
     // retirement starts. Capture a current view then; once slot release starts,
     // its retained snapshot must finish or fail without switching generations.
     std::optional<GpuVmAccess> refreshed;
-    const GpuVmAccess *access = entry.execution_access.get();
+    const GpuVmAccess *access = entry.execution_access ? &entry.execution_access->value : nullptr;
     if (!access || access->revoked()) {
       refreshed = snapshot_gpu_access(queue.address_space);
       if (!refreshed)
@@ -4042,7 +4044,7 @@ AqlAdmissionResult CommandProcessor::admit_kernel_dispatch(
   dp.enabled_cus = queue.enabled_cus;
   dp.queue_packet_id = queue_packet_id;
   dp.address_space = queue.address_space;
-  dp.execution_access = std::make_shared<GpuVmAccess>(transaction_access);
+  dp.execution_access = util::IntrusiveSharedPtr<GpuVmAccessSnapshot>::make(transaction_access);
   if (!uses_kfd_queue_abi && private_segment_fixed_size > 0 && scratch_resolver_)
     dp.scratch_backing_addr = scratch_resolver_(queue.process_id);
   dp.interrupt_sink = queue.interrupt_sink;
@@ -4304,7 +4306,7 @@ AqlAdmissionResult CommandProcessor::admit_aql_packet(const AqlPacketProcessRequ
       .dispatch_id = allocate_dispatch_id(),
       .queue_id = queue->queue_id,
       .address_space = queue->address_space,
-      .execution_access = std::make_shared<GpuVmAccess>(request.access),
+      .execution_access = util::IntrusiveSharedPtr<GpuVmAccessSnapshot>::make(request.access),
       .interrupt_sink = queue->interrupt_sink,
       .process_id = queue->process_id,
       .aql_packet_id = request.packet_index,