diff --git a/velox/common/base/Scratch.h b/velox/common/base/Scratch.h new file mode 100644 index 000000000000..a2110e01aaac --- /dev/null +++ b/velox/common/base/Scratch.h @@ -0,0 +1,159 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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 "velox/common/base/RawVector.h" + +/// A utility for reusable scoped temporary scratch areas. +namespace facebook::velox { + +/// A collection of temporary reusable scratch vectors. The vectors are accessed +/// via the ScratchPtr scoped lease. The vectors are padded so that their last +/// element can be written at full SIMD width, as with raw_vector. +class Scratch { + public: + using Item = raw_vector; + + Scratch() = default; + Scratch(const Scratch& other) = delete; + + ~Scratch() { + reserve(0); + ::free(items_); + items_ = nullptr; + capacity_ = 0; + fill_ = 0; + } + void operator=(const Scratch& other) = delete; + + /// Returns the next reusable scratch vector or makes a new one. + Item get() { + if (fill_ == 0) { + return Item(); + } + auto temp = std::move(items_[fill_ - 1]); + --fill_; + retainedSize_ -= temp.capacity(); + return temp; + } + + void release(Item&& item) { + retainedSize_ += item.capacity(); + if (fill_ == capacity_) { + reserve(std::max(16, 2 * capacity_)); + } + items_[fill_++] = std::move(item); + } + + void trim() { + reserve(0); + retainedSize_ = 0; + } + + size_t retainedSize() { + return retainedSize_; + } + + private: + void reserve(int32_t newCapacity) { + VELOX_CHECK_LE(fill_, capacity_); + // Delete the items above the new capacity. + for (auto i = newCapacity; i < fill_; ++i) { + std::destroy_at(&items_[i]); + } + if (newCapacity > capacity_) { + Item* newItems = + reinterpret_cast(::malloc(sizeof(Item) * newCapacity)); + if (fill_ > 0) { + memcpy(newItems, items_, fill_ * sizeof(Item)); + } + memset(newItems + fill_, 0, (newCapacity - fill_) * sizeof(Item)); + free(items_); + items_ = newItems; + capacity_ = newCapacity; + } + fill_ = std::min(fill_, newCapacity); + } + + Item* items_{nullptr}; + int32_t fill_{0}; + int32_t capacity_{0}; + // The total size held. If too large from outlier use cases, 'this' should be + // trimmed. + int64_t retainedSize_{0}; +}; + +/// A scoped lease for a scratch area of T. For scratch areas <= +/// 'inlineSize' the scratch area is inlined, typically on stack, and +/// no allocation will ever take place. The inline storage is padded +/// with a trailer of simd::kPadding bytes to allow writing at full +/// SIMD width at the end of the area. +template +class ScratchPtr { + public: + ScratchPtr(Scratch& scratch) : scratch_(&scratch) {} + + ScratchPtr(const ScratchPtr& other) = delete; + ScratchPtr(ScratchPtr&& other) = delete; + + inline ~ScratchPtr() { + if (data_.data()) { + scratch_->release(std::move(data_)); + } + } + + void operator=(ScratchPtr&& other) = delete; + void operator=(const ScratchPtr& other) = delete; + + /// Returns a writable pointer to at least 'size' uninitialized + /// elements of T. The last element is followed by simd::kPadding + /// bytes to allow a full width SIMD store for any element. This may + /// be called once per lifetime. + T* get(int32_t size) { + VELOX_CHECK_NULL(ptr_); + size_ = size; + if (size <= inlineSize) { + ptr_ = inline_; + return ptr_; + } + data_ = scratch_->get(); + data_.resize(size * sizeof(T)); + ptr_ = reinterpret_cast(data_.data()); + return ptr_; + } + + /// Returns the pointer returned by a previous get(int32_t). + T* get() const { + VELOX_DCHECK_NOT_NULL(ptr_); + return ptr_; + } + + /// Returns the size of the previous get(int32_t). + int32_t size() const { + return size_; + } + + private: + Scratch* scratch_{nullptr}; + raw_vector data_; + T* ptr_{nullptr}; + int32_t size_{0}; + T inline_[inlineSize]; + char padding_[inlineSize == 0 ? 0 : simd::kPadding]; +}; + +} // namespace facebook::velox diff --git a/velox/common/base/tests/CMakeLists.txt b/velox/common/base/tests/CMakeLists.txt index ff97bbe63711..5b305719e1b4 100644 --- a/velox/common/base/tests/CMakeLists.txt +++ b/velox/common/base/tests/CMakeLists.txt @@ -22,6 +22,7 @@ add_executable( FsTest.cpp RangeTest.cpp RawVectorTest.cpp + ScratchTest.cpp RuntimeMetricsTest.cpp ScopedLockTest.cpp SemaphoreTest.cpp diff --git a/velox/common/base/tests/ScratchTest.cpp b/velox/common/base/tests/ScratchTest.cpp new file mode 100644 index 000000000000..36bdd2da2f81 --- /dev/null +++ b/velox/common/base/tests/ScratchTest.cpp @@ -0,0 +1,74 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed 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 "velox/common/base/Scratch.h" + +#include + +using namespace facebook::velox; + +TEST(ScratchTest, basic) { + Scratch scratch; + { + ScratchPtr ints(scratch); + ScratchPtr longs(scratch); + auto tempInts = ints.get(1000); + auto tempLongs = longs.get(2000); + std::fill(tempInts, tempInts + 1000, -1); + std::fill(tempLongs, tempLongs + 2000, -1); + EXPECT_EQ(0, scratch.retainedSize()); + } + EXPECT_EQ(20352, scratch.retainedSize()); + { + ScratchPtr ints(scratch); + ScratchPtr longs(scratch); + auto tempLongs = longs.get(2000); + auto tempInts = ints.get(1000); + std::fill(tempInts, tempInts + 1000, -1); + std::fill(tempInts, tempInts + 2000, -1); + EXPECT_EQ(0, scratch.retainedSize()); + } + // The scratch vectors were acquired in a different order, so the smaller got + // resized to the larger size. + EXPECT_EQ(32640, scratch.retainedSize()); + scratch.trim(); + EXPECT_EQ(0, scratch.retainedSize()); + { + ScratchPtr ints(scratch); + // The size is the inline size, nothing gets returned to 'scratch'. + auto temp = ints.get(10); + temp[0] = 1; + } + EXPECT_EQ(0, scratch.retainedSize()); +} + +TEST(ScratchTest, large) { + constexpr int32_t kSize = 100; + Scratch scratch; + std::vector>> pointers; + for (auto i = 0; i < kSize; ++i) { + pointers.push_back(std::make_unique>(scratch)); + pointers.back()->get(1000); + } + pointers.clear(); + // 100 times 1000 bytes returned. + EXPECT_LT(100'000, scratch.retainedSize()); + for (auto i = 0; i < kSize; ++i) { + pointers.push_back(std::make_unique>(scratch)); + pointers.back()->get(1000); + } + EXPECT_EQ(0, scratch.retainedSize()); +}