Class ClusterBarrier

include/barrier.h:26–164

}

struct ClusterBarrier {
  using ValueType = uint64_t;

 protected:
  // Can never be initialized - can only be aliased to smem
  ValueType barrier_;

 public:
  DEVICE ClusterBarrier() = delete;

  DEVICE void init(uint32_t arrive_count) const {
    ClusterBarrier::init(&this->barrier_, arrive_count);
  }

  DEVICE uint32_t test_wait(uint32_t phase, uint32_t pred = true) const {
    return ClusterBarrier::test_wait(&this->barrier_, phase, pred);
  }

  DEVICE uint32_t try_wait(uint32_t phase) const {
    return ClusterBarrier::try_wait(&this->barrier_, phase);
  }

  DEVICE void wait(uint32_t phase) const {
    ClusterBarrier::wait(&this->barrier_, phase);
  }

  // Barrier arrive on local smem
  DEVICE void arrive() const { ClusterBarrier::arrive(&this->barrier_); }

  // Remote SMEM arrive with a perdicate (usually done to pick the thread doing
  // the arrive)
  DEVICE void arrive(uint32_t cta_id, uint32_t pred = true) const {
    ClusterBarrier::arrive(&this->barrier_, cta_id, pred);
  }

  //
  //  Static Versions
  //
  DEVICE static void init(ValueType const *smem_ptr, uint32_t arrive_count) {
    uint32_t smem_addr = cast_smem_ptr_to_uint(smem_ptr);
    asm volatile(
        "{\n\t"
        "mbarrier.init.shared.b64 [%1], %0; \n"
        "}"
        :
        : "r"(arrive_count), "r"(smem_addr));
  }

  // Static version of wait - in case we don't want to burn a register
  DEVICE static void wait(ValueType const *smem_ptr, uint32_t phase) {
    uint32_t smem_addr = cast_smem_ptr_to_uint(smem_ptr);
    // Arbitrarily large timer value after which try-wait expires and re-tries.
    uint32_t ticks = 0x989680;
    asm volatile(
        "{\n\t"
        ".reg .pred       P1; \n\t"
        "LAB_WAIT: \n\t"
        "mbarrier.try_wait.parity.shared.b64 P1, [%0], %1, %2; \n\t"