add seed and offset

Author: Bonbon-Tang

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);