From b4a60e63de57be9d0eddf050dd5179648795d070 Mon Sep 17 00:00:00 2001 From: Felipe Oliveira Carvalho Date: Thu, 1 Aug 2024 14:23:51 -0300 Subject: [PATCH 1/7] cpp: Add the DeviceAllocationTypeSet class And move `enum DeviceAllocationType` to its own lightweight header. --- cpp/src/arrow/chunked_array.cc | 12 ++++ cpp/src/arrow/chunked_array.h | 86 ++++++++++++++++++++++++++ cpp/src/arrow/datum.h | 3 + cpp/src/arrow/device.h | 19 +----- cpp/src/arrow/device_allocation_type.h | 75 ++++++++++++++++++++++ 5 files changed, 177 insertions(+), 18 deletions(-) create mode 100644 cpp/src/arrow/device_allocation_type.h diff --git a/cpp/src/arrow/chunked_array.cc b/cpp/src/arrow/chunked_array.cc index c36b736d5d5..7c72ea04c3b 100644 --- a/cpp/src/arrow/chunked_array.cc +++ b/cpp/src/arrow/chunked_array.cc @@ -86,6 +86,18 @@ Result> ChunkedArray::MakeEmpty( return std::make_shared(std::move(new_chunks)); } +DeviceAllocationTypeSet ChunkedArray::DeviceTypeSet() const { + if (chunks_.empty()) { + // An empty ChunkedArray is considered to be CPU-only. + return DeviceAllocationTypeSet::CpuOnly(); + } + DeviceAllocationTypeSet set; + for (const auto& chunk : chunks_) { + set.add(chunk->device_type()); + } + return set; +} + bool ChunkedArray::Equals(const ChunkedArray& other, const EqualOptions& opts) const { if (length_ != other.length()) { return false; diff --git a/cpp/src/arrow/chunked_array.h b/cpp/src/arrow/chunked_array.h index 5d300861d85..083da3bb46a 100644 --- a/cpp/src/arrow/chunked_array.h +++ b/cpp/src/arrow/chunked_array.h @@ -17,6 +17,7 @@ #pragma once +#include #include #include #include @@ -25,6 +26,7 @@ #include "arrow/chunk_resolver.h" #include "arrow/compare.h" +#include "arrow/device_allocation_type.h" #include "arrow/result.h" #include "arrow/status.h" #include "arrow/type_fwd.h" @@ -41,6 +43,86 @@ template class ChunkedArrayIterator; } // namespace stl +class ARROW_EXPORT DeviceAllocationTypeSet { + private: + std::bitset device_type_bitset_; + + public: + /// \brief Construct an empty set of device types. + DeviceAllocationTypeSet() = default; + + /// \brief Construct a set of device types with a single device type. + DeviceAllocationTypeSet( // NOLINT implicit construction + DeviceAllocationType accepted_device_type) { + add(accepted_device_type); + } + + /// \brief Construct a set of device types containing only "kCPU". + static DeviceAllocationTypeSet CpuOnly() { + return DeviceAllocationTypeSet{DeviceAllocationType::kCPU}; + } + + /// \brief Construct a set of device types containing all device types. + static DeviceAllocationTypeSet All() { + DeviceAllocationTypeSet all; + all.device_type_bitset_.set(); + // Don't set the invalid enum values. + all.device_type_bitset_.reset(0); + all.device_type_bitset_.reset(5); + all.device_type_bitset_.reset(6); + return all; + } + + /// \brief Add a device type to the set of device types. + void add(DeviceAllocationType device_type) { + device_type_bitset_.set(static_cast(device_type)); + } + + /// \brief Remove a device type from the set of device types. + void remove(DeviceAllocationType device_type) { + device_type_bitset_.reset(static_cast(device_type)); + } + + /// \brief Return true if the set of accepted device types includes the + /// device type. + bool contains(DeviceAllocationType device_type) const { + return device_type_bitset_.test(static_cast(device_type)); + } + + /// \brief Add all device types from another set to this set. + void Add(DeviceAllocationTypeSet other) { + device_type_bitset_ |= other.device_type_bitset_; + } + + /// \brief Return true if the set of accepted device types includes all the + /// device types in the other set. + bool Contains(DeviceAllocationTypeSet other) const { + // other \subseteq this <==> (other \intersect this == other) + return (other.device_type_bitset_ & device_type_bitset_) == other.device_type_bitset_; + } + + std::string ToString() const { + std::string result = "{"; + for (int i = 1; i <= kDeviceAllocationTypeMax; i++) { + if (device_type_bitset_.test(i)) { + // Skip all the unused values in the enum. + switch (i) { + case 0: + case 5: + case 6: + continue; + } + if (result.size() > 1) { + result += ", "; + } + result += DeviceAllocationTypeToCStr(static_cast(i)); + } + } + result += "}"; + return result; + } +}; + /// \class ChunkedArray /// \brief A data structure managing a list of primitive Arrow arrays logically /// as one large array @@ -116,6 +198,10 @@ class ARROW_EXPORT ChunkedArray { /// \return an ArrayVector of chunks const ArrayVector& chunks() const { return chunks_; } + /// \return The set of device allocation types used by the chunks in this + /// chunked array. + DeviceAllocationTypeSet DeviceTypeSet() const; + /// \brief Construct a zero-copy slice of the chunked array with the /// indicated offset and length /// diff --git a/cpp/src/arrow/datum.h b/cpp/src/arrow/datum.h index 31b2d2274c9..bd596b2af83 100644 --- a/cpp/src/arrow/datum.h +++ b/cpp/src/arrow/datum.h @@ -26,6 +26,7 @@ #include #include "arrow/array/data.h" +#include "arrow/chunked_array.h" // for DeviceAllocationTypeSet #include "arrow/scalar.h" #include "arrow/type.h" #include "arrow/type_traits.h" @@ -295,6 +296,8 @@ struct ARROW_EXPORT Datum { /// \return empty if not arraylike ArrayVector chunks() const; + DeviceAllocationTypeSet DeviceTypeSet() const; + /// \brief True if the two data are equal bool Equals(const Datum& other) const; diff --git a/cpp/src/arrow/device.h b/cpp/src/arrow/device.h index f5cca0d27d7..00a365c307f 100644 --- a/cpp/src/arrow/device.h +++ b/cpp/src/arrow/device.h @@ -22,6 +22,7 @@ #include #include +#include "arrow/device_allocation_type.h" #include "arrow/io/type_fwd.h" #include "arrow/result.h" #include "arrow/status.h" @@ -32,24 +33,6 @@ namespace arrow { -/// \brief EXPERIMENTAL: Device type enum which matches up with C Data Device types -enum class DeviceAllocationType : char { - kCPU = 1, - kCUDA = 2, - kCUDA_HOST = 3, - kOPENCL = 4, - kVULKAN = 7, - kMETAL = 8, - kVPI = 9, - kROCM = 10, - kROCM_HOST = 11, - kEXT_DEV = 12, - kCUDA_MANAGED = 13, - kONEAPI = 14, - kWEBGPU = 15, - kHEXAGON = 16, -}; - class MemoryManager; /// \brief EXPERIMENTAL: Abstract interface for hardware devices diff --git a/cpp/src/arrow/device_allocation_type.h b/cpp/src/arrow/device_allocation_type.h new file mode 100644 index 00000000000..caf74efb3d6 --- /dev/null +++ b/cpp/src/arrow/device_allocation_type.h @@ -0,0 +1,75 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +namespace arrow { + +/// \brief EXPERIMENTAL: Device type enum which matches up with C Data Device types +enum class DeviceAllocationType : char { + kCPU = 1, + kCUDA = 2, + kCUDA_HOST = 3, + kOPENCL = 4, + kVULKAN = 7, + kMETAL = 8, + kVPI = 9, + kROCM = 10, + kROCM_HOST = 11, + kEXT_DEV = 12, + kCUDA_MANAGED = 13, + kONEAPI = 14, + kWEBGPU = 15, + kHEXAGON = 16, +}; +constexpr int kDeviceAllocationTypeMax = 16; + +inline const char* DeviceAllocationTypeToCStr(DeviceAllocationType type) { + switch (type) { + case DeviceAllocationType::kCPU: + return "CPU"; + case DeviceAllocationType::kCUDA: + return "CUDA"; + case DeviceAllocationType::kCUDA_HOST: + return "CUDA_HOST"; + case DeviceAllocationType::kOPENCL: + return "OPENCL"; + case DeviceAllocationType::kVULKAN: + return "VULKAN"; + case DeviceAllocationType::kMETAL: + return "METAL"; + case DeviceAllocationType::kVPI: + return "VPI"; + case DeviceAllocationType::kROCM: + return "ROCM"; + case DeviceAllocationType::kROCM_HOST: + return "ROCM_HOST"; + case DeviceAllocationType::kEXT_DEV: + return "EXT_DEV"; + case DeviceAllocationType::kCUDA_MANAGED: + return "CUDA_MANAGED"; + case DeviceAllocationType::kONEAPI: + return "ONEAPI"; + case DeviceAllocationType::kWEBGPU: + return "WEBGPU"; + case DeviceAllocationType::kHEXAGON: + return "HEXAGON"; + } + return ""; +} + +} // namespace arrow From 3480c1bde2014c681d636428baee7b80be8bd12b Mon Sep 17 00:00:00 2001 From: Felipe Oliveira Carvalho Date: Mon, 26 Aug 2024 21:29:36 -0300 Subject: [PATCH 2/7] cpp: Put DeviceAllocationTypeSet in the cheapest possible header --- cpp/src/arrow/CMakeLists.txt | 1 + cpp/src/arrow/array/array_base.h | 1 + cpp/src/arrow/array/array_test.cc | 1 + cpp/src/arrow/array/data.cc | 1 + cpp/src/arrow/buffer.h | 1 + cpp/src/arrow/buffer_test.cc | 1 + cpp/src/arrow/c/bridge.h | 1 + cpp/src/arrow/chunked_array.cc | 1 + cpp/src/arrow/chunked_array.h | 82 +----------------- cpp/src/arrow/compute/function.cc | 1 + cpp/src/arrow/compute/kernel.cc | 1 + cpp/src/arrow/compute/kernel.h | 1 + cpp/src/arrow/datum.cc | 1 + cpp/src/arrow/datum.h | 2 +- cpp/src/arrow/device.cc | 1 + cpp/src/arrow/device_allocation_type.h | 34 -------- cpp/src/arrow/device_allocation_type_set.cc | 80 ++++++++++++++++++ cpp/src/arrow/device_allocation_type_set.h | 92 +++++++++++++++++++++ cpp/src/arrow/gpu/cuda_context.cc | 1 + cpp/src/arrow/pretty_print.cc | 1 + cpp/src/arrow/record_batch.cc | 1 + cpp/src/arrow/record_batch.h | 1 + 22 files changed, 191 insertions(+), 116 deletions(-) create mode 100644 cpp/src/arrow/device_allocation_type_set.cc create mode 100644 cpp/src/arrow/device_allocation_type_set.h diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index 6b0ac8c23c7..65343df1291 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/cpp/src/arrow/CMakeLists.txt @@ -373,6 +373,7 @@ set(ARROW_SRCS config.cc datum.cc device.cc + device_allocation_type_set.cc extension_type.cc extension/bool8.cc extension/uuid.cc diff --git a/cpp/src/arrow/array/array_base.h b/cpp/src/arrow/array/array_base.h index 716ae072206..ad8f84a4e17 100644 --- a/cpp/src/arrow/array/array_base.h +++ b/cpp/src/arrow/array/array_base.h @@ -26,6 +26,7 @@ #include "arrow/array/data.h" #include "arrow/buffer.h" #include "arrow/compare.h" +#include "arrow/device_allocation_type.h" #include "arrow/result.h" #include "arrow/status.h" #include "arrow/type.h" diff --git a/cpp/src/arrow/array/array_test.cc b/cpp/src/arrow/array/array_test.cc index 32806d9d2ed..cdc8be813d2 100644 --- a/cpp/src/arrow/array/array_test.cc +++ b/cpp/src/arrow/array/array_test.cc @@ -45,6 +45,7 @@ #include "arrow/buffer.h" #include "arrow/buffer_builder.h" #include "arrow/compare.h" +#include "arrow/device_allocation_type.h" #include "arrow/result.h" #include "arrow/scalar.h" #include "arrow/status.h" diff --git a/cpp/src/arrow/array/data.cc b/cpp/src/arrow/array/data.cc index 83eeb56c496..bf611a63141 100644 --- a/cpp/src/arrow/array/data.cc +++ b/cpp/src/arrow/array/data.cc @@ -28,6 +28,7 @@ #include "arrow/array/util.h" #include "arrow/buffer.h" #include "arrow/device.h" +#include "arrow/device_allocation_type.h" #include "arrow/scalar.h" #include "arrow/status.h" #include "arrow/type.h" diff --git a/cpp/src/arrow/buffer.h b/cpp/src/arrow/buffer.h index fbf4a22e350..adef81f4062 100644 --- a/cpp/src/arrow/buffer.h +++ b/cpp/src/arrow/buffer.h @@ -27,6 +27,7 @@ #include #include "arrow/device.h" +#include "arrow/device_allocation_type.h" #include "arrow/status.h" #include "arrow/type_fwd.h" #include "arrow/util/macros.h" diff --git a/cpp/src/arrow/buffer_test.cc b/cpp/src/arrow/buffer_test.cc index 06ed0bfba04..2e5b2ac6ea5 100644 --- a/cpp/src/arrow/buffer_test.cc +++ b/cpp/src/arrow/buffer_test.cc @@ -29,6 +29,7 @@ #include "arrow/buffer.h" #include "arrow/buffer_builder.h" #include "arrow/device.h" +#include "arrow/device_allocation_type.h" #include "arrow/io/interfaces.h" #include "arrow/memory_pool.h" #include "arrow/status.h" diff --git a/cpp/src/arrow/c/bridge.h b/cpp/src/arrow/c/bridge.h index 45367e4f930..59168eda973 100644 --- a/cpp/src/arrow/c/bridge.h +++ b/cpp/src/arrow/c/bridge.h @@ -23,6 +23,7 @@ #include "arrow/c/abi.h" #include "arrow/device.h" +#include "arrow/device_allocation_type.h" #include "arrow/result.h" #include "arrow/status.h" #include "arrow/type_fwd.h" diff --git a/cpp/src/arrow/chunked_array.cc b/cpp/src/arrow/chunked_array.cc index 7c72ea04c3b..a5e590087e2 100644 --- a/cpp/src/arrow/chunked_array.cc +++ b/cpp/src/arrow/chunked_array.cc @@ -27,6 +27,7 @@ #include "arrow/array/array_nested.h" #include "arrow/array/util.h" #include "arrow/array/validate.h" +#include "arrow/device_allocation_type_set.h" #include "arrow/pretty_print.h" #include "arrow/status.h" #include "arrow/type.h" diff --git a/cpp/src/arrow/chunked_array.h b/cpp/src/arrow/chunked_array.h index 083da3bb46a..f2d07dfed3a 100644 --- a/cpp/src/arrow/chunked_array.h +++ b/cpp/src/arrow/chunked_array.h @@ -26,7 +26,7 @@ #include "arrow/chunk_resolver.h" #include "arrow/compare.h" -#include "arrow/device_allocation_type.h" +#include "arrow/device_allocation_type_set.h" #include "arrow/result.h" #include "arrow/status.h" #include "arrow/type_fwd.h" @@ -43,86 +43,6 @@ template class ChunkedArrayIterator; } // namespace stl -class ARROW_EXPORT DeviceAllocationTypeSet { - private: - std::bitset device_type_bitset_; - - public: - /// \brief Construct an empty set of device types. - DeviceAllocationTypeSet() = default; - - /// \brief Construct a set of device types with a single device type. - DeviceAllocationTypeSet( // NOLINT implicit construction - DeviceAllocationType accepted_device_type) { - add(accepted_device_type); - } - - /// \brief Construct a set of device types containing only "kCPU". - static DeviceAllocationTypeSet CpuOnly() { - return DeviceAllocationTypeSet{DeviceAllocationType::kCPU}; - } - - /// \brief Construct a set of device types containing all device types. - static DeviceAllocationTypeSet All() { - DeviceAllocationTypeSet all; - all.device_type_bitset_.set(); - // Don't set the invalid enum values. - all.device_type_bitset_.reset(0); - all.device_type_bitset_.reset(5); - all.device_type_bitset_.reset(6); - return all; - } - - /// \brief Add a device type to the set of device types. - void add(DeviceAllocationType device_type) { - device_type_bitset_.set(static_cast(device_type)); - } - - /// \brief Remove a device type from the set of device types. - void remove(DeviceAllocationType device_type) { - device_type_bitset_.reset(static_cast(device_type)); - } - - /// \brief Return true if the set of accepted device types includes the - /// device type. - bool contains(DeviceAllocationType device_type) const { - return device_type_bitset_.test(static_cast(device_type)); - } - - /// \brief Add all device types from another set to this set. - void Add(DeviceAllocationTypeSet other) { - device_type_bitset_ |= other.device_type_bitset_; - } - - /// \brief Return true if the set of accepted device types includes all the - /// device types in the other set. - bool Contains(DeviceAllocationTypeSet other) const { - // other \subseteq this <==> (other \intersect this == other) - return (other.device_type_bitset_ & device_type_bitset_) == other.device_type_bitset_; - } - - std::string ToString() const { - std::string result = "{"; - for (int i = 1; i <= kDeviceAllocationTypeMax; i++) { - if (device_type_bitset_.test(i)) { - // Skip all the unused values in the enum. - switch (i) { - case 0: - case 5: - case 6: - continue; - } - if (result.size() > 1) { - result += ", "; - } - result += DeviceAllocationTypeToCStr(static_cast(i)); - } - } - result += "}"; - return result; - } -}; - /// \class ChunkedArray /// \brief A data structure managing a list of primitive Arrow arrays logically /// as one large array diff --git a/cpp/src/arrow/compute/function.cc b/cpp/src/arrow/compute/function.cc index e1a2e8c5d88..0478a3d1e80 100644 --- a/cpp/src/arrow/compute/function.cc +++ b/cpp/src/arrow/compute/function.cc @@ -30,6 +30,7 @@ #include "arrow/compute/kernels/common_internal.h" #include "arrow/compute/registry.h" #include "arrow/datum.h" +#include "arrow/device_allocation_type_set.h" #include "arrow/util/cpu_info.h" #include "arrow/util/logging.h" #include "arrow/util/tracing_internal.h" diff --git a/cpp/src/arrow/compute/kernel.cc b/cpp/src/arrow/compute/kernel.cc index 5c87ef2cd05..5e7461cc52d 100644 --- a/cpp/src/arrow/compute/kernel.cc +++ b/cpp/src/arrow/compute/kernel.cc @@ -24,6 +24,7 @@ #include "arrow/buffer.h" #include "arrow/compute/exec.h" +#include "arrow/device_allocation_type_set.h" #include "arrow/result.h" #include "arrow/type_traits.h" #include "arrow/util/bit_util.h" diff --git a/cpp/src/arrow/compute/kernel.h b/cpp/src/arrow/compute/kernel.h index 1adb3e96c97..cfa1cd8193f 100644 --- a/cpp/src/arrow/compute/kernel.h +++ b/cpp/src/arrow/compute/kernel.h @@ -31,6 +31,7 @@ #include "arrow/buffer.h" #include "arrow/compute/exec.h" #include "arrow/datum.h" +#include "arrow/device_allocation_type_set.h" #include "arrow/memory_pool.h" #include "arrow/result.h" #include "arrow/status.h" diff --git a/cpp/src/arrow/datum.cc b/cpp/src/arrow/datum.cc index 2ac230232e1..cd9da31982f 100644 --- a/cpp/src/arrow/datum.cc +++ b/cpp/src/arrow/datum.cc @@ -25,6 +25,7 @@ #include "arrow/array/array_base.h" #include "arrow/array/util.h" #include "arrow/chunked_array.h" +#include "arrow/device_allocation_type_set.h" #include "arrow/record_batch.h" #include "arrow/scalar.h" #include "arrow/table.h" diff --git a/cpp/src/arrow/datum.h b/cpp/src/arrow/datum.h index bd596b2af83..5dc8f9ea41b 100644 --- a/cpp/src/arrow/datum.h +++ b/cpp/src/arrow/datum.h @@ -26,7 +26,7 @@ #include #include "arrow/array/data.h" -#include "arrow/chunked_array.h" // for DeviceAllocationTypeSet +#include "arrow/device_allocation_type_set.h" #include "arrow/scalar.h" #include "arrow/type.h" #include "arrow/type_traits.h" diff --git a/cpp/src/arrow/device.cc b/cpp/src/arrow/device.cc index c2195e89e75..d08271140f0 100644 --- a/cpp/src/arrow/device.cc +++ b/cpp/src/arrow/device.cc @@ -24,6 +24,7 @@ #include "arrow/array.h" #include "arrow/buffer.h" +#include "arrow/device_allocation_type.h" #include "arrow/io/memory.h" #include "arrow/record_batch.h" #include "arrow/result.h" diff --git a/cpp/src/arrow/device_allocation_type.h b/cpp/src/arrow/device_allocation_type.h index caf74efb3d6..a8a99867bdd 100644 --- a/cpp/src/arrow/device_allocation_type.h +++ b/cpp/src/arrow/device_allocation_type.h @@ -38,38 +38,4 @@ enum class DeviceAllocationType : char { }; constexpr int kDeviceAllocationTypeMax = 16; -inline const char* DeviceAllocationTypeToCStr(DeviceAllocationType type) { - switch (type) { - case DeviceAllocationType::kCPU: - return "CPU"; - case DeviceAllocationType::kCUDA: - return "CUDA"; - case DeviceAllocationType::kCUDA_HOST: - return "CUDA_HOST"; - case DeviceAllocationType::kOPENCL: - return "OPENCL"; - case DeviceAllocationType::kVULKAN: - return "VULKAN"; - case DeviceAllocationType::kMETAL: - return "METAL"; - case DeviceAllocationType::kVPI: - return "VPI"; - case DeviceAllocationType::kROCM: - return "ROCM"; - case DeviceAllocationType::kROCM_HOST: - return "ROCM_HOST"; - case DeviceAllocationType::kEXT_DEV: - return "EXT_DEV"; - case DeviceAllocationType::kCUDA_MANAGED: - return "CUDA_MANAGED"; - case DeviceAllocationType::kONEAPI: - return "ONEAPI"; - case DeviceAllocationType::kWEBGPU: - return "WEBGPU"; - case DeviceAllocationType::kHEXAGON: - return "HEXAGON"; - } - return ""; -} - } // namespace arrow diff --git a/cpp/src/arrow/device_allocation_type_set.cc b/cpp/src/arrow/device_allocation_type_set.cc new file mode 100644 index 00000000000..afe0834864a --- /dev/null +++ b/cpp/src/arrow/device_allocation_type_set.cc @@ -0,0 +1,80 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include + +#include "arrow/device_allocation_type.h" +#include "arrow/device_allocation_type_set.h" + +namespace arrow { + +const char* DeviceAllocationTypeToCStr(DeviceAllocationType type) { + switch (type) { + case DeviceAllocationType::kCPU: + return "CPU"; + case DeviceAllocationType::kCUDA: + return "CUDA"; + case DeviceAllocationType::kCUDA_HOST: + return "CUDA_HOST"; + case DeviceAllocationType::kOPENCL: + return "OPENCL"; + case DeviceAllocationType::kVULKAN: + return "VULKAN"; + case DeviceAllocationType::kMETAL: + return "METAL"; + case DeviceAllocationType::kVPI: + return "VPI"; + case DeviceAllocationType::kROCM: + return "ROCM"; + case DeviceAllocationType::kROCM_HOST: + return "ROCM_HOST"; + case DeviceAllocationType::kEXT_DEV: + return "EXT_DEV"; + case DeviceAllocationType::kCUDA_MANAGED: + return "CUDA_MANAGED"; + case DeviceAllocationType::kONEAPI: + return "ONEAPI"; + case DeviceAllocationType::kWEBGPU: + return "WEBGPU"; + case DeviceAllocationType::kHEXAGON: + return "HEXAGON"; + } + return ""; +} + +std::string DeviceAllocationTypeSet::ToString() const { + std::string result = "{"; + for (int i = 1; i <= kDeviceAllocationTypeMax; i++) { + if (device_type_bitset_.test(i)) { + // Skip all the unused values in the enum. + switch (i) { + case 0: + case 5: + case 6: + continue; + } + if (result.size() > 1) { + result += ", "; + } + result += DeviceAllocationTypeToCStr(static_cast(i)); + } + } + result += "}"; + return result; +} + +} // namespace arrow diff --git a/cpp/src/arrow/device_allocation_type_set.h b/cpp/src/arrow/device_allocation_type_set.h new file mode 100644 index 00000000000..a52a4d73836 --- /dev/null +++ b/cpp/src/arrow/device_allocation_type_set.h @@ -0,0 +1,92 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include + +#include "arrow/device_allocation_type.h" +#include "arrow/util/visibility.h" + +namespace arrow { + +ARROW_EXPORT +const char* DeviceAllocationTypeToCStr(DeviceAllocationType type); + +class ARROW_EXPORT DeviceAllocationTypeSet { + private: + std::bitset device_type_bitset_; + + public: + /// \brief Construct an empty set of device types. + DeviceAllocationTypeSet() = default; + + /// \brief Construct a set of device types with a single device type. + DeviceAllocationTypeSet( // NOLINT implicit construction + DeviceAllocationType accepted_device_type) { + add(accepted_device_type); + } + + /// \brief Construct a set of device types containing only "kCPU". + static DeviceAllocationTypeSet CpuOnly() { + return DeviceAllocationTypeSet{DeviceAllocationType::kCPU}; + } + + /// \brief Construct a set of device types containing all device types. + static DeviceAllocationTypeSet All() { + DeviceAllocationTypeSet all; + all.device_type_bitset_.set(); + // Don't set the invalid enum values. + all.device_type_bitset_.reset(0); + all.device_type_bitset_.reset(5); + all.device_type_bitset_.reset(6); + return all; + } + + /// \brief Add a device type to the set of device types. + void add(DeviceAllocationType device_type) { + device_type_bitset_.set(static_cast(device_type)); + } + + /// \brief Remove a device type from the set of device types. + void remove(DeviceAllocationType device_type) { + device_type_bitset_.reset(static_cast(device_type)); + } + + /// \brief Return true if the set of accepted device types includes the + /// device type. + bool contains(DeviceAllocationType device_type) const { + return device_type_bitset_.test(static_cast(device_type)); + } + + /// \brief Add all device types from another set to this set. + void Add(DeviceAllocationTypeSet other) { + device_type_bitset_ |= other.device_type_bitset_; + } + + /// \brief Return true if the set of accepted device types includes all the + /// device types in the other set. + bool Contains(DeviceAllocationTypeSet other) const { + // other \subseteq this <==> (other \intersect this == other) + return (other.device_type_bitset_ & device_type_bitset_) == other.device_type_bitset_; + } + + std::string ToString() const; +}; + +} // namespace arrow diff --git a/cpp/src/arrow/gpu/cuda_context.cc b/cpp/src/arrow/gpu/cuda_context.cc index 988cc1f25b9..347cd7cb5d9 100644 --- a/cpp/src/arrow/gpu/cuda_context.cc +++ b/cpp/src/arrow/gpu/cuda_context.cc @@ -26,6 +26,7 @@ #include #include +#include "arrow/device_allocation_type.h" #include "arrow/gpu/cuda_internal.h" #include "arrow/gpu/cuda_memory.h" #include "arrow/util/checked_cast.h" diff --git a/cpp/src/arrow/pretty_print.cc b/cpp/src/arrow/pretty_print.cc index c5905d0c8c5..fb48cf71fcf 100644 --- a/cpp/src/arrow/pretty_print.cc +++ b/cpp/src/arrow/pretty_print.cc @@ -32,6 +32,7 @@ #include "arrow/array.h" #include "arrow/chunked_array.h" +#include "arrow/device_allocation_type.h" #include "arrow/record_batch.h" #include "arrow/status.h" #include "arrow/table.h" diff --git a/cpp/src/arrow/record_batch.cc b/cpp/src/arrow/record_batch.cc index e3a8c0d710c..be7d6c50e86 100644 --- a/cpp/src/arrow/record_batch.cc +++ b/cpp/src/arrow/record_batch.cc @@ -28,6 +28,7 @@ #include "arrow/array.h" #include "arrow/array/concatenate.h" #include "arrow/array/validate.h" +#include "arrow/device_allocation_type.h" #include "arrow/pretty_print.h" #include "arrow/status.h" #include "arrow/table.h" diff --git a/cpp/src/arrow/record_batch.h b/cpp/src/arrow/record_batch.h index 95596e9c155..b4b351ce22a 100644 --- a/cpp/src/arrow/record_batch.h +++ b/cpp/src/arrow/record_batch.h @@ -24,6 +24,7 @@ #include "arrow/compare.h" #include "arrow/device.h" +#include "arrow/device_allocation_type.h" #include "arrow/result.h" #include "arrow/status.h" #include "arrow/type_fwd.h" From 338ca42d5dfdc7880df88418c4ab699913f4c4b0 Mon Sep 17 00:00:00 2001 From: Felipe Oliveira Carvalho Date: Tue, 27 Aug 2024 14:14:56 -0300 Subject: [PATCH 3/7] cpp: Move DeviceAllocationType to type_fwd.h --- cpp/src/arrow/array/array_base.h | 1 - cpp/src/arrow/array/array_test.cc | 1 - cpp/src/arrow/array/data.cc | 1 - cpp/src/arrow/buffer.h | 1 - cpp/src/arrow/buffer_test.cc | 1 - cpp/src/arrow/c/bridge.h | 1 - cpp/src/arrow/device.cc | 1 - cpp/src/arrow/device.h | 1 - cpp/src/arrow/device_allocation_type.h | 41 --------------------- cpp/src/arrow/device_allocation_type_set.cc | 2 +- cpp/src/arrow/device_allocation_type_set.h | 2 +- cpp/src/arrow/gpu/cuda_context.cc | 1 - cpp/src/arrow/pretty_print.cc | 1 - cpp/src/arrow/record_batch.cc | 1 - cpp/src/arrow/record_batch.h | 1 - cpp/src/arrow/type_fwd.h | 19 ++++++++++ 16 files changed, 21 insertions(+), 55 deletions(-) delete mode 100644 cpp/src/arrow/device_allocation_type.h diff --git a/cpp/src/arrow/array/array_base.h b/cpp/src/arrow/array/array_base.h index ad8f84a4e17..716ae072206 100644 --- a/cpp/src/arrow/array/array_base.h +++ b/cpp/src/arrow/array/array_base.h @@ -26,7 +26,6 @@ #include "arrow/array/data.h" #include "arrow/buffer.h" #include "arrow/compare.h" -#include "arrow/device_allocation_type.h" #include "arrow/result.h" #include "arrow/status.h" #include "arrow/type.h" diff --git a/cpp/src/arrow/array/array_test.cc b/cpp/src/arrow/array/array_test.cc index cdc8be813d2..32806d9d2ed 100644 --- a/cpp/src/arrow/array/array_test.cc +++ b/cpp/src/arrow/array/array_test.cc @@ -45,7 +45,6 @@ #include "arrow/buffer.h" #include "arrow/buffer_builder.h" #include "arrow/compare.h" -#include "arrow/device_allocation_type.h" #include "arrow/result.h" #include "arrow/scalar.h" #include "arrow/status.h" diff --git a/cpp/src/arrow/array/data.cc b/cpp/src/arrow/array/data.cc index bf611a63141..83eeb56c496 100644 --- a/cpp/src/arrow/array/data.cc +++ b/cpp/src/arrow/array/data.cc @@ -28,7 +28,6 @@ #include "arrow/array/util.h" #include "arrow/buffer.h" #include "arrow/device.h" -#include "arrow/device_allocation_type.h" #include "arrow/scalar.h" #include "arrow/status.h" #include "arrow/type.h" diff --git a/cpp/src/arrow/buffer.h b/cpp/src/arrow/buffer.h index adef81f4062..fbf4a22e350 100644 --- a/cpp/src/arrow/buffer.h +++ b/cpp/src/arrow/buffer.h @@ -27,7 +27,6 @@ #include #include "arrow/device.h" -#include "arrow/device_allocation_type.h" #include "arrow/status.h" #include "arrow/type_fwd.h" #include "arrow/util/macros.h" diff --git a/cpp/src/arrow/buffer_test.cc b/cpp/src/arrow/buffer_test.cc index 2e5b2ac6ea5..06ed0bfba04 100644 --- a/cpp/src/arrow/buffer_test.cc +++ b/cpp/src/arrow/buffer_test.cc @@ -29,7 +29,6 @@ #include "arrow/buffer.h" #include "arrow/buffer_builder.h" #include "arrow/device.h" -#include "arrow/device_allocation_type.h" #include "arrow/io/interfaces.h" #include "arrow/memory_pool.h" #include "arrow/status.h" diff --git a/cpp/src/arrow/c/bridge.h b/cpp/src/arrow/c/bridge.h index 59168eda973..45367e4f930 100644 --- a/cpp/src/arrow/c/bridge.h +++ b/cpp/src/arrow/c/bridge.h @@ -23,7 +23,6 @@ #include "arrow/c/abi.h" #include "arrow/device.h" -#include "arrow/device_allocation_type.h" #include "arrow/result.h" #include "arrow/status.h" #include "arrow/type_fwd.h" diff --git a/cpp/src/arrow/device.cc b/cpp/src/arrow/device.cc index d08271140f0..c2195e89e75 100644 --- a/cpp/src/arrow/device.cc +++ b/cpp/src/arrow/device.cc @@ -24,7 +24,6 @@ #include "arrow/array.h" #include "arrow/buffer.h" -#include "arrow/device_allocation_type.h" #include "arrow/io/memory.h" #include "arrow/record_batch.h" #include "arrow/result.h" diff --git a/cpp/src/arrow/device.h b/cpp/src/arrow/device.h index 00a365c307f..1dbe5b4b13e 100644 --- a/cpp/src/arrow/device.h +++ b/cpp/src/arrow/device.h @@ -22,7 +22,6 @@ #include #include -#include "arrow/device_allocation_type.h" #include "arrow/io/type_fwd.h" #include "arrow/result.h" #include "arrow/status.h" diff --git a/cpp/src/arrow/device_allocation_type.h b/cpp/src/arrow/device_allocation_type.h deleted file mode 100644 index a8a99867bdd..00000000000 --- a/cpp/src/arrow/device_allocation_type.h +++ /dev/null @@ -1,41 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -#pragma once - -namespace arrow { - -/// \brief EXPERIMENTAL: Device type enum which matches up with C Data Device types -enum class DeviceAllocationType : char { - kCPU = 1, - kCUDA = 2, - kCUDA_HOST = 3, - kOPENCL = 4, - kVULKAN = 7, - kMETAL = 8, - kVPI = 9, - kROCM = 10, - kROCM_HOST = 11, - kEXT_DEV = 12, - kCUDA_MANAGED = 13, - kONEAPI = 14, - kWEBGPU = 15, - kHEXAGON = 16, -}; -constexpr int kDeviceAllocationTypeMax = 16; - -} // namespace arrow diff --git a/cpp/src/arrow/device_allocation_type_set.cc b/cpp/src/arrow/device_allocation_type_set.cc index afe0834864a..83e9e57f2ee 100644 --- a/cpp/src/arrow/device_allocation_type_set.cc +++ b/cpp/src/arrow/device_allocation_type_set.cc @@ -17,8 +17,8 @@ #include -#include "arrow/device_allocation_type.h" #include "arrow/device_allocation_type_set.h" +#include "arrow/type_fwd.h" namespace arrow { diff --git a/cpp/src/arrow/device_allocation_type_set.h b/cpp/src/arrow/device_allocation_type_set.h index a52a4d73836..19798bc5f11 100644 --- a/cpp/src/arrow/device_allocation_type_set.h +++ b/cpp/src/arrow/device_allocation_type_set.h @@ -20,7 +20,7 @@ #include #include -#include "arrow/device_allocation_type.h" +#include "arrow/type_fwd.h" #include "arrow/util/visibility.h" namespace arrow { diff --git a/cpp/src/arrow/gpu/cuda_context.cc b/cpp/src/arrow/gpu/cuda_context.cc index 347cd7cb5d9..988cc1f25b9 100644 --- a/cpp/src/arrow/gpu/cuda_context.cc +++ b/cpp/src/arrow/gpu/cuda_context.cc @@ -26,7 +26,6 @@ #include #include -#include "arrow/device_allocation_type.h" #include "arrow/gpu/cuda_internal.h" #include "arrow/gpu/cuda_memory.h" #include "arrow/util/checked_cast.h" diff --git a/cpp/src/arrow/pretty_print.cc b/cpp/src/arrow/pretty_print.cc index fb48cf71fcf..c5905d0c8c5 100644 --- a/cpp/src/arrow/pretty_print.cc +++ b/cpp/src/arrow/pretty_print.cc @@ -32,7 +32,6 @@ #include "arrow/array.h" #include "arrow/chunked_array.h" -#include "arrow/device_allocation_type.h" #include "arrow/record_batch.h" #include "arrow/status.h" #include "arrow/table.h" diff --git a/cpp/src/arrow/record_batch.cc b/cpp/src/arrow/record_batch.cc index be7d6c50e86..e3a8c0d710c 100644 --- a/cpp/src/arrow/record_batch.cc +++ b/cpp/src/arrow/record_batch.cc @@ -28,7 +28,6 @@ #include "arrow/array.h" #include "arrow/array/concatenate.h" #include "arrow/array/validate.h" -#include "arrow/device_allocation_type.h" #include "arrow/pretty_print.h" #include "arrow/status.h" #include "arrow/table.h" diff --git a/cpp/src/arrow/record_batch.h b/cpp/src/arrow/record_batch.h index b4b351ce22a..95596e9c155 100644 --- a/cpp/src/arrow/record_batch.h +++ b/cpp/src/arrow/record_batch.h @@ -24,7 +24,6 @@ #include "arrow/compare.h" #include "arrow/device.h" -#include "arrow/device_allocation_type.h" #include "arrow/result.h" #include "arrow/status.h" #include "arrow/type_fwd.h" diff --git a/cpp/src/arrow/type_fwd.h b/cpp/src/arrow/type_fwd.h index 08777d247ed..e36058fc446 100644 --- a/cpp/src/arrow/type_fwd.h +++ b/cpp/src/arrow/type_fwd.h @@ -724,4 +724,23 @@ ARROW_EXPORT MemoryPool* default_memory_pool(); constexpr int64_t kDefaultBufferAlignment = 64; +/// \brief EXPERIMENTAL: Device type enum which matches up with C Data Device types +enum class DeviceAllocationType : char { + kCPU = 1, + kCUDA = 2, + kCUDA_HOST = 3, + kOPENCL = 4, + kVULKAN = 7, + kMETAL = 8, + kVPI = 9, + kROCM = 10, + kROCM_HOST = 11, + kEXT_DEV = 12, + kCUDA_MANAGED = 13, + kONEAPI = 14, + kWEBGPU = 15, + kHEXAGON = 16, +}; +constexpr int kDeviceAllocationTypeMax = 16; + } // namespace arrow From 71722242e687d5dbf4e36bcabc54a1c94152fb76 Mon Sep 17 00:00:00 2001 From: Felipe Oliveira Carvalho Date: Tue, 27 Aug 2024 15:50:08 -0300 Subject: [PATCH 4/7] cpp: Add Datum::device_types() definition --- cpp/src/arrow/chunked_array.cc | 2 +- cpp/src/arrow/chunked_array.h | 2 +- cpp/src/arrow/datum.cc | 39 ++++++++++++++++++++++++++++++++++ cpp/src/arrow/datum.h | 2 +- 4 files changed, 42 insertions(+), 3 deletions(-) diff --git a/cpp/src/arrow/chunked_array.cc b/cpp/src/arrow/chunked_array.cc index a5e590087e2..dd6aa51534f 100644 --- a/cpp/src/arrow/chunked_array.cc +++ b/cpp/src/arrow/chunked_array.cc @@ -87,7 +87,7 @@ Result> ChunkedArray::MakeEmpty( return std::make_shared(std::move(new_chunks)); } -DeviceAllocationTypeSet ChunkedArray::DeviceTypeSet() const { +DeviceAllocationTypeSet ChunkedArray::device_types() const { if (chunks_.empty()) { // An empty ChunkedArray is considered to be CPU-only. return DeviceAllocationTypeSet::CpuOnly(); diff --git a/cpp/src/arrow/chunked_array.h b/cpp/src/arrow/chunked_array.h index f2d07dfed3a..4a3ba38c12c 100644 --- a/cpp/src/arrow/chunked_array.h +++ b/cpp/src/arrow/chunked_array.h @@ -120,7 +120,7 @@ class ARROW_EXPORT ChunkedArray { /// \return The set of device allocation types used by the chunks in this /// chunked array. - DeviceAllocationTypeSet DeviceTypeSet() const; + DeviceAllocationTypeSet device_types() const; /// \brief Construct a zero-copy slice of the chunked array with the /// indicated offset and length diff --git a/cpp/src/arrow/datum.cc b/cpp/src/arrow/datum.cc index cd9da31982f..b19d1864475 100644 --- a/cpp/src/arrow/datum.cc +++ b/cpp/src/arrow/datum.cc @@ -157,6 +157,45 @@ ArrayVector Datum::chunks() const { return this->chunked_array()->chunks(); } +DeviceAllocationTypeSet Datum::device_types() const { + switch (kind()) { + case NONE: + break; + case SCALAR: + // Scalars are asssumed as always residing in CPU memory for now. + return DeviceAllocationTypeSet::CpuOnly(); + case ARRAY: + return DeviceAllocationTypeSet{array()->device_type()}; + case CHUNKED_ARRAY: + return chunked_array()->device_types(); + case RECORD_BATCH: { + auto& columns = record_batch()->columns(); + if (columns.empty()) { + // An empty RecordBatch is considered to be CPU-only. + return DeviceAllocationTypeSet::CpuOnly(); + } + DeviceAllocationTypeSet set; + for (const auto& column : columns) { + set.add(column->device_type()); + } + return set; + } + case TABLE: { + auto& columns = table()->columns(); + if (columns.empty()) { + // An empty Table is considered to be CPU-only. + return DeviceAllocationTypeSet::CpuOnly(); + } + DeviceAllocationTypeSet set; + for (const auto& column : columns) { + set.Add(column->device_types()); + } + return set; + } + } + return {}; +} + bool Datum::Equals(const Datum& other) const { if (this->kind() != other.kind()) return false; diff --git a/cpp/src/arrow/datum.h b/cpp/src/arrow/datum.h index 5dc8f9ea41b..4a88e7a8112 100644 --- a/cpp/src/arrow/datum.h +++ b/cpp/src/arrow/datum.h @@ -296,7 +296,7 @@ struct ARROW_EXPORT Datum { /// \return empty if not arraylike ArrayVector chunks() const; - DeviceAllocationTypeSet DeviceTypeSet() const; + DeviceAllocationTypeSet device_types() const; /// \brief True if the two data are equal bool Equals(const Datum& other) const; From 46b0400fb780a100d8ff66720e5822bce4646c35 Mon Sep 17 00:00:00 2001 From: Felipe Oliveira Carvalho Date: Tue, 27 Aug 2024 15:59:23 -0300 Subject: [PATCH 5/7] fixup! cpp: Add the DeviceAllocationTypeSet class --- cpp/src/arrow/chunked_array.h | 1 - 1 file changed, 1 deletion(-) diff --git a/cpp/src/arrow/chunked_array.h b/cpp/src/arrow/chunked_array.h index 4a3ba38c12c..1424a991a67 100644 --- a/cpp/src/arrow/chunked_array.h +++ b/cpp/src/arrow/chunked_array.h @@ -17,7 +17,6 @@ #pragma once -#include #include #include #include From 2b1bd8210482cdbb931be169590cd588065b5850 Mon Sep 17 00:00:00 2001 From: Felipe Oliveira Carvalho Date: Tue, 27 Aug 2024 17:51:59 -0300 Subject: [PATCH 6/7] cpp: Forward-declare DeviceAllocationTypeSet --- cpp/src/arrow/type_fwd.h | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cpp/src/arrow/type_fwd.h b/cpp/src/arrow/type_fwd.h index e36058fc446..8faebe217f1 100644 --- a/cpp/src/arrow/type_fwd.h +++ b/cpp/src/arrow/type_fwd.h @@ -743,4 +743,6 @@ enum class DeviceAllocationType : char { }; constexpr int kDeviceAllocationTypeMax = 16; +class DeviceAllocationTypeSet; + } // namespace arrow From 8cd3da67249628204ded22361205d8e47f228bd1 Mon Sep 17 00:00:00 2001 From: Felipe Oliveira Carvalho Date: Tue, 27 Aug 2024 18:15:21 -0300 Subject: [PATCH 7/7] cpp: Add ChunkedArray::is_cpu() and tests --- cpp/src/arrow/chunked_array.h | 3 +++ cpp/src/arrow/chunked_array_test.cc | 5 +++++ cpp/src/arrow/device_allocation_type_set.h | 5 +++++ 3 files changed, 13 insertions(+) diff --git a/cpp/src/arrow/chunked_array.h b/cpp/src/arrow/chunked_array.h index 1424a991a67..c65b6cb6e22 100644 --- a/cpp/src/arrow/chunked_array.h +++ b/cpp/src/arrow/chunked_array.h @@ -121,6 +121,9 @@ class ARROW_EXPORT ChunkedArray { /// chunked array. DeviceAllocationTypeSet device_types() const; + /// \return true if all chunks are allocated on CPU-accessible memory. + bool is_cpu() const { return device_types().is_cpu_only(); } + /// \brief Construct a zero-copy slice of the chunked array with the /// indicated offset and length /// diff --git a/cpp/src/arrow/chunked_array_test.cc b/cpp/src/arrow/chunked_array_test.cc index e9cc283b53c..b796e925000 100644 --- a/cpp/src/arrow/chunked_array_test.cc +++ b/cpp/src/arrow/chunked_array_test.cc @@ -61,12 +61,17 @@ TEST_F(TestChunkedArray, Make) { ChunkedArray::Make({}, int64())); AssertTypeEqual(*int64(), *result->type()); ASSERT_EQ(result->num_chunks(), 0); + // Empty chunked arrays are treated as CPU-allocated. + ASSERT_TRUE(result->is_cpu()); auto chunk0 = ArrayFromJSON(int8(), "[0, 1, 2]"); auto chunk1 = ArrayFromJSON(int16(), "[3, 4, 5]"); ASSERT_OK_AND_ASSIGN(result, ChunkedArray::Make({chunk0, chunk0})); ASSERT_OK_AND_ASSIGN(auto result2, ChunkedArray::Make({chunk0, chunk0}, int8())); + // All chunks are CPU-accessible. + ASSERT_TRUE(result->is_cpu()); + ASSERT_TRUE(result2->is_cpu()); AssertChunkedEqual(*result, *result2); ASSERT_RAISES(TypeError, ChunkedArray::Make({chunk0, chunk1})); diff --git a/cpp/src/arrow/device_allocation_type_set.h b/cpp/src/arrow/device_allocation_type_set.h index 19798bc5f11..974367307e6 100644 --- a/cpp/src/arrow/device_allocation_type_set.h +++ b/cpp/src/arrow/device_allocation_type_set.h @@ -68,6 +68,11 @@ class ARROW_EXPORT DeviceAllocationTypeSet { device_type_bitset_.reset(static_cast(device_type)); } + /// \brief Return true iff the set only contains the CPU device type. + bool is_cpu_only() const { + return device_type_bitset_ == CpuOnly().device_type_bitset_; + } + /// \brief Return true if the set of accepted device types includes the /// device type. bool contains(DeviceAllocationType device_type) const {