diff --git a/cpp/core/jni/JniWrapper.cc b/cpp/core/jni/JniWrapper.cc index 29036132a4c..a01b6a8d0a9 100644 --- a/cpp/core/jni/JniWrapper.cc +++ b/cpp/core/jni/JniWrapper.cc @@ -1048,7 +1048,7 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleWriterJniWrappe } ObjectStore::release(partitionWriterHandle); - auto shuffleWriterOptions = std::make_shared( + auto shuffleWriterOptions = std::make_shared( toPartitioning(jStringToCString(env, partitioningNameJstr)), startPartitionId, splitBufferSize, diff --git a/cpp/velox/CMakeLists.txt b/cpp/velox/CMakeLists.txt index b9cffc92fa0..636b073045d 100644 --- a/cpp/velox/CMakeLists.txt +++ b/cpp/velox/CMakeLists.txt @@ -227,7 +227,6 @@ if(ENABLE_GPU) operators/plannodes/CudfVectorStream.cc shuffle/VeloxGpuAsyncShuffleReader.cc shuffle/VeloxGpuShuffleReader.cc - shuffle/VeloxGpuShuffleWriter.cc operators/serializer/VeloxGpuColumnarBatchSerializer.cc utils/GpuBufferBatchResizer.cc memory/GpuBufferColumnarBatch.cc) diff --git a/cpp/velox/memory/GpuBufferColumnarBatch.cc b/cpp/velox/memory/GpuBufferColumnarBatch.cc index b76a849c82e..eab90eb2f81 100644 --- a/cpp/velox/memory/GpuBufferColumnarBatch.cc +++ b/cpp/velox/memory/GpuBufferColumnarBatch.cc @@ -28,7 +28,7 @@ namespace gluten { namespace { -enum class BufferType { kNull, kLength, kValue }; +enum class BufferType { kNull, kLength, kValue, kBooleanValue, kTimestampValue }; } using namespace facebook; @@ -67,16 +67,20 @@ std::shared_ptr GpuBufferColumnarBatch::compose( GLUTEN_CHECK(!batches.empty(), "No batches to compose"); // Compute the returned GpuBufferColumnarBatch buffers. auto& type = batches[0]->getRowType(); - const auto bufferSize = batches[0]->buffers().size(); + const auto numBuffers = batches[0]->buffers().size(); std::vector bufferSizes; - bufferSizes.resize(bufferSize); + bufferSizes.resize(numBuffers); std::vector bufferTypes; - bufferTypes.reserve(bufferSize); + bufferTypes.reserve(numBuffers); for (const auto& colType : type->children()) { bufferSizes[bufferTypes.size()] = arrow::bit_util::BytesForBits(numRows); bufferTypes.push_back(BufferType::kNull); - if (colType->isFixedWidth()) { + if (colType->isTimestamp()) { + bufferTypes.push_back(BufferType::kTimestampValue); + } else if (colType->isBoolean()) { + bufferTypes.push_back(BufferType::kBooleanValue); + } else if (colType->isFixedWidth()) { bufferTypes.push_back(BufferType::kValue); } else { // Add the first offset 0. @@ -86,27 +90,37 @@ std::shared_ptr GpuBufferColumnarBatch::compose( bufferTypes.push_back(BufferType::kValue); } } - VELOX_CHECK_EQ(bufferTypes.size(), bufferSize); + VELOX_CHECK_EQ(bufferTypes.size(), numBuffers); // This buffer may be more than the actual reauired buffer for null buffer. for (const auto& batch : batches) { if (batch->numRows() == 0) { continue; } - for (auto i = 0; i < bufferSize; ++i) { + for (auto i = 0; i < numBuffers; ++i) { // The null buffer may be null or length = 0. // Maybe optimize later, detect if the null buffer is all true. And set the return null buffer to 0. if (bufferTypes[i] == BufferType::kNull || bufferTypes[i] == BufferType::kLength) { continue; } + auto& buffer = batch->bufferAt(i); VELOX_CHECK_NOT_NULL(buffer); - bufferSizes[i] += buffer->size(); + + if (bufferTypes[i] == BufferType::kTimestampValue) { + // Velox Timestamp value is 16 bytes. CUDF Timestamp value is 8 bytes. + bufferSizes[i] += static_cast(batch->numRows()) * sizeof(int64_t); + } else if (bufferTypes[i] == BufferType::kBooleanValue) { + // Velox Boolean value is 1 bit. CUDF Boolean value is 1 byte. + bufferSizes[i] += batch->numRows(); + } else { + bufferSizes[i] += buffer->size(); + } } } std::vector> returnBuffers; - returnBuffers.reserve(bufferSize); - for (auto i = 0; i < bufferSize; ++i) { + returnBuffers.reserve(numBuffers); + for (auto i = 0; i < numBuffers; ++i) { // Defer the null buffer to really contains null. if (bufferTypes[i] == BufferType::kNull) { returnBuffers.emplace_back(nullptr); @@ -121,8 +135,8 @@ std::shared_ptr GpuBufferColumnarBatch::compose( int32_t bufferIdx = 0; for (const auto& colType : type->children()) { size_t rowNumber = 0; - // Also records the value buffer offset. - size_t stringOffset = 0; + // Also records the length buffer offset for strings. + size_t valueBufferOffset = 0; for (auto i = 0; i < batches.size(); ++i) { const auto& batch = batches[i]; if (batch->numRows() == 0) { @@ -147,28 +161,47 @@ std::shared_ptr GpuBufferColumnarBatch::compose( arrow::internal::CopyBitmap(batch->bufferAt(bufferIdx)->data(), 0, batch->numRows(), dst, rowNumber); } - if (colType->isFixedWidth()) { + if (colType->isTimestamp()) { + const auto bufferSize = batch->numRows() * sizeof(int64_t); + VELOX_CHECK_LE(valueBufferOffset + bufferSize, returnBuffers[bufferIdx + 1]->size()); + const auto* src = reinterpret_cast(batch->bufferAt(bufferIdx + 1)->data()); + auto* dst = reinterpret_cast(returnBuffers[bufferIdx + 1]->mutable_data() + valueBufferOffset); + for (auto j = 0; j < batch->numRows(); ++j) { + // src[0] is seconds, src[1] is nanoseconds in Timestamp. + dst[j] = src[j << 1] * 1'000'000'000L + src[j << 1 | 1]; + } + valueBufferOffset += bufferSize; + } else if (colType->isBoolean()) { + const auto bufferSize = batch->numRows(); + VELOX_CHECK_LE(valueBufferOffset + bufferSize, returnBuffers[bufferIdx + 1]->size()); + const auto* src = batch->bufferAt(bufferIdx + 1)->data(); + auto* dst = returnBuffers[bufferIdx + 1]->mutable_data() + valueBufferOffset; + for (auto j = 0; j < batch->numRows(); ++j) { + dst[j] = (src[j >> 3] >> (j & 7)) & 1; + } + valueBufferOffset += bufferSize; + } else if (colType->isFixedWidth()) { // The buffer is values. const auto bufferSize = batch->bufferAt(bufferIdx + 1)->size(); - VELOX_CHECK_LE(stringOffset + bufferSize, returnBuffers[bufferIdx + 1]->size()); + VELOX_CHECK_LE(valueBufferOffset + bufferSize, returnBuffers[bufferIdx + 1]->size()); memcpy( - returnBuffers[bufferIdx + 1]->mutable_data() + stringOffset, + returnBuffers[bufferIdx + 1]->mutable_data() + valueBufferOffset, batch->bufferAt(bufferIdx + 1)->data(), bufferSize); - stringOffset += bufferSize; + valueBufferOffset += bufferSize; } else { // String, lengths, values memcpy( - returnBuffers[bufferIdx + 2]->mutable_data() + stringOffset, + returnBuffers[bufferIdx + 2]->mutable_data() + valueBufferOffset, batch->bufferAt(bufferIdx + 2)->data(), batch->bufferAt(bufferIdx + 2)->size()); const auto* lengths = reinterpret_cast(batch->bufferAt(bufferIdx + 1)->data()); auto* offsetBuffer = reinterpret_cast(returnBuffers[bufferIdx + 1]->mutable_data()); for (auto j = 0; j < batch->numRows(); ++j) { - offsetBuffer[rowNumber + j] = stringOffset; - stringOffset += lengths[j]; + offsetBuffer[rowNumber + j] = valueBufferOffset; + valueBufferOffset += lengths[j]; } - offsetBuffer[rowNumber + batch->numRows()] = stringOffset; + offsetBuffer[rowNumber + batch->numRows()] = valueBufferOffset; } rowNumber += batch->numRows(); } diff --git a/cpp/velox/shuffle/VeloxGpuShuffleWriter.cc b/cpp/velox/shuffle/VeloxGpuShuffleWriter.cc deleted file mode 100644 index 1df6cb91ad8..00000000000 --- a/cpp/velox/shuffle/VeloxGpuShuffleWriter.cc +++ /dev/null @@ -1,72 +0,0 @@ -/* - * 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 "VeloxGpuShuffleWriter.h" - -namespace gluten { - -// Split bool bit to bytes. -void VeloxGpuHashShuffleWriter::splitBoolValueType(const uint8_t* srcAddr, const std::vector& dstAddrs) { - for (auto& pid : partitionUsed_) { - auto dstaddr = dstAddrs[pid]; - if (dstaddr == nullptr) { - continue; - } - auto dstPidBase = reinterpret_cast(dstaddr + partitionBufferBase_[pid] * sizeof(uint8_t)); - auto pos = partition2RowOffsetBase_[pid]; - auto end = partition2RowOffsetBase_[pid + 1]; - for (; pos < end; ++pos) { - auto rowId = rowOffset2RowId_[pos]; - // === Extract one bool value from the bit-packed source === - uint8_t byte = srcAddr[rowId >> 3]; // find the source byte - uint8_t bit = (byte >> (rowId & 7)) & 0x01; // extract 0 or 1 - *dstPidBase++ = bit; // copy - } - } -} - -// Split timestamp from int128_t to int64_t, both of them represents the timestamp nanoseconds. -arrow::Status VeloxGpuHashShuffleWriter::splitTimestamp(const uint8_t* srcAddr, const std::vector& dstAddrs) { - for (auto& pid : partitionUsed_) { - auto dstPidBase = reinterpret_cast(dstAddrs[pid] + partitionBufferBase_[pid] * sizeof(int64_t)); - auto pos = partition2RowOffsetBase_[pid]; - auto end = partition2RowOffsetBase_[pid + 1]; - for (; pos < end; ++pos) { - auto rowId = rowOffset2RowId_[pos]; - auto* src = reinterpret_cast(srcAddr) + rowId * 2; - // src[0] is seconds, src[1] is nanoseconds in Timestamp. - *dstPidBase++ = src[0] * 1'000'000'000LL + src[1]; - } - } - return arrow::Status::OK(); -} - -arrow::Result> VeloxGpuHashShuffleWriter::create( - uint32_t numPartitions, - const std::shared_ptr& partitionWriter, - const std::shared_ptr& options, - MemoryManager* memoryManager) { - if (auto hashOptions = std::dynamic_pointer_cast(options)) { - std::shared_ptr res = - std::make_shared(numPartitions, partitionWriter, hashOptions, memoryManager); - RETURN_NOT_OK(res->init()); - return res; - } - return arrow::Status::Invalid("Error casting ShuffleWriterOptions to GpuHashShuffleWriterOptions. "); -} - -} // namespace gluten diff --git a/cpp/velox/shuffle/VeloxGpuShuffleWriter.h b/cpp/velox/shuffle/VeloxGpuShuffleWriter.h deleted file mode 100644 index e6f0e1a3586..00000000000 --- a/cpp/velox/shuffle/VeloxGpuShuffleWriter.h +++ /dev/null @@ -1,57 +0,0 @@ -/* - * 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. - */ - -#pragma once - -#include "VeloxHashShuffleWriter.h" - -namespace gluten { - -class VeloxGpuHashShuffleWriter : public VeloxHashShuffleWriter { - public: - static arrow::Result> create( - uint32_t numPartitions, - const std::shared_ptr& partitionWriter, - const std::shared_ptr& options, - MemoryManager* memoryManager); - - VeloxGpuHashShuffleWriter( - uint32_t numPartitions, - const std::shared_ptr& partitionWriter, - const std::shared_ptr& options, - MemoryManager* memoryManager) - : VeloxHashShuffleWriter(numPartitions, partitionWriter, options, memoryManager) {} - - private: - // Split the bool to byte. - void splitBoolValueType(const uint8_t* srcAddr, const std::vector& dstAddrs) override; - - uint64_t valueBufferSizeForBool(uint32_t newSize) override { - return newSize; - } - - bool boolIsBit() override { - return false; - } - - arrow::Status splitTimestamp(const uint8_t* srcAddr, const std::vector& dstAddrs) override; - - uint64_t valueBufferSizeForTimestamp(uint32_t newSize) override { - return sizeof(int64_t) * newSize; - } -}; -} // namespace gluten diff --git a/cpp/velox/shuffle/VeloxShuffleWriter.cc b/cpp/velox/shuffle/VeloxShuffleWriter.cc index 25ed0990141..a83e3a19471 100644 --- a/cpp/velox/shuffle/VeloxShuffleWriter.cc +++ b/cpp/velox/shuffle/VeloxShuffleWriter.cc @@ -20,10 +20,6 @@ #include "shuffle/VeloxRssSortShuffleWriter.h" #include "shuffle/VeloxSortShuffleWriter.h" -#ifdef GLUTEN_ENABLE_GPU -#include "VeloxGpuShuffleWriter.h" -#endif - namespace gluten { arrow::Result> VeloxShuffleWriter::create( @@ -35,15 +31,12 @@ arrow::Result> VeloxShuffleWriter::create( std::shared_ptr shuffleWriter; switch (type) { case ShuffleWriterType::kHashShuffle: + case ShuffleWriterType::kGpuHashShuffle: return VeloxHashShuffleWriter::create(numPartitions, std::move(partitionWriter), options, memoryManager); case ShuffleWriterType::kSortShuffle: return VeloxSortShuffleWriter::create(numPartitions, std::move(partitionWriter), options, memoryManager); case ShuffleWriterType::kRssSortShuffle: return VeloxRssSortShuffleWriter::create(numPartitions, std::move(partitionWriter), options, memoryManager); -#ifdef GLUTEN_ENABLE_GPU - case ShuffleWriterType::kGpuHashShuffle: - return VeloxGpuHashShuffleWriter::create(numPartitions, std::move(partitionWriter), options, memoryManager); -#endif default: return arrow::Status::Invalid("Unsupported shuffle writer type: ", typeToString(type)); } diff --git a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc index cc846eb4eb6..896fa9e6265 100644 --- a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc +++ b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc @@ -20,8 +20,7 @@ #include "config/GlutenConfig.h" #include "memory/GpuBufferColumnarBatch.h" -#include "shuffle/VeloxGpuShuffleWriter.h" -#include "shuffle/VeloxHashShuffleWriter.h" +#include "shuffle/VeloxShuffleWriter.h" #include "tests/VeloxShuffleWriterTestBase.h" #include "tests/utils/TestAllocationListener.h" #include "tests/utils/TestStreamReader.h" @@ -227,7 +226,7 @@ class GpuVeloxShuffleWriterTest : public ::testing::TestWithParam(); + auto hashOptions = std::make_shared(); hashOptions->splitBufferSize = splitBufferSize; options = hashOptions; } break;