marin-ma commented on code in PR #13337:
URL: https://github.com/apache/gluten/pull/13337#discussion_r4312115323
##########
cpp/core/shuffle/Utils.cc:
##########
@@ -265,3 +265,18 @@ std::shared_ptr<arrow::Buffer>
gluten::zeroLengthNullBuffer() {
static std::shared_ptr<arrow::Buffer> kNullBuffer =
std::make_shared<arrow::Buffer>(nullptr, 0);
return kNullBuffer;
}
+
+arrow::Result<std::shared_ptr<arrow::ResizableBuffer>>
gluten::allocatePaddedBuffer(
+ int64_t size,
+ arrow::MemoryPool* pool) {
+ ARROW_ASSIGN_OR_RAISE(auto buffer, arrow::AllocateResizableBuffer(size +
kVeloxBufferPadding, pool));
+ // Shrinking without shrink_to_fit keeps the capacity.
+ RETURN_NOT_OK(buffer->Resize(size, /*shrink_to_fit=*/false));
Review Comment:
Perhaps add `buffer->ZeroPadding()` here for the output from shuffle reader.
##########
cpp/core/shuffle/Utils.h:
##########
@@ -73,6 +73,19 @@ int64_t getMaxCompressedBufferSize(
std::shared_ptr<arrow::Buffer> zeroLengthNullBuffer();
+// Velox may read past the end of a buffer with SIMD loads, up to
simd::kPadding bytes,
+// which its AlignedBuffer keeps addressable. The shuffle reader hands
deserialized buffers
+// to Velox without a copy, so they reserve the same spare capacity. 64 covers
AVX-512.
+constexpr int64_t kVeloxBufferPadding = 64;
Review Comment:
Maybe consider changing the name to `kSimdPadding`, since the point of this
change is that Gluten guarantees SIMD padding on its output.
##########
cpp/core/shuffle/Utils.cc:
##########
@@ -265,3 +265,18 @@ std::shared_ptr<arrow::Buffer>
gluten::zeroLengthNullBuffer() {
static std::shared_ptr<arrow::Buffer> kNullBuffer =
std::make_shared<arrow::Buffer>(nullptr, 0);
return kNullBuffer;
}
+
+arrow::Result<std::shared_ptr<arrow::ResizableBuffer>>
gluten::allocatePaddedBuffer(
+ int64_t size,
+ arrow::MemoryPool* pool) {
+ ARROW_ASSIGN_OR_RAISE(auto buffer, arrow::AllocateResizableBuffer(size +
kVeloxBufferPadding, pool));
+ // Shrinking without shrink_to_fit keeps the capacity.
+ RETURN_NOT_OK(buffer->Resize(size, /*shrink_to_fit=*/false));
+ return std::shared_ptr<arrow::ResizableBuffer>(std::move(buffer));
Review Comment:
We can keep the same return type as the `arrow::AllocateResizableBuffer`,
which is `arrow::Result<std::unique_ptr<arrow::ResizableBuffer>>`
##########
cpp/core/shuffle/Utils.cc:
##########
@@ -265,3 +265,18 @@ std::shared_ptr<arrow::Buffer>
gluten::zeroLengthNullBuffer() {
static std::shared_ptr<arrow::Buffer> kNullBuffer =
std::make_shared<arrow::Buffer>(nullptr, 0);
return kNullBuffer;
}
+
+arrow::Result<std::shared_ptr<arrow::ResizableBuffer>>
gluten::allocatePaddedBuffer(
+ int64_t size,
+ arrow::MemoryPool* pool) {
+ ARROW_ASSIGN_OR_RAISE(auto buffer, arrow::AllocateResizableBuffer(size +
kVeloxBufferPadding, pool));
+ // Shrinking without shrink_to_fit keeps the capacity.
+ RETURN_NOT_OK(buffer->Resize(size, /*shrink_to_fit=*/false));
+ return std::shared_ptr<arrow::ResizableBuffer>(std::move(buffer));
+}
+
+arrow::Status gluten::resizeWithPadding(arrow::ResizableBuffer& buffer,
int64_t size) {
+ // Reserve() only reallocates if the capacity is too small.
+ RETURN_NOT_OK(buffer.Reserve(size + kVeloxBufferPadding));
+ return buffer.Resize(size, /*shrink_to_fit=*/false);
Review Comment:
Ditto: `buffer->ZeroPadding()`
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
