add seed and offset

DeepLink-org · Apr 1, 2024 · 7a4c1ea · 7a4c1ea
1 parent 8d2a550
commit 7a4c1ea
Showing 1 changed file with 24 additions and 0 deletions.
diff --git a/dipu/torch_dipu/csrc_dipu/diopirt/diopirt_impl.cpp b/dipu/torch_dipu/csrc_dipu/diopirt/diopirt_impl.cpp
@@ -183,6 +183,30 @@ DIOPI_RT_API diopiError_t diopiGeneratorSetState(
   return diopiSuccess;
 }
 
+DIOPI_RT_API diopiError_t diopiGeneratorGetSeedAndOffset(
+    diopiGeneratorHandle_t th, uint64_t& seed, uint64_t& offset) {
+  auto generator = reinterpret_cast<at::Generator*>(th);
+  auto gen_impl = at::check_generator<dipu::DIPUGeneratorImpl>(*generator);
+  {
+    offset = gen_impl->get_offset();
+    seed = gen_impl->current_seed();
+  }
+
+  return diopiSuccess;
+}
+
+DIOPI_RT_API diopiError_t diopiGeneratorSetSeedAndOffset(
+    diopiGeneratorHandle_t th, uint64_t seed, uint64_t offset) {
+  auto generator = reinterpret_cast<at::Generator*>(th);
+  auto gen_impl = at::check_generator<dipu::DIPUGeneratorImpl>(*generator);
+  {
+    gen_impl->set_offset(offset);
+    gen_impl->set_current_seed(seed);
+  }
+
+  return diopiSuccess;
+}
+
 DIOPI_RT_API diopiError_t diopiRecordStart(const char* record_name,
                                            void** record) {
   *record = new RecordBlockCreator(record_name);