1 #pragma once
2
3 #include <cstdint>
4 #include <deque>
5 #include <mutex>
6 #include <utility>
7
8 #include <c10/util/Exception.h>
9 #include <c10/util/intrusive_ptr.h>
10 #include <c10/core/Device.h>
11 #include <c10/core/DispatchKeySet.h>
12
13 // For the record I don't think this is a correct pimpl idiom.
14 // Including Impl header in interface header defeats the purpose
15 // because you can't change Impl private members without forcing
16 // everything that included the interface to rebuild.
17 // Impl should be forward-declared in the interface header instead.
18 #include <c10/core/GeneratorImpl.h>
19
20 /**
21 * Note [Generator]
22 * ~~~~~~~~~~~~~~~~
23 * A Pseudo Random Number Generator (PRNG) is an engine that uses an algorithm to
24 * generate a seemingly random sequence of numbers, that may be later be used in creating
25 * a random distribution. Such an engine almost always maintains a state and requires a
26 * seed to start off the creation of random numbers. Often times, users have
27 * found it beneficial to be able to explicitly create, retain, and destroy
28 * PRNG states and also be able to have control over the seed value.
29 *
30 * A Generator in ATen gives users the ability to read, write and modify a PRNG engine.
31 * For instance, it does so by letting users seed a PRNG engine, fork the state of the
32 * engine, etc.
33 *
34 * By default, there is one generator per device, and a device's generator is
35 * lazily created. A user can use the torch.Generator() api to create their own generator.
36 */
37
38 /**
39 * Note [Acquire lock when using random generators]
40 * ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
41 * Generator and its derived classes are NOT thread-safe. Please note that most of the
42 * places where we have inserted locking for generators are historically based, and we
43 * haven't actually checked that everything is truly thread safe (and it probably isn't).
44 * Please use the public mutex_ when using any methods from these classes, except for the
45 * read-only methods. You can learn about the usage by looking into the unittests
46 * (aten/src/ATen/cpu_generator_test.cpp) and other places where we have used lock_guard.
47 *
48 * TODO: Look into changing the threading semantics of Generators in ATen (e.g., making
49 * them non-thread safe and instead making the generator state splittable, to accommodate
50 * forks into other threads).
51 */
52
53 namespace at {
54
55 class Tensor;
56
57 struct TORCH_API Generator {
58 Generator() = default;
59
GeneratorGenerator60 explicit Generator(c10::intrusive_ptr<c10::GeneratorImpl> gen_impl)
61 : impl_(std::move(gen_impl)) {
62 if (impl_.get() == nullptr) {
63 throw std::runtime_error("GeneratorImpl with nullptr is not supported");
64 }
65 }
66
67 bool operator==(const Generator& rhs) const {
68 return this->impl_ == rhs.impl_;
69 }
70
71 bool operator!=(const Generator& rhs) const {
72 return !((*this) == rhs);
73 }
74
definedGenerator75 bool defined() const {
76 return static_cast<bool>(impl_);
77 }
78
unsafeGetGeneratorImplGenerator79 c10::GeneratorImpl* unsafeGetGeneratorImpl() const {
80 return impl_.get();
81 }
82
unsafeReleaseGeneratorImplGenerator83 c10::GeneratorImpl* unsafeReleaseGeneratorImpl() {
84 return impl_.release();
85 }
86
getIntrusivePtrGenerator87 const c10::intrusive_ptr<c10::GeneratorImpl>& getIntrusivePtr() const {
88 return impl_;
89 }
90
set_current_seedGenerator91 void set_current_seed(uint64_t seed) { impl_->set_current_seed(seed); }
92 // Sets the offset of Generator state to the desired offset. This is currently
93 // supported for only Philox based Generators, i.e., CUDA and MPS.
set_offsetGenerator94 void set_offset(uint64_t offset) { impl_->set_offset(offset); }
95
96 // Returns the offset of Generator state. This is currently supported for only
97 // Philox based Generators, i.e., CUDA and MPS.
get_offsetGenerator98 uint64_t get_offset() const { return impl_->get_offset(); }
99
current_seedGenerator100 uint64_t current_seed() const { return impl_->current_seed(); }
101
seedGenerator102 uint64_t seed() { return impl_->seed(); }
103
104 // Implementation not inlined to prevent cycle reference between
105 // `ATen/core/Generator.h` and `ATen/core/Tensor.h`
106 void set_state(const at::Tensor& new_state);
107
108 at::Tensor get_state() const;
109
110 void graphsafe_set_state(const Generator& new_state);
111
112 Generator graphsafe_get_state() const;
113
mutexGenerator114 std::mutex& mutex() {
115 return impl_->mutex_;
116 }
117
key_setGenerator118 DispatchKeySet key_set() const {
119 return impl_->key_set();
120 }
121
deviceGenerator122 Device device() const { return impl_->device(); }
123
set_pyobjGenerator124 inline void set_pyobj(PyObject* pyobj) const noexcept {
125 impl_->set_pyobj(pyobj);
126 }
127
pyobjGenerator128 inline PyObject* pyobj() const noexcept {
129 return impl_->pyobj();
130 }
131
132 template<typename T>
getGenerator133 T* get() const { return static_cast<T*>(impl_.get()); }
134
cloneGenerator135 Generator clone() const {
136 return Generator(impl_->clone());
137 }
138
139 private:
140 c10::intrusive_ptr<c10::GeneratorImpl> impl_;
141 };
142
143 template<class Impl, class... Args>
make_generator(Args &&...args)144 Generator make_generator(Args&&... args) {
145 return Generator(c10::make_intrusive<Impl>(std::forward<Args>(args)...));
146 }
147
148 /**
149 * Utility function to static cast input Generator* to
150 * the backend generator type (CPU/CUDAGeneratorImpl etc.)
151 */
152 template <typename T>
check_generator(std::optional<Generator> gen)153 inline T * check_generator(std::optional<Generator> gen) {
154 TORCH_CHECK(gen.has_value(), "Expected Generator but received nullopt");
155 TORCH_CHECK(gen->defined(), "Generator with undefined implementation is not allowed");
156 TORCH_CHECK(T::device_type() == gen->device().type(), "Expected a '", T::device_type(), "' device type for generator but found '", gen->device().type(), "'");
157 return gen->get<T>();
158 }
159
160 /**
161 * Utility function used in tensor implementations, which
162 * supplies the default generator to tensors, if an input generator
163 * is not supplied. The input Generator* is also static casted to
164 * the backend generator type (CPU/CUDAGeneratorImpl etc.)
165 */
166 template <typename T>
get_generator_or_default(const std::optional<Generator> & gen,const Generator & default_gen)167 inline T* get_generator_or_default(const std::optional<Generator>& gen, const Generator& default_gen) {
168 return gen.has_value() && gen->defined() ? check_generator<T>(gen) : check_generator<T>(default_gen);
169 }
170
171 namespace detail {
172
173 /**
174 * Helper function for checking the validity of new random generator
175 * state. Right now following conditions are checked:
176 *
177 * - The new state tensor must be a torch.ByteTensor
178 * - Data of the new state tensor must be contiguous
179 */
check_rng_state(const c10::TensorImpl & new_state)180 inline void check_rng_state(const c10::TensorImpl& new_state) {
181 TORCH_CHECK_TYPE(
182 new_state.layout() == kStrided && new_state.device().type() == kCPU && new_state.dtype() == kByte,
183 "RNG state must be a torch.ByteTensor"
184 );
185
186 TORCH_CHECK(new_state.is_contiguous(), "RNG state must be contiguous");
187 }
188
189 } // namespace detail
190
191 } // namespace at
192