Barrier.h 771 Bytes
Newer Older
Li Zhang's avatar
Li Zhang committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
// Copyright (c) OpenMMLab. All rights reserved.

#pragma once

#include "src/fastertransformer/utils/logger.h"
#include <pthread.h>

namespace fastertransformer {

class Barrier {
public:
    Barrier(unsigned count)
    {
        FT_LOG_INFO("Barrier(%d)", (int)count);
        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_{};
};

AllentDan's avatar
AllentDan committed
37
}  // namespace fastertransformer