diff --git a/cudax/include/cuda/experimental/__places/exec/locality_domain.cuh b/cudax/include/cuda/experimental/__places/exec/locality_domain.cuh index 49ecb63f4213..2a695a8f0b0f 100644 --- a/cudax/include/cuda/experimental/__places/exec/locality_domain.cuh +++ b/cudax/include/cuda/experimental/__places/exec/locality_domain.cuh @@ -91,6 +91,8 @@ # pragma system_header #endif // no system header +#include +#include #include #include @@ -187,64 +189,6 @@ inline bool locality_domain_memory_disabled() return disabled; } -/** - * @brief Cache of per-(device, domain) localized memory pools. - * - * Localized stream-ordered allocation goes through memory pools - * (`cuMemPoolCreate` + `cuMemAllocFromPoolAsync`). Pools are created lazily - * and reused for the lifetime of the process. Thread-safe. - */ -class locality_domain_mem_pool_cache -{ -public: - static locality_domain_mem_pool_cache& instance() - { - static locality_domain_mem_pool_cache inst; - return inst; - } - - CUmemoryPool get(int dev_id, int domain_id) - { - ::std::lock_guard<::std::mutex> lock(mtx_); - auto key = ::std::make_pair(dev_id, domain_id); - auto it = pools_.find(key); - if (it != pools_.end()) - { - return it->second; - } - - CUmemPoolProps props = {}; - props.allocType = CU_MEM_ALLOCATION_TYPE_PINNED; - // Plain device memory when localization is disabled, or when the driver - // cannot answer the locality-domain query (whole-device degrade: the - // localized location type would be rejected). - if (locality_domain_memory_disabled() || locality_domain_native_raw_count(dev_id) <= 0) - { - props.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - props.location.id = dev_id; - } - else - { - props.location.type = CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN; - props.location.localized.deviceId = static_cast(dev_id); - props.location.localized.localityDomainId = static_cast(domain_id); - } - - CUmemoryPool pool = nullptr; - cuda_try(cuMemPoolCreate(&pool, &props)); - pools_[key] = pool; - return pool; - } - -private: - locality_domain_mem_pool_cache() = default; - locality_domain_mem_pool_cache(const locality_domain_mem_pool_cache&) = delete; - locality_domain_mem_pool_cache& operator=(const locality_domain_mem_pool_cache&) = delete; - - ::std::map<::std::pair, CUmemoryPool> pools_; - ::std::mutex mtx_; -}; - /** * @brief Cache of per-domain green contexts and stream pools. * @@ -596,9 +540,12 @@ inline unsigned int locality_domain_count(int dev_id) * * With the native backend, `mem_create` produces a VMM physical handle whose * backing store lives in the requested domain, and `allocate` hands out - * stream-ordered memory from a per-domain localized memory pool. With the - * fallback backend, both delegate to the plain device data place. Identity - * (device ordinal, domain ordinal) is preserved by both backends. + * stream-ordered memory from the driver's default memory pool for that + * domain's location (obtained through `cuda::__get_default_memory_pool`, + * which also owns that pool's release-threshold policy — this layer creates + * and owns no pool of its own). With the fallback backend, both delegate to + * the plain device data place. Identity (device ordinal, domain ordinal) is + * preserved by both backends. */ class locality_domain_data_place_impl : public data_place_interface { @@ -646,31 +593,43 @@ public: #if _CUDAX_PLACES_LOCALITY_DOMAIN_NATIVE /** - * @brief Create physical memory localized to this domain (VMM API). + * @brief This domain's memory location: the locality domain itself, or + * plain device memory when localization is disabled or the driver cannot + * answer the locality-domain query (whole-device degrade — the localized + * location type would be rejected). * - * Uses `CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN` so the backing store is - * placed in the requested domain. When `CUDASTF_DISABLE_LOCALIZED_MEMORY` is - * set, falls back to plain device memory. + * Shared by the VMM and stream-ordered allocation paths, so both agree on + * which memory this place refers to. */ - CUresult mem_create(CUmemGenericAllocationHandle* handle, size_t size) const override + [[nodiscard]] _CCCL_HOST_API CUmemLocation __pool_location() const noexcept { - CUmemAllocationProp prop = {}; - prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; - - // Plain device memory when localization is disabled, or when the driver - // cannot answer the locality-domain query (whole-device degrade). + CUmemLocation location = {}; if (locality_domain_memory_disabled() || locality_domain_native_raw_count(view_.devid) <= 0) { - prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - prop.location.id = view_.devid; + location.type = CU_MEM_LOCATION_TYPE_DEVICE; + location.id = view_.devid; } else { - prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN; - prop.location.localized.deviceId = static_cast(view_.devid); - prop.location.localized.localityDomainId = static_cast(view_.domain_id); + location.type = CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN; + location.localized.deviceId = static_cast(view_.devid); + location.localized.localityDomainId = static_cast(view_.domain_id); } + return location; + } + /** + * @brief Create physical memory localized to this domain (VMM API). + * + * Uses `CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN` so the backing store is + * placed in the requested domain. When `CUDASTF_DISABLE_LOCALIZED_MEMORY` is + * set, falls back to plain device memory. + */ + CUresult mem_create(CUmemGenericAllocationHandle* handle, size_t size) const override + { + CUmemAllocationProp prop = {}; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.location = __pool_location(); return cuMemCreate(handle, size, &prop, 0); } @@ -679,12 +638,33 @@ public: */ void* allocate(::std::ptrdiff_t size, cudaStream_t stream) const override { - // No cudaSetDevice here: unlike the cudaMallocAsync-based places (device, + // The driver keeps one default pool per (location, allocation type), so + // there is nothing to create, own or cache here: the same handle comes + // back on every call, shared with every other consumer of that location + // in the process. + // + // Which accessor depends on WHOSE pool it is. A locality-domain location + // is this workload's own: `cuda::__get_default_memory_pool` is the + // library-wide site for such pools and settles their release-threshold + // policy (retaining freed memory instead of returning it to the OS at + // every synchronization, which would re-back algorithm-scale scratch on + // every call). The whole-device degrade location is NOT ours — it is the + // process-global device default pool that every `cudaMallocAsync` user + // shares — so it is fetched raw, inheriting whatever policy the process + // already has: this place should not decide retention on other + // components' behalf, least of all on machines whose memory is not + // partitioned into domains at all. + // + // No cudaSetDevice either: unlike the cudaMallocAsync-based places (device, // green_ctx), which draw from the *current* device's default pool, the pool - // is passed explicitly and was created with props.location.id == devid, so - // placement does not depend on the current device. This also keeps - // allocate() symmetric with deallocate(), which never switched. - CUmemoryPool pool = locality_domain_mem_pool_cache::instance().get(view_.devid, view_.domain_id); + // is passed explicitly and belongs to this domain's location, so placement + // does not depend on the current device. This also keeps allocate() + // symmetric with deallocate(), which never switched. + const CUmemLocation location = __pool_location(); + const CUmemoryPool pool = + (location.type == CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN) + ? ::cuda::__get_default_memory_pool(location, ::CU_MEM_ALLOCATION_TYPE_PINNED) + : ::cuda::__driver::__getDefaultMemPool(location, ::CU_MEM_ALLOCATION_TYPE_PINNED); CUdeviceptr ptr = 0; cuda_try(cuMemAllocFromPoolAsync(&ptr, static_cast(size), pool, reinterpret_cast(stream))); diff --git a/cudax/test/places/data_place_alloc.cu b/cudax/test/places/data_place_alloc.cu index 27defb2da2de..63e9bcd00990 100644 --- a/cudax/test/places/data_place_alloc.cu +++ b/cudax/test/places/data_place_alloc.cu @@ -18,6 +18,7 @@ * outside of the task-based programming model. */ +#include #include #include @@ -182,6 +183,41 @@ void test_managed_allocation() printf(" Managed allocation test PASSED\n"); } +// A locality-domain data place must never change the release-threshold policy +// of the process-global device default pool. That pool is shared with every +// other `cudaMallocAsync` user in the process, so retention there is not this +// place's decision to make — most visibly on machines with no locality +// domains, where the place degrades to whole-device memory and would +// otherwise be configuring a pool it does not own. (The domain's own default +// pool is a different location and is configured by the library-wide +// accessor.) +void test_device_default_pool_policy_untouched() +{ + printf("Testing that a locality-domain place leaves the device default pool alone...\n"); + + cudaMemPool_t default_pool; + cuda_try(cudaDeviceGetDefaultMemPool(&default_pool, 0)); + cuuint64_t before = 0; + cuda_try(cudaMemPoolGetAttribute(default_pool, cudaMemPoolAttrReleaseThreshold, &before)); + + cudaStream_t stream; + cuda_try(cudaStreamCreate(&stream)); + + auto place = data_place::locality_domain(0, 0); + const size_t n = size_t{1} << 20; + void* ptr = place.allocate(static_cast<::std::ptrdiff_t>(n), stream); + cuda_try(cudaStreamSynchronize(stream)); + place.deallocate(ptr, n, stream); + cuda_try(cudaStreamSynchronize(stream)); + + cuuint64_t after = 0; + cuda_try(cudaMemPoolGetAttribute(default_pool, cudaMemPoolAttrReleaseThreshold, &after)); + EXPECT(before == after); + + cuda_try(cudaStreamDestroy(stream)); + printf(" Device default pool policy test PASSED\n"); +} + int main() { printf("=== Testing data_place direct allocation (no context) ===\n\n"); @@ -189,6 +225,7 @@ int main() test_host_allocation(); test_device_allocation(); test_managed_allocation(); + test_device_default_pool_policy_untouched(); printf("\n=== All tests PASSED ===\n"); return 0;