1 | #pragma once |
2 | |
3 | #include <vector> |
4 | #include <c10/macros/Macros.h> |
5 | #include <c10/util/Optional.h> |
6 | #include <c10/util/irange.h> |
7 | |
8 | namespace c10d { |
9 | const int kUnsetSeqNum = 0; |
10 | |
11 | namespace { |
12 | constexpr int kByteOffset = 8; |
13 | } |
14 | |
15 | // Converts from int to char vec to write in store |
16 | template <typename T> |
17 | inline std::vector<T> toVec(uint64_t num, int numBytes) { |
18 | std::vector<T> values; |
19 | // Read off bytes from right to left, pushing them into |
20 | // char array. |
21 | for (const auto i : c10::irange(numBytes)) { |
22 | uint8_t x = (num >> (kByteOffset * i)) & 0xff; |
23 | values.push_back(static_cast<T>(x)); |
24 | } |
25 | return values; |
26 | } |
27 | |
28 | // Converts from char vec (such as from store read) to int. |
29 | template <typename T> |
30 | inline uint64_t fromVec(const std::vector<T>& values) { |
31 | uint64_t num = 0; |
32 | // Set each byte at the correct location on num |
33 | for (const auto i : c10::irange(values.size())) { |
34 | uint8_t x = static_cast<uint8_t>(values[i]); |
35 | num |= (static_cast<int64_t>(x) << (kByteOffset * i)); |
36 | } |
37 | return num; |
38 | } |
39 | |
40 | class TORCH_API SequenceNum { |
41 | public: |
42 | SequenceNum(); |
43 | explicit SequenceNum(const uint64_t num); |
44 | // Retrieve num_. Will throw if not set. |
45 | uint64_t get() const; |
46 | // Increment num_. Will throw if not set. |
47 | void increment(); |
48 | // Increment num_ and return the old value. Will throw if not set. |
49 | uint64_t getAndIncrement(); |
50 | // Sets num_ |
51 | void set(const uint64_t num); |
52 | // Returns true if this SequenceNum is properly initialized with a value, else |
53 | // false. |
54 | bool isSet() const; |
55 | |
56 | SequenceNum& operator=(const SequenceNum& other); |
57 | |
58 | SequenceNum(const SequenceNum& other); |
59 | |
60 | private: |
61 | c10::optional<uint64_t> num_; |
62 | mutable std::mutex lock_; |
63 | }; |
64 | |
65 | } // namespace c10d |
66 | |