[AMDGPU] Trim default mempool on qd.reset()

quadrants/rhi/amdgpu/amdgpu_context.cpp

@@ -205,6 +205,16 @@ void AMDGPUContext::launch(void *func,
   }
 }
 
+void AMDGPUContext::trim_default_mem_pool() {
+  if (!supports_mem_pool_) {
+    return;
+  }
+  void *default_mem_pool = nullptr;
+  driver_.device_get_default_mem_pool(&default_mem_pool, 0 /*device ordinal*/);
+  // `min_bytes_to_keep = 0` releases every cached page back to the driver.
+  driver_.mem_pool_trim_to(default_mem_pool, 0u);
+}
+
 AMDGPUContext::~AMDGPUContext() {
   // Currently unreachable: singleton is heap-allocated via `new` in get_instance() and never deleted.
   for (auto *s : stream_pool_) {

quadrants/rhi/amdgpu/amdgpu_context.h

@@ -79,6 +79,14 @@ class AMDGPUContext {
     return supports_mem_pool_;
   }
 
+  // Force the default device memory pool to release every cached page back to the driver. Called by
+  // `LlvmRuntimeExecutor::finalize` (i.e. `qd.reset()`) to align actual driver-visible free VRAM with the user-facing
+  // contract that `qd.reset()` releases everything Quadrants allocated. Without this, up to the configured release
+  // threshold (128 MiB at construction) of freed pages stays cached in the pool and shows up as "used" to other
+  // processes on the same VF, materially raising the chance of multi-process `HSA_STATUS_ERROR_OUT_OF_RESOURCES`
+  // failures across pytest-xdist workers. No-op if the device does not advertise mempool support.
+  void trim_default_mem_pool();
+
   ~AMDGPUContext();
 
   class ContextGuard {

quadrants/rhi/amdgpu/amdgpu_driver_functions.inc.h

@@ -42,6 +42,7 @@ PER_AMDGPU_FUNCTION(mem_alloc_host, hipHostMalloc, void **, std::size_t, uint32)
 PER_AMDGPU_FUNCTION(mem_free_host, hipHostFree, void *);
 PER_AMDGPU_FUNCTION(device_get_default_mem_pool, hipDeviceGetDefaultMemPool, void **, int);
 PER_AMDGPU_FUNCTION(mem_pool_set_attribute, hipMemPoolSetAttribute, void *, uint32, void *);
+PER_AMDGPU_FUNCTION(mem_pool_trim_to, hipMemPoolTrimTo, void *, std::size_t);
 PER_AMDGPU_FUNCTION(mem_get_info, hipMemGetInfo, std::size_t *, std::size_t *);
 PER_AMDGPU_FUNCTION(mem_get_attribute, hipPointerGetAttribute, void *, uint32, void *);
 PER_AMDGPU_FUNCTION(mem_get_attributes, hipPointerGetAttributes, void *, void *);

quadrants/rhi/cuda/cuda_context.cpp

@@ -109,6 +109,15 @@ std::string CUDAContext::get_device_name() {
   return str;
 }
 
+void CUDAContext::trim_default_mem_pool() {
+  if (!supports_mem_pool_) {
+    return;
+  }
+  void *default_mem_pool = nullptr;
+  driver_.device_get_default_mem_pool(&default_mem_pool, device_);
+  driver_.mem_pool_trim_to(default_mem_pool, 0u);
+}
+
 int64_t CUDAContext::get_clock_rate_khz() const {
   int clock_rate_khz = 0;
   driver_.device_get_attribute(&clock_rate_khz, CU_DEVICE_ATTRIBUTE_CLOCK_RATE, device_);

quadrants/rhi/cuda/cuda_context.h

@@ -89,6 +89,11 @@ class CUDAContext {
     return supports_mem_pool_;
   }
 
+  // Force the default device memory pool to release every cached page back to the driver. Symmetric with the AMDGPU
+  // side and called from `LlvmRuntimeExecutor::finalize` (i.e. `qd.reset()`); see
+  // `AMDGPUContext::trim_default_mem_pool` for the rationale. No-op when mempool support is unavailable.
+  void trim_default_mem_pool();
+
   // True when the device can coherently dereference plain host pointers (`malloc` / `new`) from kernel code via HMM /
   // system-allocated memory. Maps `CU_DEVICE_ATTRIBUTE_PAGEABLE_MEMORY_ACCESS` directly - 1 on Linux with an
   // HMM-capable driver + kernel (open-source nvidia module or 535+ with HMM enabled), 0 on Turing and older parts,

quadrants/rhi/cuda/cuda_driver_functions.inc.h

@@ -41,6 +41,7 @@ PER_CUDA_FUNCTION(mem_advise, cuMemAdvise, void *, std::size_t, uint32, uint32);
 PER_CUDA_FUNCTION(mem_get_info, cuMemGetInfo_v2, std::size_t *, std::size_t *);
 PER_CUDA_FUNCTION(mem_get_attribute, cuPointerGetAttribute, void *, uint32, void *);
 PER_CUDA_FUNCTION(mem_pool_set_attribute, cuMemPoolSetAttribute, void *, uint32, void *);
+PER_CUDA_FUNCTION(mem_pool_trim_to, cuMemPoolTrimTo, void *, std::size_t);
 
 // Module and kernels
 PER_CUDA_FUNCTION(module_get_function, cuModuleGetFunction, void **, void *, const char *);

quadrants/runtime/llvm/llvm_runtime_executor.cpp

@@ -612,6 +612,23 @@ void LlvmRuntimeExecutor::finalize() {
 
     // Release unused memory from cuda memory pool
     synchronize();
+
+    // Trim the driver's default async-alloc mempool back to the driver. The mempool's release-threshold (set to 128
+    // MiB at context construction) would otherwise keep up to that many bytes of freed pages cached, which shows up to
+    // other co-tenant processes on the same device as memory still in use. Aligning post-`qd.reset()` driver-visible
+    // free VRAM with the user-facing contract that `qd.reset()` releases everything Quadrants allocated materially
+    // reduces multi-process `HSA_STATUS_ERROR_OUT_OF_RESOURCES` failures across `pytest-xdist` workers (see
+    // perso_hugh/doc/amdgpu_test_suite_oom_followup.md). No-op on devices without mempool support.
+#if defined(QD_WITH_CUDA)
+    if (config_.arch == Arch::cuda) {
+      CUDAContext::get_instance().trim_default_mem_pool();
+    }
+#endif
+#if defined(QD_WITH_AMDGPU)
+    if (config_.arch == Arch::amdgpu) {
+      AMDGPUContext::get_instance().trim_default_mem_pool();
+    }
+#endif
   }
   finalized_ = true;
 }