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