latch.h

Go to the documentation of this file.

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */
 
#pragma once
 
#include <dispenso/platform.h>
 
#include <dispenso/detail/completion_event_impl.h>
 
namespace dispenso {
 
class Latch {
 public:
  explicit Latch(uint32_t threadGroupCount) noexcept : impl_(threadGroupCount) {}
 
  void count_down(uint32_t n = 1) noexcept {
    if (impl_.intrusiveStatus().fetch_sub(n, std::memory_order_acq_rel) == 1) {
      impl_.notify(0);
    }
  }
 
  bool try_wait() const noexcept {
    return impl_.intrusiveStatus().load(std::memory_order_acquire) == 0;
  }
 
  void wait() const noexcept {
    impl_.wait(0);
  }
 
  void arrive_and_wait() noexcept {
    if (impl_.intrusiveStatus().fetch_sub(1, std::memory_order_acq_rel) > 1) {
      impl_.wait(0);
    } else {
      impl_.notify(0);
    }
  }
 
 private:
  detail::CompletionEventImpl impl_;
};
 
} // namespace dispenso