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
8namespace c10d {
9const int kUnsetSeqNum = 0;
10
11namespace {
12 constexpr int kByteOffset = 8;
13}
14
15// Converts from int to char vec to write in store
16template <typename T>
17inline 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.
29template <typename T>
30inline 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
40class 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