Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
147 changes: 125 additions & 22 deletions src/multiprocess/multiprocess_memory_limit.c
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
#include <stdlib.h>
#include <errno.h>
#include <stddef.h>
#include <sched.h>
#include <stdint.h>
#include <semaphore.h>
#include <unistd.h>
Expand Down Expand Up @@ -307,6 +308,17 @@ size_t get_gpu_memory_monitor(const int dev) {
}

// Lock-free memory usage aggregation with seqlock for consistent snapshots
static inline void seqlock_cpu_relax(void) {
#if defined(__x86_64__) || defined(__i386__)
__asm__ __volatile__("pause" ::: "memory");
#elif defined(__aarch64__)
__asm__ __volatile__("yield" ::: "memory");
#else
__asm__ __volatile__("" ::: "memory"); /* barrier only, no pause opcode */
#endif
}


size_t get_gpu_memory_usage(const int dev) {
LOG_INFO("get_gpu_memory_usage_lockfree dev=%d", dev);
ensure_initialized();
Expand All @@ -332,12 +344,9 @@ size_t get_gpu_memory_usage(const int dev) {
while (seq1 & 1) {
// Exponential backoff to reduce contention
if (retry_count < 5) {
// First 5 retries: just CPU pause (fast path)
#if defined(__x86_64__) || defined(__i386__)
__asm__ __volatile__("pause" ::: "memory");
#elif defined(__aarch64__)
__asm__ __volatile__("yield" ::: "memory");
#endif
// First 5 retries: just CPU pause (fast path). Shared with
// the writer so both sides get the portable fallback.
seqlock_cpu_relax();
} else if (retry_count < 20) {
// Next 15 retries: 1μs delay
usleep(1);
Expand Down Expand Up @@ -476,6 +485,100 @@ uint64_t nvml_get_device_memory_usage(const int dev) {
return usage;
}

/* Seqlock writers must exclude each other, or the sequence counter can be
* driven back to even while another writer's stores are still in flight and a
* reader will accept a torn snapshot.
*
* A process-local mutex is not sufficient here: the slow paths below write
* another process's slot, so two writers to the same slot can live in
* different processes. Exclusion therefore has to sit in the shared region
* itself, which is what the CAS provides.
*
* Only the writer that wins the even->odd transition enters; a waiting writer
* never publishes an odd value, so the window in which readers see a write in
* progress is no longer than it is today.
*/
/* The critical section is a handful of atomic adds with no syscall in it, so a
* real wait is nanoseconds. Escalate anyway so a waiter never sits on a core:
* pause, then yield, then short sleeps. */
#define SEQLOCK_SPIN_LIMIT 1000u /* pause only */
#define SEQLOCK_YIELD_LIMIT 11000u /* then sched_yield */
#define SEQLOCK_RECLAIM_AFTER 31000u /* then start trying to recover */
#define SEQLOCK_RECLAIM_EVERY 10000u /* and retry about once a second */

/* A holder killed between begin and end leaves the sequence odd for good, which
* blocks every later writer and every reader of that slot. Recovery has to prove
* the holder is gone rather than infer it from a timeout, because a live writer
* can be descheduled for longer than any threshold we could pick.
*
* The writer therefore records its pid in the slot while it holds the sequence,
* and this only resets the counter when that pid no longer exists. Taking
* lock_shrreg() serialises the reclaim so two waiters cannot both perform it.
*
* Known residual: a writer killed in the window between winning the CAS and
* storing its pid leaves the sequence odd with seq_owner still 0. That is
* treated as an unknown holder and never reclaimed, so the slot stays wedged as
* it would have before this change. Closing it means publishing the sequence
* and the owner in one atomic word, which changes how the reader interprets the
* counter. */
static void seqlock_try_reclaim(shrreg_proc_slot_t* slot) {
lock_shrreg();
uint64_t seq = atomic_load_explicit(&slot->seqlock, memory_order_acquire);
int32_t owner = atomic_load_explicit(&slot->seq_owner, memory_order_relaxed);
if ((seq & 1) && owner > 0 && kill(owner, 0) == -1 && errno == ESRCH) {
LOG_WARN("seqlock owner %d is gone, reclaiming slot", owner);
/* the dead writer stopped partway through, so this slot's counters no
* longer describe anything and are cleared with the sequence */
int dev;
for (dev = 0; dev < CUDA_DEVICE_MAX_COUNT; dev++) {
atomic_store_explicit(&slot->used[dev].total, 0, memory_order_relaxed);
atomic_store_explicit(&slot->used[dev].context_size, 0, memory_order_relaxed);
atomic_store_explicit(&slot->used[dev].module_size, 0, memory_order_relaxed);
atomic_store_explicit(&slot->used[dev].data_size, 0, memory_order_relaxed);
}
atomic_store_explicit(&slot->seq_owner, 0, memory_order_relaxed);
atomic_store_explicit(&slot->seqlock, seq + 1, memory_order_release);
}
unlock_shrreg();
}

static inline void seqlock_write_begin(shrreg_proc_slot_t* slot) {
uint64_t seq;
unsigned spins = 0;
for (;;) {
seq = atomic_load_explicit(&slot->seqlock, memory_order_relaxed);
if (seq & 1) { /* another writer holds the slot */
if (spins < SEQLOCK_SPIN_LIMIT) {
seqlock_cpu_relax();
} else if (spins < SEQLOCK_YIELD_LIMIT) {
sched_yield();
} else {
usleep(100);
if (spins >= SEQLOCK_RECLAIM_AFTER &&
(spins - SEQLOCK_RECLAIM_AFTER) % SEQLOCK_RECLAIM_EVERY == 0)
seqlock_try_reclaim(slot);
}
spins++;
continue;
}
/* Acquire is all this needs: it stops the stores in the critical section
* being hoisted above the claim. Nothing is published here, and a reader
* that sees an odd sequence only backs off, so release would pair with
* nothing. seqlock_write_end carries the release. */
if (atomic_compare_exchange_weak_explicit(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why acq_rel on success? begin publishes nothing yet. would acquire not be enough?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, it should be enough, I'll update

&slot->seqlock, &seq, seq + 1,
memory_order_acquire, memory_order_relaxed)) {
atomic_store_explicit(&slot->seq_owner, getpid(), memory_order_relaxed);
return;
}
}
}

static inline void seqlock_write_end(shrreg_proc_slot_t* slot) {
atomic_store_explicit(&slot->seq_owner, 0, memory_order_relaxed);
atomic_fetch_add_explicit(&slot->seqlock, 1, memory_order_release);
}
Comment thread
Jjateen marked this conversation as resolved.

// Lock-free memory add using atomics with seqlock for consistent reads
int add_gpu_device_memory_usage(int32_t pid, int cudadev, size_t usage, int type) {
LOG_INFO("add_gpu_device_memory_lockfree:%d %d->%d %lu", pid, cudadev, cuda_to_nvml_map(cudadev), usage);
Expand All @@ -491,8 +594,8 @@ int add_gpu_device_memory_usage(int32_t pid, int cudadev, size_t usage, int type
if (self_slot != NULL) {
shrreg_proc_slot_t* slot = self_slot;

// Seqlock protocol: increment to odd (write in progress)
atomic_fetch_add_explicit(&slot->seqlock, 1, memory_order_release);
// Seqlock protocol: take the slot for writing (even -> odd)
seqlock_write_begin(slot);

// Perform updates with release semantics for visibility
atomic_fetch_add_explicit(&slot->used[dev].total, usage, memory_order_release);
Expand All @@ -508,8 +611,8 @@ int add_gpu_device_memory_usage(int32_t pid, int cudadev, size_t usage, int type
break;
}

// Seqlock protocol: increment to even (write complete)
atomic_fetch_add_explicit(&slot->seqlock, 1, memory_order_release);
// Seqlock protocol: release the slot (odd -> even)
seqlock_write_end(slot);

LOG_INFO("gpu_device_memory_added_lockfree:%d %d %lu", pid, dev, usage);
return 0;
Expand All @@ -523,8 +626,8 @@ int add_gpu_device_memory_usage(int32_t pid, int cudadev, size_t usage, int type
if (slot_pid == pid) {
shrreg_proc_slot_t* slot = &region_info.shared_region->procs[i];

// Seqlock protocol: increment to odd (write in progress)
atomic_fetch_add_explicit(&slot->seqlock, 1, memory_order_release);
// Seqlock protocol: take the slot for writing (even -> odd)
seqlock_write_begin(slot);

// Perform updates
atomic_fetch_add_explicit(&slot->used[dev].total, usage, memory_order_release);
Expand All @@ -540,8 +643,8 @@ int add_gpu_device_memory_usage(int32_t pid, int cudadev, size_t usage, int type
break;
}

// Seqlock protocol: increment to even (write complete)
atomic_fetch_add_explicit(&slot->seqlock, 1, memory_order_release);
// Seqlock protocol: release the slot (odd -> even)
seqlock_write_end(slot);

LOG_INFO("gpu_device_memory_added_lockfree:%d %d %lu", pid, dev, usage);
return 0;
Expand All @@ -566,8 +669,8 @@ int rm_gpu_device_memory_usage(int32_t pid, int cudadev, size_t usage, int type)
if (self_slot != NULL) {
shrreg_proc_slot_t* slot = self_slot;

// Seqlock protocol: increment to odd (write in progress)
atomic_fetch_add_explicit(&slot->seqlock, 1, memory_order_release);
// Seqlock protocol: take the slot for writing (even -> odd)
seqlock_write_begin(slot);

// Perform updates with release semantics
atomic_fetch_sub_explicit(&slot->used[dev].total, usage, memory_order_release);
Expand All @@ -583,8 +686,8 @@ int rm_gpu_device_memory_usage(int32_t pid, int cudadev, size_t usage, int type)
break;
}

// Seqlock protocol: increment to even (write complete)
atomic_fetch_add_explicit(&slot->seqlock, 1, memory_order_release);
// Seqlock protocol: release the slot (odd -> even)
seqlock_write_end(slot);

uint64_t new_total = atomic_load_explicit(&slot->used[dev].total, memory_order_acquire);
LOG_INFO("after delete_lockfree:%lu", new_total);
Expand All @@ -599,8 +702,8 @@ int rm_gpu_device_memory_usage(int32_t pid, int cudadev, size_t usage, int type)
if (slot_pid == pid) {
shrreg_proc_slot_t* slot = &region_info.shared_region->procs[i];

// Seqlock protocol: increment to odd (write in progress)
atomic_fetch_add_explicit(&slot->seqlock, 1, memory_order_release);
// Seqlock protocol: take the slot for writing (even -> odd)
seqlock_write_begin(slot);

// Perform updates
atomic_fetch_sub_explicit(&slot->used[dev].total, usage, memory_order_release);
Expand All @@ -616,8 +719,8 @@ int rm_gpu_device_memory_usage(int32_t pid, int cudadev, size_t usage, int type)
break;
}

// Seqlock protocol: increment to even (write complete)
atomic_fetch_add_explicit(&slot->seqlock, 1, memory_order_release);
// Seqlock protocol: release the slot (odd -> even)
seqlock_write_end(slot);

uint64_t new_total = atomic_load_explicit(&slot->used[dev].total, memory_order_acquire);
LOG_INFO("after delete_lockfree:%lu", new_total);
Expand Down
8 changes: 7 additions & 1 deletion src/multiprocess/multiprocess_memory_limit.h
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,13 @@ typedef struct {
device_util_t device_util[CUDA_DEVICE_MAX_COUNT];
_Atomic int32_t status;
_Atomic uint64_t seqlock; // Sequence lock for consistent snapshots
uint64_t unused[2];
// Pid of the writer currently holding seqlock, 0 when no one holds it.
// Taken from the space reserved by unused[], so sizeof this struct and the
// shared region layout are unchanged: a process built without it leaves the
// field 0 and the reclaim path simply declines to act.
_Atomic int32_t seq_owner;
int32_t seq_owner_pad;
uint64_t unused[1];
} shrreg_proc_slot_t;

typedef char uuid[96];
Expand Down
40 changes: 33 additions & 7 deletions test/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,26 @@ foreach(TEST_SCRIPT ${TEST_SCRIPTS})
get_filename_component(TEST_TARGET_NAME ${RELATIVE_TEST_PATH} NAME_WE)
set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/${TEST_TARGET_DIR})

if (TEST_TARGET_NAME STREQUAL "test_postinit_owner_death")
if (TEST_TARGET_NAME STREQUAL "test_postinit_owner_death" OR
TEST_TARGET_NAME STREQUAL "test_seqlock_writer_exclusion" OR
TEST_TARGET_NAME STREQUAL "test_seqlock_multiprocess")
# Build this focused regression test with the production shared-region
# implementation. It does not invoke any CUDA/NVML entry point at
# runtime. Section garbage collection drops unrelated GPU-facing
# production functions.
add_executable(${TEST_TARGET_NAME}
${TEST_SCRIPT}
${CMAKE_CURRENT_SOURCE_DIR}/../src/multiprocess/multiprocess_memory_limit.c
${CMAKE_CURRENT_SOURCE_DIR}/../src/log_utils.c)
if (TEST_TARGET_NAME STREQUAL "test_seqlock_writer_exclusion" OR
TEST_TARGET_NAME STREQUAL "test_seqlock_multiprocess")
# includes the translation unit directly to observe the sequence
# counter, so it must not also be linked in
add_executable(${TEST_TARGET_NAME}
${TEST_SCRIPT}
${CMAKE_CURRENT_SOURCE_DIR}/../src/log_utils.c)
else()
add_executable(${TEST_TARGET_NAME}
${TEST_SCRIPT}
${CMAKE_CURRENT_SOURCE_DIR}/../src/multiprocess/multiprocess_memory_limit.c
${CMAKE_CURRENT_SOURCE_DIR}/../src/log_utils.c)
endif()
target_compile_definitions(${TEST_TARGET_NAME} PRIVATE
_GNU_SOURCE)
target_compile_options(${TEST_TARGET_NAME} PRIVATE
Expand All @@ -38,7 +49,9 @@ foreach(TEST_SCRIPT ${TEST_SCRIPTS})
endif()

list(APPEND TEST_TARGET_NAMES_LIST ${TEST_TARGET_NAME})
if (TEST_TARGET_NAME STREQUAL "test_postinit_owner_death")
if (TEST_TARGET_NAME STREQUAL "test_postinit_owner_death" OR
TEST_TARGET_NAME STREQUAL "test_seqlock_writer_exclusion" OR
TEST_TARGET_NAME STREQUAL "test_seqlock_multiprocess")
target_link_libraries(${TEST_TARGET_NAME} -lrt -lpthread)
elseif (TEST_TARGET_NAME STREQUAL "test_dlsym_rtld_next")
target_link_libraries(${TEST_TARGET_NAME} -ldl)
Expand All @@ -56,7 +69,20 @@ set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})

add_test(NAME postinit_owner_death
COMMAND test_postinit_owner_death)
set_tests_properties(postinit_owner_death PROPERTIES TIMEOUT 20)
set_tests_properties(postinit_owner_death PROPERTIES TIMEOUT 20
RESOURCE_LOCK shared_region)

add_test(NAME seqlock_writer_exclusion
COMMAND test_seqlock_writer_exclusion)
set_tests_properties(seqlock_writer_exclusion PROPERTIES TIMEOUT 60
SKIP_RETURN_CODE 77
RESOURCE_LOCK shared_region)

add_test(NAME seqlock_multiprocess
COMMAND test_seqlock_multiprocess)
set_tests_properties(seqlock_multiprocess PROPERTIES TIMEOUT 120
SKIP_RETURN_CODE 77
RESOURCE_LOCK shared_region)

add_test(NAME dlsym_rtld_next
COMMAND test_dlsym_rtld_next)
Expand Down
Loading
Loading