diff --git a/cpp/src/arrow/buffer_test.cc b/cpp/src/arrow/buffer_test.cc index 9b71aea1addcd..06ed0bfba0497 100644 --- a/cpp/src/arrow/buffer_test.cc +++ b/cpp/src/arrow/buffer_test.cc @@ -1027,13 +1027,13 @@ TEST(TestDeviceRegistry, Basics) { // Test the error cases for the device registry // CPU is already registered - ASSERT_RAISES(KeyError, RegisterDeviceMemoryManager( - DeviceAllocationType::kCPU, [](int64_t device_id) { - return default_cpu_memory_manager(); - })); + ASSERT_RAISES(KeyError, + RegisterDeviceMapper(DeviceAllocationType::kCPU, [](int64_t device_id) { + return default_cpu_memory_manager(); + })); // VPI is not registered - ASSERT_RAISES(KeyError, GetDeviceMemoryManager(DeviceAllocationType::kVPI)); + ASSERT_RAISES(KeyError, GetDeviceMapper(DeviceAllocationType::kVPI)); } } // namespace arrow diff --git a/cpp/src/arrow/c/bridge.cc b/cpp/src/arrow/c/bridge.cc index e825f34e712ec..d004de7a2ea9f 100644 --- a/cpp/src/arrow/c/bridge.cc +++ b/cpp/src/arrow/c/bridge.cc @@ -1967,10 +1967,10 @@ Result> ImportRecordBatch(struct ArrowArray* array, return ImportRecordBatch(array, *maybe_schema); } -Result> DefaultDeviceMapper(ArrowDeviceType device_type, - int64_t device_id) { - ARROW_ASSIGN_OR_RAISE(auto mapper, GetDeviceMemoryManager( - static_cast(device_type))); +Result> DefaultDeviceMemoryMapper( + ArrowDeviceType device_type, int64_t device_id) { + ARROW_ASSIGN_OR_RAISE(auto mapper, + GetDeviceMapper(static_cast(device_type))); return mapper(device_id); } diff --git a/cpp/src/arrow/c/bridge.h b/cpp/src/arrow/c/bridge.h index 0ced3d38cd1e6..74a302be4c27d 100644 --- a/cpp/src/arrow/c/bridge.h +++ b/cpp/src/arrow/c/bridge.h @@ -219,8 +219,8 @@ using DeviceMemoryMapper = std::function>(ArrowDeviceType, int64_t)>; ARROW_EXPORT -Result> DefaultDeviceMapper(ArrowDeviceType device_type, - int64_t device_id); +Result> DefaultDeviceMemoryMapper( + ArrowDeviceType device_type, int64_t device_id); /// \brief EXPERIMENTAL: Import C++ device array from the C data interface. /// @@ -236,7 +236,7 @@ Result> DefaultDeviceMapper(ArrowDeviceType devic ARROW_EXPORT Result> ImportDeviceArray( struct ArrowDeviceArray* array, std::shared_ptr type, - const DeviceMemoryMapper& mapper = DefaultDeviceMapper); + const DeviceMemoryMapper& mapper = DefaultDeviceMemoryMapper); /// \brief EXPERIMENTAL: Import C++ device array and its type from the C data interface. /// @@ -253,7 +253,7 @@ Result> ImportDeviceArray( ARROW_EXPORT Result> ImportDeviceArray( struct ArrowDeviceArray* array, struct ArrowSchema* type, - const DeviceMemoryMapper& mapper = DefaultDeviceMapper); + const DeviceMemoryMapper& mapper = DefaultDeviceMemoryMapper); /// \brief EXPERIMENTAL: Import C++ record batch with buffers on a device from the C data /// interface. @@ -271,7 +271,7 @@ Result> ImportDeviceArray( ARROW_EXPORT Result> ImportDeviceRecordBatch( struct ArrowDeviceArray* array, std::shared_ptr schema, - const DeviceMemoryMapper& mapper = DefaultDeviceMapper); + const DeviceMemoryMapper& mapper = DefaultDeviceMemoryMapper); /// \brief EXPERIMENTAL: Import C++ record batch with buffers on a device and its schema /// from the C data interface. @@ -291,7 +291,7 @@ Result> ImportDeviceRecordBatch( ARROW_EXPORT Result> ImportDeviceRecordBatch( struct ArrowDeviceArray* array, struct ArrowSchema* schema, - const DeviceMemoryMapper& mapper = DefaultDeviceMapper); + const DeviceMemoryMapper& mapper = DefaultDeviceMemoryMapper); /// @} diff --git a/cpp/src/arrow/device.cc b/cpp/src/arrow/device.cc index 0e8420099a6e5..172136df2c9eb 100644 --- a/cpp/src/arrow/device.cc +++ b/cpp/src/arrow/device.cc @@ -272,11 +272,11 @@ std::shared_ptr CPUDevice::default_memory_manager() { namespace { -class DeviceMemoryManagerRegistryImpl { +class DeviceMapperRegistryImpl { public: - DeviceMemoryManagerRegistryImpl() {} + DeviceMapperRegistryImpl() {} - Status RegisterDevice(DeviceAllocationType device_type, MemoryMapper memory_mapper) { + Status RegisterDevice(DeviceAllocationType device_type, DeviceMapper memory_mapper) { std::lock_guard lock(lock_); auto [_, inserted] = registry_.try_emplace(device_type, std::move(memory_mapper)); if (!inserted) { @@ -286,7 +286,7 @@ class DeviceMemoryManagerRegistryImpl { return Status::OK(); } - Result GetMemoryManager(DeviceAllocationType device_type) { + Result GetMapper(DeviceAllocationType device_type) { std::lock_guard lock(lock_); auto it = registry_.find(device_type); if (it == registry_.end()) { @@ -298,38 +298,38 @@ class DeviceMemoryManagerRegistryImpl { private: std::mutex lock_; - std::unordered_map registry_; + std::unordered_map registry_; }; -Result> DefaultCPUMemoryMapper(int64_t device_id) { +Result> DefaultCPUDeviceMapper(int64_t device_id) { return default_cpu_memory_manager(); } -static std::unique_ptr CreateDeviceRegistry() { - auto registry = std::make_unique(); +static std::unique_ptr CreateDeviceRegistry() { + auto registry = std::make_unique(); // Always register the CPU device - DCHECK_OK(registry->RegisterDevice(DeviceAllocationType::kCPU, DefaultCPUMemoryMapper)); + DCHECK_OK(registry->RegisterDevice(DeviceAllocationType::kCPU, DefaultCPUDeviceMapper)); return registry; } -DeviceMemoryManagerRegistryImpl* GetDeviceRegistry() { +DeviceMapperRegistryImpl* GetDeviceRegistry() { static auto g_registry = CreateDeviceRegistry(); return g_registry.get(); } } // namespace -Status RegisterDeviceMemoryManager(DeviceAllocationType device_type, - MemoryMapper memory_mapper) { +Status RegisterDeviceMapper(DeviceAllocationType device_type, + DeviceMapper mapper) { auto registry = GetDeviceRegistry(); - return registry->RegisterDevice(device_type, std::move(memory_mapper)); + return registry->RegisterDevice(device_type, std::move(mapper)); } -Result GetDeviceMemoryManager(DeviceAllocationType device_type) { +Result GetDeviceMapper(DeviceAllocationType device_type) { auto registry = GetDeviceRegistry(); - return registry->GetMemoryManager(device_type); + return registry->GetMapper(device_type); } } // namespace arrow diff --git a/cpp/src/arrow/device.h b/cpp/src/arrow/device.h index f15e078c2afc9..b4d19c95129ac 100644 --- a/cpp/src/arrow/device.h +++ b/cpp/src/arrow/device.h @@ -363,7 +363,7 @@ class ARROW_EXPORT CPUMemoryManager : public MemoryManager { ARROW_EXPORT std::shared_ptr default_cpu_memory_manager(); -using MemoryMapper = +using DeviceMapper = std::function>(int64_t device_id)>; /// \brief Register a function to retrieve a MemoryManager for a Device type @@ -380,8 +380,7 @@ using MemoryMapper = /// MemoryManager for the registered device type and given device id /// \return Status ARROW_EXPORT -Status RegisterDeviceMemoryManager(DeviceAllocationType device_type, - MemoryMapper memory_mapper); +Status RegisterDeviceMapper(DeviceAllocationType device_type, DeviceMapper mapper); /// \brief Get the registered function to retrieve a MemoryManager for the /// given Device type @@ -390,6 +389,6 @@ Status RegisterDeviceMemoryManager(DeviceAllocationType device_type, /// \return function that takes a device id and returns the appropriate /// MemoryManager for the registered device type and given device id ARROW_EXPORT -Result GetDeviceMemoryManager(DeviceAllocationType device_type); +Result GetDeviceMapper(DeviceAllocationType device_type); } // namespace arrow diff --git a/cpp/src/arrow/gpu/cuda_memory.cc b/cpp/src/arrow/gpu/cuda_memory.cc index 7da7b2f50500b..6972321006a9a 100644 --- a/cpp/src/arrow/gpu/cuda_memory.cc +++ b/cpp/src/arrow/gpu/cuda_memory.cc @@ -504,14 +504,13 @@ Result> DefaultMemoryMapper(ArrowDeviceType devic namespace { -Result> DefaultCUDAMemoryMapper(int64_t device_id) { +Result> DefaultCUDADeviceMapper(int64_t device_id) { ARROW_ASSIGN_OR_RAISE(auto device, arrow::cuda::CudaDevice::Make(device_id)); return device->default_memory_manager(); } bool RegisterCUDADeviceInternal() { - DCHECK_OK( - RegisterDeviceMemoryManager(DeviceAllocationType::kCUDA, DefaultCUDAMemoryMapper)); + DCHECK_OK(RegisterDeviceMapper(DeviceAllocationType::kCUDA, DefaultCUDADeviceMapper)); // TODO add the CUDA_HOST and CUDA_MANAGED allocation types when they are supported in // the CudaDevice return true;