This is an automated email from the ASF dual-hosted git repository.
zhouyuan pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 5dc5f002b0 [VL] Keep Velox SIMD padding on shuffle reader buffers
(#13337)
5dc5f002b0 is described below
commit 5dc5f002b0f5c67116e8137b2fccfad1d2a5b60c
Author: Yuan <[email protected]>
AuthorDate: Thu Oct 29 15:06:30 2026 +0000
[VL] Keep Velox SIMD padding on shuffle reader buffers (#13337)
---
cpp/core/shuffle/Payload.cc | 21 ++--
cpp/core/shuffle/Utils.cc | 18 ++++
cpp/core/shuffle/Utils.h | 13 +++
cpp/core/tests/CMakeLists.txt | 1 +
cpp/core/tests/PayloadTest.cc | 163 ++++++++++++++++++++++++++++++++
cpp/velox/shuffle/VeloxShuffleReader.cc | 19 +++-
6 files changed, 221 insertions(+), 14 deletions(-)
diff --git a/cpp/core/shuffle/Payload.cc b/cpp/core/shuffle/Payload.cc
index 5d02a84483..5c830246ee 100644
--- a/cpp/core/shuffle/Payload.cc
+++ b/cpp/core/shuffle/Payload.cc
@@ -170,7 +170,7 @@ readUncompressedBuffer(arrow::io::InputStream* inputStream,
arrow::MemoryPool* p
if (bufferLength == kNullBuffer) {
return nullptr;
}
- ARROW_ASSIGN_OR_RAISE(auto buffer,
arrow::AllocateResizableBuffer(bufferLength, pool));
+ ARROW_ASSIGN_OR_RAISE(auto buffer, allocatePaddedBuffer(bufferLength, pool));
RETURN_NOT_OK(inputStream->Read(bufferLength, buffer->mutable_data()));
return buffer;
}
@@ -204,7 +204,7 @@ arrow::Result<std::shared_ptr<arrow::Buffer>>
readCompressedBuffer(
RETURN_NOT_OK(inputStream->Read(actualCompressedLen,
compressed->mutable_data()));
timer.switchTo(&decompressTime);
- ARROW_ASSIGN_OR_RAISE(auto output,
arrow::AllocateResizableBuffer(uncompressedLength, pool));
+ ARROW_ASSIGN_OR_RAISE(auto output,
allocatePaddedBuffer(uncompressedLength, pool));
RETURN_NOT_OK(TypeAwareCompressCodec::decompress(
compressed->data(), actualCompressedLen,
output->mutable_data(), uncompressedLength)
.status());
@@ -212,7 +212,7 @@ arrow::Result<std::shared_ptr<arrow::Buffer>>
readCompressedBuffer(
}
if (compressedLength == kUncompressedBuffer) {
- ARROW_ASSIGN_OR_RAISE(auto uncompressed,
arrow::AllocateResizableBuffer(uncompressedLength, pool));
+ ARROW_ASSIGN_OR_RAISE(auto uncompressed,
allocatePaddedBuffer(uncompressedLength, pool));
RETURN_NOT_OK(inputStream->Read(uncompressedLength,
uncompressed->mutable_data()));
return uncompressed;
}
@@ -220,7 +220,7 @@ arrow::Result<std::shared_ptr<arrow::Buffer>>
readCompressedBuffer(
RETURN_NOT_OK(inputStream->Read(compressedLength,
compressed->mutable_data()));
timer.switchTo(&decompressTime);
- ARROW_ASSIGN_OR_RAISE(auto output,
arrow::AllocateResizableBuffer(uncompressedLength, pool));
+ ARROW_ASSIGN_OR_RAISE(auto output, allocatePaddedBuffer(uncompressedLength,
pool));
RETURN_NOT_OK(codec->Decompress(compressedLength, compressed->data(),
uncompressedLength, output->mutable_data()));
return output;
}
@@ -451,8 +451,7 @@ arrow::Result<std::unique_ptr<InMemoryPayload>>
InMemoryPayload::merge(
if (!appendBuffer) {
merged[i] = nullptr;
} else {
- ARROW_ASSIGN_OR_RAISE(
- auto buffer,
arrow::AllocateResizableBuffer(arrow::bit_util::BytesForBits(mergedRows),
pool));
+ ARROW_ASSIGN_OR_RAISE(auto buffer,
allocatePaddedBuffer(arrow::bit_util::BytesForBits(mergedRows), pool));
// Source is null, fill all true.
arrow::bit_util::SetBitsTo(buffer->mutable_data(), 0,
source->numRows(), true);
// Write append bits.
@@ -467,10 +466,10 @@ arrow::Result<std::unique_ptr<InMemoryPayload>>
InMemoryPayload::merge(
auto mergedBytes = arrow::bit_util::BytesForBits(mergedRows);
if (resizable) {
// If source is resizable, resize and reuse source.
- RETURN_NOT_OK(resizable->Resize(mergedBytes));
+ RETURN_NOT_OK(resizeWithPadding(*resizable, mergedBytes));
} else {
// Otherwise copy source.
- ARROW_ASSIGN_OR_RAISE(resizable,
arrow::AllocateResizableBuffer(mergedBytes, pool));
+ ARROW_ASSIGN_OR_RAISE(resizable, allocatePaddedBuffer(mergedBytes,
pool));
memcpy(resizable->mutable_data(), sourceBuffer->data(),
sourceBufferSize);
}
if (!appendBuffer) {
@@ -491,10 +490,10 @@ arrow::Result<std::unique_ptr<InMemoryPayload>>
InMemoryPayload::merge(
auto resizable =
std::dynamic_pointer_cast<arrow::ResizableBuffer>(sourceBuffer);
if (resizable) {
// If source is resizable, resize and reuse source.
- RETURN_NOT_OK(resizable->Resize(mergedSize));
+ RETURN_NOT_OK(resizeWithPadding(*resizable, mergedSize));
} else {
// Otherwise copy source.
- ARROW_ASSIGN_OR_RAISE(resizable,
arrow::AllocateResizableBuffer(mergedSize, pool));
+ ARROW_ASSIGN_OR_RAISE(resizable, allocatePaddedBuffer(mergedSize,
pool));
memcpy(resizable->mutable_data(), sourceBuffer->data(),
sourceBufferSize);
}
// Copy append.
@@ -539,7 +538,7 @@ arrow::Status
InMemoryPayload::copyBuffers(arrow::MemoryPool* pool) {
buffer = zeroLengthNullBuffer();
continue;
}
- ARROW_ASSIGN_OR_RAISE(auto copy,
arrow::AllocateResizableBuffer(buffer->size(), pool));
+ ARROW_ASSIGN_OR_RAISE(auto copy, allocatePaddedBuffer(buffer->size(),
pool));
memcpy(copy->mutable_data(), buffer->data(), buffer->size());
buffer = std::move(copy);
}
diff --git a/cpp/core/shuffle/Utils.cc b/cpp/core/shuffle/Utils.cc
index d2c9c6fe6d..49f82357ad 100644
--- a/cpp/core/shuffle/Utils.cc
+++ b/cpp/core/shuffle/Utils.cc
@@ -265,3 +265,21 @@ 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::unique_ptr<arrow::ResizableBuffer>>
gluten::allocatePaddedBuffer(
+ int64_t size,
+ arrow::MemoryPool* pool) {
+ ARROW_ASSIGN_OR_RAISE(auto buffer, arrow::AllocateResizableBuffer(size +
kSimdPadding, pool));
+ // Shrinking without shrink_to_fit keeps the capacity.
+ RETURN_NOT_OK(buffer->Resize(size, /*shrink_to_fit=*/false));
+ buffer->ZeroPadding();
+ return 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 + kSimdPadding));
+ RETURN_NOT_OK(buffer.Resize(size, /*shrink_to_fit=*/false));
+ buffer.ZeroPadding();
+ return arrow::Status::OK();
+}
diff --git a/cpp/core/shuffle/Utils.h b/cpp/core/shuffle/Utils.h
index 540662313e..a0eb27bad0 100644
--- a/cpp/core/shuffle/Utils.h
+++ b/cpp/core/shuffle/Utils.h
@@ -73,6 +73,19 @@ int64_t getMaxCompressedBufferSize(
std::shared_ptr<arrow::Buffer> zeroLengthNullBuffer();
+// Spare capacity, zero-filled, that shuffle guarantees after the buffers it
outputs, so that
+// consumers can read past the end with SIMD loads. E.g. Velox reads up to
simd::kPadding bytes
+// past a buffer, and the shuffle reader hands buffers to Velox without a
copy. 64 covers AVX-512.
+constexpr int64_t kSimdPadding = 64;
+
+// Allocates a buffer of 'size' bytes with at least kSimdPadding bytes of
zeroed spare
+// capacity after it.
+arrow::Result<std::unique_ptr<arrow::ResizableBuffer>>
allocatePaddedBuffer(int64_t size, arrow::MemoryPool* pool);
+
+// Resizes 'buffer' to 'size' bytes, keeping at least kSimdPadding bytes of
zeroed spare
+// capacity after it.
+arrow::Status resizeWithPadding(arrow::ResizableBuffer& buffer, int64_t size);
+
// MmapFileStream is used to optimize sequential file reading. It uses madvise
// to prefetch and release memory timely.
class MmapFileStream : public arrow::io::InputStream {
diff --git a/cpp/core/tests/CMakeLists.txt b/cpp/core/tests/CMakeLists.txt
index c022f37aea..26000453f6 100644
--- a/cpp/core/tests/CMakeLists.txt
+++ b/cpp/core/tests/CMakeLists.txt
@@ -18,3 +18,4 @@ add_test_case(object_store_test SOURCES ObjectStoreTest.cc)
add_test_case(memory_allocator_test SOURCES MemoryAllocatorTest.cc)
add_test_case(ffor_codec_test SOURCES FForCodecTest.cc)
add_test_case(print_config_test SOURCES PrintConfigTest.cc)
+add_test_case(payload_test SOURCES PayloadTest.cc)
diff --git a/cpp/core/tests/PayloadTest.cc b/cpp/core/tests/PayloadTest.cc
new file mode 100644
index 0000000000..defbfd9c98
--- /dev/null
+++ b/cpp/core/tests/PayloadTest.cc
@@ -0,0 +1,163 @@
+/*
+ * 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 "shuffle/Payload.h"
+#include "shuffle/Utils.h"
+#include "utils/Compression.h"
+
+#include <arrow/buffer.h>
+#include <arrow/io/memory.h>
+#include <gtest/gtest.h>
+
+#include <cstring>
+
+namespace gluten {
+namespace {
+
+// 4096 rows: the validity buffer is 512 bytes, a multiple of 64, so Arrow's
+// capacity rounding alone leaves no spare bytes after it.
+constexpr uint32_t kNumRows = 4096;
+const std::vector<bool> kIsValidityBuffer = {true, false};
+
+std::shared_ptr<arrow::Buffer> makeBuffer(int64_t size, uint8_t seed) {
+ auto buffer = arrow::AllocateBuffer(size).ValueOrDie();
+ for (int64_t i = 0; i < size; ++i) {
+ buffer->mutable_data()[i] = static_cast<uint8_t>(seed + i * 7);
+ }
+ return std::shared_ptr<arrow::Buffer>(std::move(buffer));
+}
+
+std::vector<std::shared_ptr<arrow::Buffer>> makeBuffers(uint8_t seed) {
+ return {makeBuffer(kNumRows / 8, seed), makeBuffer(kNumRows *
sizeof(int64_t), seed + 1)};
+}
+
+void expectZeroPadding(const arrow::Buffer& buffer) {
+ for (int64_t i = buffer.size(); i < buffer.size() + kSimdPadding; ++i) {
+ ASSERT_EQ(buffer.data()[i], 0) << "padding byte " << i << " of buffer with
size " << buffer.size();
+ }
+}
+
+void expectPadded(const std::vector<std::shared_ptr<arrow::Buffer>>& buffers) {
+ for (const auto& buffer : buffers) {
+ ASSERT_NE(buffer, nullptr);
+ EXPECT_GE(buffer->capacity() - buffer->size(), kSimdPadding)
+ << "size " << buffer->size() << ", capacity " << buffer->capacity();
+ expectZeroPadding(*buffer);
+ }
+}
+
+void expectSameData(const arrow::Buffer& actual, const arrow::Buffer&
expected) {
+ ASSERT_EQ(actual.size(), expected.size());
+ EXPECT_EQ(std::memcmp(actual.data(), expected.data(), expected.size()), 0);
+}
+
+// Serializes and deserializes 'buffers' as a block payload.
+std::vector<std::shared_ptr<arrow::Buffer>> roundTrip(
+ const std::vector<std::shared_ptr<arrow::Buffer>>& buffers,
+ Payload::Type type,
+ const std::shared_ptr<arrow::util::Codec>& codec) {
+ auto* pool = arrow::default_memory_pool();
+ auto payload = BlockPayload::fromBuffers(type, kNumRows, buffers,
&kIsValidityBuffer, pool, codec.get()).ValueOrDie();
+ auto out = arrow::io::BufferOutputStream::Create(1024, pool).ValueOrDie();
+ EXPECT_TRUE(payload->serialize(out.get()).ok());
+ auto in =
std::make_shared<arrow::io::BufferReader>(out->Finish().ValueOrDie());
+
+ uint32_t numRows = 0;
+ int64_t deserializeTime = 0;
+ int64_t decompressTime = 0;
+ auto result = BlockPayload::deserialize(in.get(), codec, pool, numRows,
deserializeTime, decompressTime).ValueOrDie();
+ EXPECT_EQ(numRows, kNumRows);
+ return result;
+}
+
+} // namespace
+
+TEST(PayloadTest, allocatePaddedBuffer) {
+ for (int64_t size : {0, 1, 63, 64, 512, 4096}) {
+ auto buffer = allocatePaddedBuffer(size,
arrow::default_memory_pool()).ValueOrDie();
+ EXPECT_EQ(buffer->size(), size);
+ EXPECT_GE(buffer->capacity() - buffer->size(), kSimdPadding);
+ expectZeroPadding(*buffer);
+ }
+}
+
+TEST(PayloadTest, resizeWithPaddingKeepsPadding) {
+ auto buffer = allocatePaddedBuffer(512,
arrow::default_memory_pool()).ValueOrDie();
+ std::memset(buffer->mutable_data(), 0xab, 512);
+ for (int64_t size : {520, 1024, 100}) {
+ ASSERT_TRUE(resizeWithPadding(*buffer, size).ok());
+ EXPECT_EQ(buffer->size(), size);
+ EXPECT_GE(buffer->capacity() - buffer->size(), kSimdPadding);
+ expectZeroPadding(*buffer);
+ }
+ // The leading bytes survive reallocation.
+ EXPECT_EQ(buffer->data()[0], 0xab);
+ EXPECT_EQ(buffer->data()[99], 0xab);
+}
+
+TEST(PayloadTest, deserializedUncompressedBuffersArePadded) {
+ auto buffers = makeBuffers(1);
+ auto result = roundTrip(buffers, Payload::kUncompressed, nullptr);
+ ASSERT_EQ(result.size(), buffers.size());
+ expectPadded(result);
+ for (size_t i = 0; i < buffers.size(); ++i) {
+ expectSameData(*result[i], *buffers[i]);
+ }
+}
+
+TEST(PayloadTest, deserializedCompressedBuffersArePadded) {
+ std::shared_ptr<arrow::util::Codec> codec =
createCompressionCodec(arrow::Compression::LZ4_FRAME, CodecBackend::NONE);
+ ASSERT_NE(codec, nullptr);
+ auto buffers = makeBuffers(2);
+ auto result = roundTrip(buffers, Payload::kCompressed, codec);
+ ASSERT_EQ(result.size(), buffers.size());
+ expectPadded(result);
+ for (size_t i = 0; i < buffers.size(); ++i) {
+ expectSameData(*result[i], *buffers[i]);
+ }
+}
+
+TEST(PayloadTest, mergedBuffersArePadded) {
+ auto* pool = arrow::default_memory_pool();
+ auto sourceBuffers = roundTrip(makeBuffers(3), Payload::kUncompressed,
nullptr);
+ auto appendBuffers = roundTrip(makeBuffers(4), Payload::kUncompressed,
nullptr);
+ auto source = std::make_unique<InMemoryPayload>(kNumRows,
&kIsValidityBuffer, nullptr, sourceBuffers);
+ auto append = std::make_unique<InMemoryPayload>(kNumRows,
&kIsValidityBuffer, nullptr, appendBuffers);
+
+ auto merged = InMemoryPayload::merge(std::move(source), std::move(append),
pool).ValueOrDie();
+ ASSERT_EQ(merged->numRows(), 2 * kNumRows);
+ std::vector<std::shared_ptr<arrow::Buffer>> mergedBuffers;
+ for (uint32_t i = 0; i < merged->numBuffers(); ++i) {
+ mergedBuffers.push_back(merged->readBufferAt(i).ValueOrDie());
+ }
+ expectPadded(mergedBuffers);
+ // Merged validity is 1024 bytes and values are 64 KiB, both multiples of 64.
+ EXPECT_EQ(mergedBuffers[0]->size(), 2 * kNumRows / 8);
+ EXPECT_EQ(mergedBuffers[1]->size(), 2 * kNumRows * sizeof(int64_t));
+}
+
+TEST(PayloadTest, copiedBuffersArePadded) {
+ auto payload = std::make_unique<InMemoryPayload>(kNumRows,
&kIsValidityBuffer, nullptr, makeBuffers(5));
+ ASSERT_TRUE(payload->copyBuffers(arrow::default_memory_pool()).ok());
+ std::vector<std::shared_ptr<arrow::Buffer>> buffers;
+ for (uint32_t i = 0; i < payload->numBuffers(); ++i) {
+ buffers.push_back(payload->readBufferAt(i).ValueOrDie());
+ }
+ expectPadded(buffers);
+}
+
+} // namespace gluten
diff --git a/cpp/velox/shuffle/VeloxShuffleReader.cc
b/cpp/velox/shuffle/VeloxShuffleReader.cc
index 478d619773..3baffa7941 100644
--- a/cpp/velox/shuffle/VeloxShuffleReader.cc
+++ b/cpp/velox/shuffle/VeloxShuffleReader.cc
@@ -28,6 +28,8 @@
#include "utils/Timer.h"
#include "utils/VeloxArrowUtils.h"
+#include "velox/buffer/Buffer.h"
+#include "velox/common/base/SimdUtil.h"
#include "velox/common/memory/ByteStream.h"
#include "velox/row/CompactRow.h"
#include "velox/serializers/PrestoHeader.h"
@@ -39,6 +41,7 @@
#include <algorithm>
#include <array>
+#include <cstring>
#include <sstream>
#include "VeloxGpuAsyncShuffleReader.h"
@@ -88,10 +91,20 @@ BufferPtr wrapInBufferViewAsOwner(const void* buffer,
size_t length, std::shared
static_cast<const uint8_t*>(buffer), length,
{std::move(bufferReleaser)});
}
-BufferPtr convertToVeloxBuffer(std::shared_ptr<arrow::Buffer> buffer) {
+static_assert(kSimdPadding >= simd::kPadding, "Shuffle buffers must keep
Velox's SIMD padding");
+
+BufferPtr convertToVeloxBuffer(std::shared_ptr<arrow::Buffer> buffer,
memory::MemoryPool* pool) {
if (buffer == nullptr) {
return nullptr;
}
+ // Velox may read up to simd::kPadding bytes past the end with SIMD loads,
e.g. when
+ // gathering null bits. Wrap the buffer without a copy only if that is
addressable.
+ // Payload deserialization allocates with the padding, so the copy is a
fallback.
+ if (buffer->size() > 0 && buffer->capacity() - buffer->size() <
simd::kPadding) {
+ auto copy = AlignedBuffer::allocate<char>(buffer->size(), pool);
+ std::memcpy(copy->asMutable<char>(), buffer->data(), buffer->size());
+ return copy;
+ }
return wrapInBufferViewAsOwner(buffer->data(), buffer->size(), buffer);
}
@@ -308,7 +321,7 @@ std::shared_ptr<VeloxColumnarBatch> makeColumnarBatch(
std::vector<BufferPtr> veloxBuffers;
veloxBuffers.reserve(arrowBuffers.size());
for (auto& buffer : arrowBuffers) {
- veloxBuffers.push_back(convertToVeloxBuffer(std::move(buffer)));
+ veloxBuffers.push_back(convertToVeloxBuffer(std::move(buffer), pool));
}
auto rowVector = deserialize(type, numRows, veloxBuffers, dictionaryFields,
dictionaries, pool);
return std::make_shared<VeloxColumnarBatch>(std::move(rowVector));
@@ -325,7 +338,7 @@ std::shared_ptr<VeloxColumnarBatch> makeColumnarBatch(
veloxBuffers.reserve(numBuffers);
for (size_t i = 0; i < numBuffers; ++i) {
GLUTEN_ASSIGN_OR_THROW(auto buffer, payload->readBufferAt(i));
- veloxBuffers.push_back(convertToVeloxBuffer(std::move(buffer)));
+ veloxBuffers.push_back(convertToVeloxBuffer(std::move(buffer), pool));
}
auto rowVector = deserialize(type, payload->numRows(), veloxBuffers, {}, {},
pool);
return std::make_shared<VeloxColumnarBatch>(std::move(rowVector));
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
