pinned_pool.cuh

Go to the documentation of this file.

#pragma once
 
#include <cuda_runtime.h>
#include <cuda/devices>
#include <cuda/memory_pool>
 
#include <array>
#include <cstdint>
#include <limits>
#include <mutex>
#include <vector>
 
#include <cuddl/cuda_error.hpp>
#include <cuddl/detail/ablation.hpp>
 
namespace cuddl::detail {
 
enum class pinned_placement : uint8_t {
    device_node,
    host,
};
 
[[nodiscard]] inline cuda::pinned_memory_pool_ref
device_pinned_pool(cuda::device_ref device, pinned_placement placement) {
    static std::mutex mutex;
    static std::vector<std::array<cudaMemPool_t, 2>> pools(cuda::devices.size());
    std::lock_guard lock(mutex);
    auto& pool = pools.at(static_cast<size_t>(device.get()))[static_cast<size_t>(placement)];
    if (pool == nullptr) {
        // Ablated, the pool behaves like CCCL's default pinned pool: mapped into every visible
        // device and returning freed pages to the OS at every synchronisation. numa-pool
        // separately puts device_node pages on any node.
        bool const ablate = ablated("pinned-pool");
        int node = -1;
        CUDDL_CUDA_CALL(cudaDeviceGetAttribute(&node, cudaDevAttrHostNumaId, device.get()));
        bool const numa =
            !ablated("numa-pool") && placement == pinned_placement::device_node && node >= 0 &&
            device.attribute(cuda::device_attributes::host_numa_memory_pools_supported) != 0;
        cudaMemPoolProps properties{};
        properties.allocType = cudaMemAllocationTypePinned;
        properties.location = numa ? cudaMemLocation{cudaMemLocationTypeHostNuma, node}
                                   : cudaMemLocation{cudaMemLocationTypeHost, 0};
        cudaMemPool_t created = nullptr;
        CUDDL_CUDA_CALL(cudaMemPoolCreate(&created, &properties));
        cudaMemAccessDesc const access{
            {cudaMemLocationTypeDevice, device.get()}, cudaMemAccessFlagsProtReadWrite
        };
        CUDDL_CUDA_CALL(cudaMemPoolSetAccess(created, &access, 1));
        if (ablate) {
            for (auto const& other : cuda::devices) {
                cudaMemAccessDesc const shared{
                    {cudaMemLocationTypeDevice, other.get()}, cudaMemAccessFlagsProtReadWrite
                };
                CUDDL_CUDA_CALL(cudaMemPoolSetAccess(created, &shared, 1));
            }
        }
        // Pinning is the expensive part of an allocation: with the default threshold of zero,
        // every synchronisation returns freed pages to the OS and the next allocation pins them
        // again, which cost sketch loaders more time than inflating their files.
        auto keep = ablate ? uint64_t{0} : std::numeric_limits<uint64_t>::max();
        CUDDL_CUDA_CALL(cudaMemPoolSetAttribute(created, cudaMemPoolAttrReleaseThreshold, &keep));
        pool = created;
    }
    return cuda::pinned_memory_pool_ref(pool);
}
 
class pinned_pool_trim {
   public:
    pinned_pool_trim(cuda::device_ref device, pinned_placement placement)
        : pool_(device_pinned_pool(device, placement)) {}
    pinned_pool_trim(pinned_pool_trim const&) = delete;
    pinned_pool_trim& operator=(pinned_pool_trim const&) = delete;
    ~pinned_pool_trim() {
        (void)cudaMemPoolTrimTo(pool_.get(), 0);
    }
 
   private:
    cuda::pinned_memory_pool_ref pool_;
};
 
}  // namespace cuddl::detail