-
Notifications
You must be signed in to change notification settings - Fork 3.6k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
GH-36488: [C++] Import/Export ArrowDeviceArray #36489
Changes from 2 commits
ab98249
4d0ac93
30256c1
d591208
58c15b0
2d3bcb6
ff2dd19
a858ccb
6e7de49
f0deaab
849457b
9244ea4
f79d2a2
5e1939f
1ca39b8
d90ce9e
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -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 | ||
zeroshade marked this conversation as resolved.
Show resolved
Hide resolved
|
||
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. | ||
westonpace marked this conversation as resolved.
Show resolved
Hide resolved
|
||
if (device_type != DeviceType::UNKNOWN) { | ||
device_type_ = device_type; | ||
} | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Would a DCHECK be useful here to confirm that There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Possibly, we'd have to add a |
||
} | ||
|
||
Buffer(uintptr_t address, int64_t size, std::shared_ptr<MemoryManager> mm, | ||
|
@@ -240,6 +247,8 @@ class ARROW_EXPORT Buffer { | |
|
||
const std::shared_ptr<MemoryManager>& memory_manager() const { return memory_manager_; } | ||
|
||
DeviceType device_type() const { return device_type_; } | ||
|
||
std::shared_ptr<Buffer> parent() const { return parent_; } | ||
|
||
/// \brief Get a RandomAccessFile for reading a buffer | ||
|
@@ -294,6 +303,7 @@ class ARROW_EXPORT Buffer { | |
const uint8_t* data_; | ||
int64_t size_; | ||
int64_t capacity_; | ||
DeviceType device_type_; | ||
|
||
// null by default, but may be set | ||
std::shared_ptr<Buffer> parent_; | ||
|
@@ -309,6 +319,7 @@ class ARROW_EXPORT Buffer { | |
void SetMemoryManager(std::shared_ptr<MemoryManager> mm) { | ||
memory_manager_ = std::move(mm); | ||
is_cpu_ = memory_manager_->is_cpu(); | ||
device_type_ = memory_manager_->device()->device_type(); | ||
} | ||
|
||
private: | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Do we still need the
is_cpu_
boolean? Is there a situation in whichis_cpu_
isfalse
anddevice_type_ != CPU
?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The
is_cpu_
boolean is still useful because of pinned host memory. BothCUDA_HOST
andROCM_HOST
are pinned CPU memory which would have the corresponding device type (since the memory is managed by the device drivers) but is accessible to the CPU.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
While the pointers are accessible to the CPU, there's no guarantee that there isn't stream ordered work happening on them where it could cause a race condition to CPU code accessing the buffer. Given
is_cpu_
is existing ABI, I'm wondering if it wouldn't be a safer option to have it returnFalse
in these cases where if someone wants to do CPU things againstCUDA_HOST
orROCM_HOST
buffers they should explicitly check the device type instead?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hmm, I think the intent for the
is_cpu_
flag was solely to know whether or not the data is accessible via CPU rather than caring about the stream ordered work synchronizations. I intend on expanding the existing classes/objects to make it easier to synchronize for stream ordered work as a separate abstraction, so I lean towards leaving it marked astrue
in this case, but I'm open to changing it.@pitrou @felipeblazing @Rhonda85 what do you all think?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I had not envisioned this issue, but I think your approach is correct:
is_cpu() == true
means thataddress()
is a CPU-accessible pointer to the data. Whether or not the "expected" data is already there will depend on the producer.And, yes, a generalized version of
Buffer
could have an optional sync event or something...