include/arrow/device_allocation_type_set.h

0001 // Licensed to the Apache Software Foundation (ASF) under one
0002 // or more contributor license agreements.  See the NOTICE file
0003 // distributed with this work for additional information
0004 // regarding copyright ownership.  The ASF licenses this file
0005 // to you under the Apache License, Version 2.0 (the
0006 // "License"); you may not use this file except in compliance
0007 // with the License.  You may obtain a copy of the License at
0008 //
0009 //   http://www.apache.org/licenses/LICENSE-2.0
0010 //
0011 // Unless required by applicable law or agreed to in writing,
0012 // software distributed under the License is distributed on an
0013 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
0014 // KIND, either express or implied.  See the License for the
0015 // specific language governing permissions and limitations
0016 // under the License.
0017
0018 #pragma once
0019
0020 #include <bitset>
0021 #include <string>
0022
0023 #include "arrow/type_fwd.h"
0024 #include "arrow/util/visibility.h"
0025
0026 namespace arrow {
0027
0028 ARROW_EXPORT
0029 const char* DeviceAllocationTypeToCStr(DeviceAllocationType type);
0030
0031 class ARROW_EXPORT DeviceAllocationTypeSet {
0032  private:
0033   std::bitset<kDeviceAllocationTypeMax + 1> device_type_bitset_;
0034
0035  public:
0036   /// \brief Construct an empty set of device types.
0037   DeviceAllocationTypeSet() = default;
0038
0039   /// \brief Construct a set of device types with a single device type.
0040   DeviceAllocationTypeSet(  // NOLINT implicit construction
0041       DeviceAllocationType accepted_device_type) {
0042     add(accepted_device_type);
0043   }
0044
0045   /// \brief Construct a set of device types containing only "kCPU".
0046   static DeviceAllocationTypeSet CpuOnly() {
0047     return DeviceAllocationTypeSet{DeviceAllocationType::kCPU};
0048   }
0049
0050   /// \brief Construct a set of device types containing all device types.
0051   static DeviceAllocationTypeSet All() {
0052     DeviceAllocationTypeSet all;
0053     all.device_type_bitset_.set();
0054     // Don't set the invalid enum values.
0055     all.device_type_bitset_.reset(0);
0056     all.device_type_bitset_.reset(5);
0057     all.device_type_bitset_.reset(6);
0058     return all;
0059   }
0060
0061   /// \brief Add a device type to the set of device types.
0062   void add(DeviceAllocationType device_type) {
0063     device_type_bitset_.set(static_cast<int>(device_type));
0064   }
0065
0066   /// \brief Remove a device type from the set of device types.
0067   void remove(DeviceAllocationType device_type) {
0068     device_type_bitset_.reset(static_cast<int>(device_type));
0069   }
0070
0071   /// \brief Return true iff the set only contains the CPU device type.
0072   bool is_cpu_only() const {
0073     return device_type_bitset_ == CpuOnly().device_type_bitset_;
0074   }
0075
0076   /// \brief Return true if the set of accepted device types includes the
0077   /// device type.
0078   bool contains(DeviceAllocationType device_type) const {
0079     return device_type_bitset_.test(static_cast<int>(device_type));
0080   }
0081
0082   /// \brief Add all device types from another set to this set.
0083   void Add(DeviceAllocationTypeSet other) {
0084     device_type_bitset_ |= other.device_type_bitset_;
0085   }
0086
0087   /// \brief Return true if the set of accepted device types includes all the
0088   /// device types in the other set.
0089   bool Contains(DeviceAllocationTypeSet other) const {
0090     // other \subseteq this <==> (other \intersect this == other)
0091     return (other.device_type_bitset_ & device_type_bitset_) == other.device_type_bitset_;
0092   }
0093
0094   std::string ToString() const;
0095 };
0096
0097 }  // namespace arrow