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/chunked_array.cc b/cpp/src/arrow/chunked_array.cc index c36b736d5d5..dd6aa51534f 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" @@ -86,6 +87,18 @@ Result> ChunkedArray::MakeEmpty( return std::make_shared(std::move(new_chunks)); } +DeviceAllocationTypeSet ChunkedArray::device_types() 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..c65b6cb6e22 100644 --- a/cpp/src/arrow/chunked_array.h +++ b/cpp/src/arrow/chunked_array.h @@ -25,6 +25,7 @@ #include "arrow/chunk_resolver.h" #include "arrow/compare.h" +#include "arrow/device_allocation_type_set.h" #include "arrow/result.h" #include "arrow/status.h" #include "arrow/type_fwd.h" @@ -116,6 +117,13 @@ 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 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/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..b19d1864475 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" @@ -156,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 31b2d2274c9..4a88e7a8112 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/device_allocation_type_set.h" #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 device_types() 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..1dbe5b4b13e 100644 --- a/cpp/src/arrow/device.h +++ b/cpp/src/arrow/device.h @@ -32,24 +32,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_set.cc b/cpp/src/arrow/device_allocation_type_set.cc new file mode 100644 index 00000000000..83e9e57f2ee --- /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_set.h" +#include "arrow/type_fwd.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..974367307e6 --- /dev/null +++ b/cpp/src/arrow/device_allocation_type_set.h @@ -0,0 +1,97 @@ +// 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/type_fwd.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 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 { + 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/type_fwd.h b/cpp/src/arrow/type_fwd.h index 08777d247ed..8faebe217f1 100644 --- a/cpp/src/arrow/type_fwd.h +++ b/cpp/src/arrow/type_fwd.h @@ -724,4 +724,25 @@ 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; + +class DeviceAllocationTypeSet; + } // namespace arrow