zeroshade commented on code in PR #36489:
URL: https://github.com/apache/arrow/pull/36489#discussion_r1254563903
##########
cpp/src/arrow/buffer.h:
##########
@@ -57,18 +57,25 @@ class ARROW_EXPORT Buffer {
///
/// \note The passed memory must be kept alive through some other means
Buffer(const uint8_t* data, int64_t size)
- : is_mutable_(false), is_cpu_(true), data_(data), size_(size),
capacity_(size) {
+ : is_mutable_(false), is_cpu_(true), data_(data), size_(size),
capacity_(size), device_type_(DeviceType::CPU) {
SetMemoryManager(default_cpu_memory_manager());
}
Buffer(const uint8_t* data, int64_t size, std::shared_ptr<MemoryManager> mm,
- std::shared_ptr<Buffer> parent = NULLPTR)
+ std::shared_ptr<Buffer> parent = NULLPTR, DeviceType device_type =
DeviceType::UNKNOWN)
: is_mutable_(false),
data_(data),
size_(size),
capacity_(size),
parent_(std::move(parent)) {
+ // will set device_type from the memory manager
SetMemoryManager(std::move(mm));
+ // if a device type is specified, use that instead. for example:
+ // CUDA_HOST. The CudaMemoryManager will set device_type_ to CUDA,
+ // but you can specify CUDA_HOST as the device type to override it.
+ if (device_type != DeviceType::UNKNOWN) {
+ device_type_ = device_type;
+ }
Review Comment:
Possibly, we'd have to add a `Validate` or `Valid` method to the memory
manager to do so, to allow the memory manager to specify if the device_type is
valid for it. Not sure how worthwhile that is though.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
