xref: /aosp_15_r20/external/pytorch/torch/csrc/distributed/c10d/sequence_num.hpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
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)18 inline 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)31 inline 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