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