Barrier.h 724 Bytes
Newer Older
Li Zhang's avatar
Li Zhang committed
1
2
3
4
// Copyright (c) OpenMMLab. All rights reserved.

#pragma once

lvhan028's avatar
lvhan028 committed
5
#include "src/turbomind/utils/logger.h"
Li Zhang's avatar
Li Zhang committed
6
7
#include <pthread.h>

lvhan028's avatar
lvhan028 committed
8
namespace turbomind {
Li Zhang's avatar
Li Zhang committed
9
10
11
12
13

class Barrier {
public:
    Barrier(unsigned count)
    {
lvhan028's avatar
lvhan028 committed
14
        TM_LOG_INFO("Barrier(%d)", (int)count);
Li Zhang's avatar
Li Zhang committed
15
16
17
        pthread_barrier_init(&barrier_, nullptr, count);
    }

AllentDan's avatar
AllentDan committed
18
19
20
    Barrier(const Barrier&) = delete;
    Barrier& operator=(const Barrier&) = delete;
    Barrier(Barrier&&) noexcept        = delete;
Li Zhang's avatar
Li Zhang committed
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
    Barrier& operator=(Barrier&&) noexcept = delete;

    void wait()
    {
        pthread_barrier_wait(&barrier_);
    }

    ~Barrier()
    {
        pthread_barrier_destroy(&barrier_);
    }

private:
    pthread_barrier_t barrier_{};
};

lvhan028's avatar
lvhan028 committed
37
}  // namespace turbomind