[K/N] Non-worker work source in ParalleProcessor

To be used as thread-local mark queue in concurrent mark.


Merge-request: KT-MR-12723
Merged-by: Alexey Glushko <aleksei.glushko@jetbrains.com>

kotlin-native/runtime/src/gc/cms/cpp/ConcurrentMarkAndSweep.cpp

@@ -83,10 +83,6 @@ void gc::ConcurrentMarkAndSweep::ThreadData::clearMarkFlags() {
     rootSetLocked_.store(false, std::memory_order_release);
 }
 
-mm::ThreadData& gc::ConcurrentMarkAndSweep::ThreadData::commonThreadData() const {
-    return threadData_;
-}
-
 gc::ConcurrentMarkAndSweep::ConcurrentMarkAndSweep(
         alloc::Allocator& allocator, gcScheduler::GCScheduler& gcScheduler, bool mutatorsCooperate, std::size_t auxGCThreads) noexcept :
     allocator_(allocator),

kotlin-native/runtime/src/gc/cms/cpp/ConcurrentMarkAndSweep.hpp

@@ -43,20 +43,22 @@ public:
 
         void onThreadRegistration() noexcept { barriers_.onThreadRegistration(); }
 
-        BarriersThreadData& barriers() noexcept { return barriers_; }
-
         bool tryLockRootSet();
         void publish();
         bool published() const;
         void clearMarkFlags();
 
-        mm::ThreadData& commonThreadData() const;
+        auto& commonThreadData() const noexcept { return threadData_; }
+        auto& barriers() noexcept { return barriers_; }
+        // TODO use in concurrent mark
+        [[maybe_unused]] auto& markQueue() noexcept { return markQueue_; }
 
     private:
         friend ConcurrentMarkAndSweep;
         ConcurrentMarkAndSweep& gc_;
         mm::ThreadData& threadData_;
         BarriersThreadData barriers_;
+        ManuallyScoped<mark::ParallelMark::MutatorQueue> markQueue_;
 
         std::atomic<bool> rootSetLocked_ = false;
         std::atomic<bool> published_ = false;

kotlin-native/runtime/src/gc/cms/cpp/ParallelMark.hpp

@@ -78,6 +78,9 @@ class ParallelMark : private Pinned {
     // work balancing parameters were chosen pretty arbitrary
     using ParallelProcessor = ParallelProcessor<MarkStackImpl, 512, 4096>;
 public:
+
+    using MutatorQueue = ParallelProcessor::WorkSource;
+
     class MarkTraits {
     public:
         using MarkQueue = ParallelProcessor::Worker;

kotlin-native/runtime/src/main/cpp/BoundedQueue.hpp

@@ -35,14 +35,12 @@ namespace kotlin {
 
 /**
  * A fixed-size concurrent multi-producer/multi-consumer queue.
- * @tparam kCapacity must be a power of 2.
+ * @tparam kCapacity is suggested to set to a power of 2.
  */
 template<typename T, std::size_t kCapacity>
 class BoundedQueue : private Pinned {
 public:
     BoundedQueue() {
-        static_assert((kCapacity >= 2) && ((kCapacity & (kCapacity - 1)) == 0), "Queue capacity must be a power of 2");
-
         for (size_t i = 0; i < kCapacity; ++i) {
             buffer_[i].sequence_.store(i, std::memory_order_relaxed);
         }
@@ -54,7 +52,7 @@ public:
         Cell* cell;
         std::size_t pos = enqueuePos_.load(std::memory_order_relaxed);
         while (true) {
-            cell = &buffer_[pos & (kCapacity - 1)];
+            cell = &buffer_[pos % kCapacity];
             std::size_t seq = cell->sequence_.load(std::memory_order_acquire);
             std::intptr_t dif = static_cast<std::intptr_t>(seq) - static_cast<std::intptr_t>(pos);
             if (dif == 0) {
@@ -76,7 +74,7 @@ public:
         Cell* cell;
         std::size_t pos = dequeuePos_.load(std::memory_order_relaxed);
         while (true) {
-            cell = &buffer_[pos & (kCapacity - 1)];
+            cell = &buffer_[pos % kCapacity];
             std::size_t seq = cell->sequence_.load(std::memory_order_acquire);
             std::intptr_t dif = static_cast<std::intptr_t>(seq) - static_cast<std::intptr_t>(pos + 1);
             if (dif == 0) {

kotlin-native/runtime/src/main/cpp/BoundedQueueTest.cpp

@@ -88,6 +88,8 @@ TEST(BoundedQueueTest, ConcurrentDequeue) {
         });
     }
     start = true;
+
+    threads.clear();
 }
 
 TEST(BoundedQueueTest, PingPongWithOverflow) {

kotlin-native/runtime/src/main/cpp/ParallelProcessor.hpp

@@ -24,6 +24,7 @@ namespace kotlin {
  */
 template <typename ListImpl, std::size_t kBatchSize, std::size_t kBatchesPoolSize>
 class ParallelProcessor : private Pinned {
+private:
     class Batch {
     public:
         ALWAYS_INLINE bool empty() const noexcept {
@@ -71,25 +72,32 @@ class ParallelProcessor : private Pinned {
         std::size_t elemsCount_ = 0;
     };
 
-public:
-    class Worker : private Pinned {
-        friend ParallelProcessor;
+    class LocalQueue : private Pinned {
     public:
-        explicit Worker(ParallelProcessor& dispatcher) : dispatcher_(dispatcher) {
-            dispatcher_.registeredWorkers_.fetch_add(1, std::memory_order_relaxed);
-            RuntimeLogDebug({ kTagBalancing }, "Worker registered");
-        }
-
         ALWAYS_INLINE bool localEmpty() const noexcept {
-            return batch_.empty() && overflowList_.empty();
+            return localQueue_.empty();
         }
 
         ALWAYS_INLINE bool tryPushLocal(typename ListImpl::reference value) noexcept {
-            return overflowList_.try_push_front(value);
+            return localQueue_.try_push_front(value);
         }
 
         ALWAYS_INLINE typename ListImpl::pointer tryPopLocal() noexcept {
-            return overflowList_.try_pop_front();
+            return localQueue_.try_pop_front();
+        }
+
+    protected:
+        ListImpl localQueue_;
+    };
+
+public:
+    class WorkSource : public LocalQueue {
+        friend ParallelProcessor;
+    public:
+        explicit WorkSource(ParallelProcessor& dispatcher) : dispatcher_(dispatcher) {}
+
+        ALWAYS_INLINE bool retainsNoWork() const noexcept {
+            return batch_.empty() && this->localEmpty();
         }
 
         ALWAYS_INLINE bool tryPush(typename ListImpl::reference value) noexcept {
@@ -97,75 +105,113 @@ public:
                 bool released = dispatcher_.releaseBatch(std::move(batch_));
                 if (!released) {
                     RuntimeLogDebug({ kTagBalancing }, "Batches pool overflow");
-                    batch_.transferAllInto(overflowList_);
+                    batch_.transferAllInto(overflowList());
                 }
                 batch_ = Batch{};
             }
             return batch_.tryPush(value);
         }
 
+        /**
+         * Tries to transfer all the tasks stored in this WorkSource locally into the shared ParallelProcessor's storage.
+         * @return `true` iff this WorkSource doesn't contain any local tasks anymore.
+         */
+        ALWAYS_INLINE bool forceFlush() noexcept {
+            while (true) {
+                if (!batch_.empty()) {
+                    bool released = dispatcher_.releaseBatch(std::move(batch_));
+                    if (released) {
+                        RuntimeLogDebug({ kTagBalancing }, "Work butch flushed");
+                        batch_ = Batch{};
+                    } else {
+                        RuntimeLogDebug({ kTagBalancing }, "Failed to force flush work queue");
+                        return false;
+                    };
+                }
+                RuntimeAssert(batch_.empty(), "Now must be empty");
+                if (overflowList().empty()) {
+                    return true;
+                } else {
+                    RuntimeLogDebug({ kTagBalancing }, "Refiling batch from overflow list");
+                    batch_.fillFrom(overflowList());
+                }
+            }
+        }
+
+    protected:
+        ListImpl& overflowList() noexcept {
+            return this->localQueue_;
+        }
+
+        ParallelProcessor& dispatcher_;
+        Batch batch_;
+    };
+
+    class Worker : public WorkSource {
+        friend ParallelProcessor;
+    public:
+        explicit Worker(ParallelProcessor& dispatcher) : WorkSource(dispatcher) {
+            this->dispatcher_.registeredWorkers_.fetch_add(1, std::memory_order_relaxed);
+            RuntimeLogDebug({ kTagBalancing }, "Worker registered");
+        }
+
         ALWAYS_INLINE typename ListImpl::pointer tryPop() noexcept {
-            if (batch_.empty()) {
+            if (this->batch_.empty()) {
                 while (true) {
-                    bool acquired = dispatcher_.acquireBatch(batch_);
+                    bool acquired = this->dispatcher_.acquireBatch(this->batch_);
                     if (!acquired) {
-                        if (!overflowList_.empty()) {
-                            batch_.fillFrom(overflowList_);
-                            RuntimeLogDebug({ kTagBalancing }, "Acquired %zu elements from the overflow list", batch_.elementsCount());
+                        if (!this->overflowList().empty()) {
+                            this->batch_.fillFrom(this->overflowList());
+                            RuntimeLogDebug({ kTagBalancing }, "Acquired %zu elements from the overflow list", this->batch_.elementsCount());
                         } else {
                             bool newWorkAvailable = waitForMoreWork();
                             if (newWorkAvailable) continue;
                             return nullptr;
                         }
                     }
-                    RuntimeAssert(!batch_.empty(), "Must have acquired some elements");
+                    RuntimeAssert(!this->batch_.empty(), "Must have acquired some elements");
                     break;
                 }
             }
 
-            return batch_.tryPop();
+            return this->batch_.tryPop();
         }
 
     private:
         bool waitForMoreWork() noexcept {
-            RuntimeAssert(batch_.empty(), "Local batch must be depleted before waiting for shared work");
-            RuntimeAssert(overflowList_.empty(), "Local overflow list must be depleted before waiting for shared work");
+            RuntimeAssert(this->batch_.empty(), "Local batch must be depleted before waiting for shared work");
+            RuntimeAssert(this->overflowList().empty(), "Local overflow list must be depleted before waiting for shared work");
 
-            std::unique_lock lock(dispatcher_.waitMutex_);
+            std::unique_lock lock(this->dispatcher_.waitMutex_);
 
-            auto nowWaiting = dispatcher_.waitingWorkers_.fetch_add(1, std::memory_order_relaxed) + 1;
+            auto nowWaiting = this->dispatcher_.waitingWorkers_.fetch_add(1, std::memory_order_relaxed) + 1;
             RuntimeLogDebug({ kTagBalancing }, "Worker goes to sleep (now sleeping %zu of %zu)",
-                            nowWaiting, dispatcher_.registeredWorkers_.load(std::memory_order_relaxed));
+                            nowWaiting, this->dispatcher_.registeredWorkers_.load(std::memory_order_relaxed));
 
-            if (dispatcher_.allDone_) {
-                dispatcher_.waitingWorkers_.fetch_sub(1, std::memory_order_relaxed);
+            if (this->dispatcher_.allDone_) {
+                this->dispatcher_.waitingWorkers_.fetch_sub(1, std::memory_order_relaxed);
                 return false;
             }
 
-            if (nowWaiting == dispatcher_.registeredWorkers_.load(std::memory_order_relaxed)) {
+            if (nowWaiting == this->dispatcher_.registeredWorkers_.load(std::memory_order_relaxed)) {
                 // we are the last ones awake
                 RuntimeLogDebug({ kTagBalancing }, "Worker has detected termination");
-                dispatcher_.allDone_ = true;
+                this->dispatcher_.allDone_ = true;
+                this->dispatcher_.waitingWorkers_.fetch_sub(1, std::memory_order_relaxed);
                 lock.unlock();
-                dispatcher_.waitCV_.notify_all();
-                dispatcher_.waitingWorkers_.fetch_sub(1, std::memory_order_relaxed);
+                this->dispatcher_.waitCV_.notify_all();
                 return false;
             }
 
-            dispatcher_.waitCV_.wait(lock);
-            dispatcher_.waitingWorkers_.fetch_sub(1, std::memory_order_relaxed);
-            if (dispatcher_.allDone_) {
+            this->dispatcher_.waitCV_.wait(lock);
+            this->dispatcher_.waitingWorkers_.fetch_sub(1, std::memory_order_relaxed);
+            if (this->dispatcher_.allDone_) {
                 return false;
             }
             RuntimeLogDebug({ kTagBalancing }, "Worker woke up");
 
             return true;
         }
-
-        ParallelProcessor& dispatcher_;
-
-        Batch batch_;
-        ListImpl overflowList_;
     };
 
     ParallelProcessor() = default;

kotlin-native/runtime/src/main/cpp/ParallelProcessorTest.cpp

@@ -49,7 +49,7 @@ struct Task {
     Task* next_ = nullptr;
 };
 
-auto workBatch(std::size_t size) {
+auto createWork(std::size_t size) {
     std::list<Task> batch;
     for (size_t i = 0; i < size; ++i) {
         batch.emplace_back();
@@ -66,9 +66,11 @@ void offerWork(WorkList& wl, Iterable& batch) {
 }
 
 using ListImpl = intrusive_forward_list<Task>;
-static constexpr auto kBatchSize = 256;
-using Processor = ParallelProcessor<ListImpl, kBatchSize, 1024>;
+constexpr auto kBatchSize = 256;
+constexpr auto kBatchPoolSize = 4;
+using Processor = ParallelProcessor<ListImpl, kBatchSize, kBatchPoolSize>;
 using Worker = typename Processor::Worker;
+using WorkSource = typename Processor::WorkSource;
 
 } // namespace
 
@@ -99,16 +101,102 @@ TEST(ParallelProcessorTest, ContededRegistration) {
 
 TEST(ParallelProcessorTest, Sharing) {
     Processor processor;
+
     Worker giver(processor);
     Worker taker(processor);
+    EXPECT_THAT(processor.registeredWorkers(), 2);
 
-    auto work = workBatch(kBatchSize * 2);
-    offerWork(giver, work);
+    auto twoBatches = createWork(kBatchSize * 2);
+    offerWork(giver, twoBatches);
 
-    EXPECT_TRUE(taker.localEmpty());
+    EXPECT_TRUE(taker.retainsNoWork());
 
     // have to steal from giver
     EXPECT_NE(taker.tryPop(), nullptr);
 
-    EXPECT_FALSE(taker.localEmpty());
+    EXPECT_FALSE(taker.retainsNoWork());
+}
+
+TEST(ParallelProcessorTest, SharingFromNonWorkerSource) {
+    Processor processor;
+
+    WorkSource giver(processor);
+    EXPECT_THAT(processor.registeredWorkers(), 0);
+
+    Worker taker(processor);
+    EXPECT_THAT(processor.registeredWorkers(), 1);
+
+    auto work = createWork(kBatchSize * 2);
+    offerWork(giver, work);
+
+    EXPECT_TRUE(taker.retainsNoWork());
+
+    // have to steal from giver
+    EXPECT_NE(taker.tryPop(), nullptr);
+
+    EXPECT_FALSE(taker.retainsNoWork());
+}
+
+TEST(ParallelProcessorTest, Overflow) {
+    Processor processor;
+    Worker worker(processor);
+
+    auto workSize = kBatchSize * (kBatchPoolSize + 2);
+    auto work = createWork(workSize);
+    offerWork(worker, work);
+
+    std::size_t poppedCount = 0;
+    while (worker.tryPop() != nullptr) {
+        ++poppedCount;
+    }
+    EXPECT_THAT(poppedCount, workSize);
+    EXPECT_THAT(worker.retainsNoWork(), true);
+}
+
+TEST(ParallelProcessorTest, ForceFlush) {
+    Processor processor;
+
+    WorkSource source(processor);
+
+    auto workSize = kBatchSize / 2;
+    auto halfBatch = createWork(workSize);
+    offerWork(source, halfBatch);
+
+    EXPECT_THAT(source.forceFlush(), true);
+    EXPECT_THAT(source.retainsNoWork(), true);
+
+    Worker checker(processor);
+    std::size_t poppedCount = 0;
+    while (checker.tryPop() != nullptr) {
+        ++poppedCount;
+    }
+    EXPECT_THAT(poppedCount, workSize);
+}
+
+TEST(ParallelProcessorTest, ForceFlushWithOverflow) {
+    Processor processor;
+
+    // Fill up the processor's work pool
+    Worker overflower(processor);
+
+    auto poolSizeBatches = createWork(kBatchSize * kBatchPoolSize);
+    offerWork(overflower, poolSizeBatches);
+
+    EXPECT_THAT(overflower.forceFlush(), true);
+
+    // No overflow a local work source
+    WorkSource source(processor);
+
+    auto twoBatches = createWork(kBatchSize * 2);
+    offerWork(source, twoBatches);
+
+    // no space to flush into
+    EXPECT_THAT(source.forceFlush(), false);
+
+    // drain the processor's pool
+    while (overflower.tryPop() != nullptr) {}
+
+    // now can flush
+    EXPECT_THAT(source.forceFlush(), true);
+    EXPECT_THAT(source.retainsNoWork(), true);
 }