From 10fdbd6c037cb5c04085ce8ab0bec80fbf3c326f Mon Sep 17 00:00:00 2001 From: zetongzhao Date: Thu, 14 Oct 2021 14:25:19 -0400 Subject: [PATCH] resize queue and test --- .../ccsrc/minddata/dataset/engine/connector.h | 2 +- mindspore/ccsrc/minddata/dataset/util/queue.h | 77 +++++++++++++------ tests/ut/cpp/dataset/queue_test.cc | 48 ++++++++++++ 3 files changed, 103 insertions(+), 24 deletions(-) diff --git a/mindspore/ccsrc/minddata/dataset/engine/connector.h b/mindspore/ccsrc/minddata/dataset/engine/connector.h index 98c2333e202..f966f251381 100644 --- a/mindspore/ccsrc/minddata/dataset/engine/connector.h +++ b/mindspore/ccsrc/minddata/dataset/engine/connector.h @@ -138,7 +138,7 @@ class Connector { // starting from the beginning. void Reset() { for (int i = 0; i < queues_.size(); ++i) { - queues_[i]->ResetQue(); + queues_[i]->Reset(); } expect_consumer_ = 0; pop_from_ = 0; diff --git a/mindspore/ccsrc/minddata/dataset/util/queue.h b/mindspore/ccsrc/minddata/dataset/util/queue.h index 3a62c246264..c7cd414bfab 100644 --- a/mindspore/ccsrc/minddata/dataset/util/queue.h +++ b/mindspore/ccsrc/minddata/dataset/util/queue.h @@ -66,7 +66,10 @@ class Queue { bool empty() const { return head_ == tail_; } - void Reset() { ResetQue(); } + void Reset() { + std::unique_lock _lock(mux_); + ResetQue(); + } // Producer Status Add(const_reference ele) noexcept { @@ -74,8 +77,7 @@ class Queue { // Block when full Status rc = full_cv_.Wait(&_lock, [this]() -> bool { return (size() != capacity()); }); if (rc.IsOk()) { - auto k = tail_++ % sz_; - *(arr_[k]) = ele; + RETURN_IF_NOT_OK(AddWhileHoldingLock(ele)); empty_cv_.NotifyAll(); _lock.unlock(); } else { @@ -121,8 +123,7 @@ class Queue { // Block when empty Status rc = empty_cv_.Wait(&_lock, [this]() -> bool { return !empty(); }); if (rc.IsOk()) { - auto k = head_++ % sz_; - *p = std::move(*(arr_[k])); + RETURN_IF_NOT_OK(PopFrontWhileHoldingLock(p)); full_cv_.NotifyAll(); _lock.unlock(); } else { @@ -131,24 +132,6 @@ class Queue { return rc; } - void ResetQue() noexcept { - std::unique_lock _lock(mux_); - // If there are elements in the queue, drain them. We won't call PopFront directly - // because we have got the lock already. We will deadlock if we call PopFront - for (auto i = head_; i < tail_; ++i) { - auto k = i % sz_; - auto val = std::move(*(arr_[k])); - // Let val go out of scope and its destructor will be invoked automatically. - // But our compiler may complain val is not in use. So let's do some useless - // stuff. - MS_LOG(DEBUG) << "Address of val: " << &val; - } - empty_cv_.ResetIntrpState(); - full_cv_.ResetIntrpState(); - head_ = 0; - tail_ = 0; - } - Status Register(TaskGroup *vg) { Status rc1 = empty_cv_.Register(vg->GetIntrpService()); Status rc2 = full_cv_.Register(vg->GetIntrpService()); @@ -159,6 +142,28 @@ class Queue { } } + Status Resize(int32_t new_capacity) { + std::unique_lock _lock(mux_); + CHECK_FAIL_RETURN_UNEXPECTED( + new_capacity >= static_cast(size()), + "New capacity: " + std::to_string(new_capacity) + ", is smaller than queue size:" + std::to_string(size())); + std::vector queue; + while (head_ < tail_) { + T temp; + RETURN_IF_NOT_OK(this->PopFrontWhileHoldingLock(&temp)); + queue.push_back(temp); + } + this->ResetQue(); + RETURN_IF_NOT_OK(arr_.allocate(new_capacity)); + sz_ = new_capacity; + for (int i = 0; i < queue.size(); ++i) { + RETURN_IF_NOT_OK(this->AddWhileHoldingLock(queue[i])); + } + queue.clear(); + _lock.unlock(); + return Status::OK(); + } + private: size_t sz_; MemGuard> arr_; @@ -168,6 +173,32 @@ class Queue { std::mutex mux_; CondVar empty_cv_; CondVar full_cv_; + + // Helper function for Add, must be called when holding a lock + Status AddWhileHoldingLock(const_reference ele) { + auto k = tail_++ % sz_; + *(arr_[k]) = ele; + return Status::OK(); + } + + // Helper function for PopFront, must be called when holding a lock + Status PopFrontWhileHoldingLock(pointer p) { + auto k = head_++ % sz_; + *p = std::move(*(arr_[k])); + return Status::OK(); + } + + void ResetQue() noexcept { + while (head_ < tail_) { + T val; + this->PopFrontWhileHoldingLock(&val); + MS_LOG(DEBUG) << "Address of val: " << &val; + } + empty_cv_.ResetIntrpState(); + full_cv_.ResetIntrpState(); + head_ = 0; + tail_ = 0; + } }; // A container of queues with [] operator accessors. Basically this is a wrapper over of a vector of queues diff --git a/tests/ut/cpp/dataset/queue_test.cc b/tests/ut/cpp/dataset/queue_test.cc index 3f2e1ad4af2..93aeaff09af 100644 --- a/tests/ut/cpp/dataset/queue_test.cc +++ b/tests/ut/cpp/dataset/queue_test.cc @@ -177,3 +177,51 @@ TEST_F(MindDataTestQueue, Test6) { MS_LOG(INFO) << "Popped value " << *pepped_value << " from queue index " << chosen_queue_index; ASSERT_EQ(*pepped_value, 99); } + +// Feature: Check resize is finished without changing elements and influencing operations. +// Description: Compare elements in queue before and after resize, and test add/pop/reset. +// Expectation: Elements in queue after resize are the same as the original queue. +TEST_F(MindDataTestQueue, TestResize) { + // Create a list of queues with capacity = 3 + Queue queue(3); + ASSERT_EQ(3, queue.capacity()); + // Add 3 rows into the queue + TensorRow a; + std::shared_ptr test_tensor1; + std::vector input = {1.1, 0.2, 0.3, 0.4, 0.5, 0.6, 1.2, 0.7, 0.8, 0.9, 1.0, 2.0, 1.3, 3.0, 4.0}; + ASSERT_OK(Tensor::CreateFromVector(input, TensorShape{3, 5}, &test_tensor1)); + a.push_back(test_tensor1); + EXPECT_OK(queue.Add(a)); + + TensorRow b; + std::shared_ptr test_tensor2; + ASSERT_OK(Tensor::CreateScalar(true, &test_tensor2)); + b.push_back(test_tensor2); + EXPECT_OK(queue.Add(b)); + + TensorRow c; + std::shared_ptr test_tensor3; + ASSERT_OK(Tensor::CreateFromVector(input, &test_tensor3)); + c.push_back(test_tensor3); + EXPECT_OK(queue.Add(c)); + ASSERT_EQ(3, queue.size()); + // Check false if the resize is smaller than current size + EXPECT_ERROR(queue.Resize(2)); + // Check true if the resize is larger than current size, and capacity is changed + EXPECT_OK(queue.Resize(12)); + ASSERT_EQ(12, queue.capacity()); + TensorRow d = a; + EXPECT_OK(queue.Add(d)); + ASSERT_EQ(4, queue.size()); + // Expect the rows after resize are the same as original input + TensorRow e; + EXPECT_OK(queue.PopFront(&e)); + EXPECT_EQ(a.getRow(), e.getRow()); + EXPECT_OK(queue.PopFront(&e)); + EXPECT_EQ(b.getRow(), e.getRow()); + EXPECT_OK(queue.PopFront(&e)); + EXPECT_EQ(c.getRow(), e.getRow()); + ASSERT_EQ(1, queue.size()); + queue.Reset(); + ASSERT_EQ(0, queue.size()); +} \ No newline at end of file