Integrate xpu into torch.Generator and torch.seed (pytorch#109866)

Integrate torch.xpu.Generator into torch.Generator
Integrate torch.xpu.seed into torch.seed
Pull Request resolved: pytorch#109866
Approved by: https://github.com/ezyang

Author: leizhenyuan

aten/src/ATen/Context.h

@@ -44,6 +44,8 @@ class TORCH_API Context {
       return at::detail::getCUDAHooks().getDefaultCUDAGenerator(device.index());
     } else if (device_type == at::kMPS) {
       return at::detail::getMPSHooks().getDefaultMPSGenerator();
+    } else if (device_type == at::kXPU) {
+      return at::detail::getXPUHooks().getDefaultXPUGenerator(device.index());
     } else if (device_type == at::kPrivateUse1) {
       return at::GetPrivateUse1HooksInterface()->getDefaultGenerator(
           device.index());
@@ -445,9 +447,9 @@ static inline void manual_seed(uint64_t seed) {
   }
   // NB: Sometimes we build with CUDA, but we don't have any GPUs
   // available. In that case, we must not seed CUDA; it will fail!
-  const auto num_gpus = detail::getCUDAHooks().getNumGPUs();
-  if (hasCUDA() && num_gpus > 0) {
-    for (const auto i : c10::irange(num_gpus)) {
+  const auto cuda_num_gpus = detail::getCUDAHooks().getNumGPUs();
+  if (hasCUDA() && cuda_num_gpus > 0) {
+    for (const auto i : c10::irange(cuda_num_gpus)) {
       auto cuda_gen = globalContext().defaultGenerator(
           Device(at::kCUDA, static_cast<c10::DeviceIndex>(i)));
       {
@@ -458,6 +460,19 @@ static inline void manual_seed(uint64_t seed) {
     }
   }
 
+  const auto xpu_num_gpus = detail::getXPUHooks().getNumGPUs();
+  if (hasXPU() && xpu_num_gpus > 0) {
+    for (const auto i : c10::irange(xpu_num_gpus)) {
+      auto xpu_gen = globalContext().defaultGenerator(
+          Device(at::kXPU, static_cast<c10::DeviceIndex>(i)));
+      {
+        // See Note [Acquire lock when using random generators]
+        std::lock_guard<std::mutex> lock(xpu_gen.mutex());
+        xpu_gen.set_current_seed(seed);
+      }
+    }
+  }
+
   if (hasMPS()) {
     auto mps_gen = globalContext().defaultGenerator(c10::DeviceType::MPS);
     // See Note [Acquire lock when using random generators]

aten/src/ATen/detail/XPUHooksInterface.h

@@ -2,7 +2,7 @@
 
 #include <c10/core/Device.h>
 #include <c10/util/Exception.h>
-
+#include <ATen/core/Generator.h>
 #include <c10/util/Registry.h>
 
 #include <cstddef>
@@ -55,7 +55,7 @@ struct TORCH_API XPUHooksInterface {
         false,
         "Cannot get XPU device without Intel Extension for Pytorch. ",
         XPU_HELP);
-  }
+  };
 
   virtual DLDevice_& getDLPackDeviceFromATenDevice(
       DLDevice_& dl_device,
@@ -65,6 +65,20 @@ struct TORCH_API XPUHooksInterface {
         false,
         "Cannot get XPU DL device without Intel Extension for Pytorch. ",
         XPU_HELP);
+  };
+
+  virtual Generator getXPUGenerator(DeviceIndex device_index = -1) const {
+    (void)device_index; // Suppress unused variable warning
+    TORCH_CHECK(false, "Cannot get XPU generator without Intel Extension for Pytorch. ", XPU_HELP);
+  }
+
+    const Generator& getDefaultXPUGenerator(DeviceIndex device_index = -1) const {
+    (void)device_index; // Suppress unused variable warning
+    TORCH_CHECK(false, "Cannot get default XPU generator without Intel Extension for Pytorch. ", XPU_HELP);
+  }
+
+  virtual int getNumGPUs() const {
+    return 0;
   }
 };
 

torch/csrc/Generator.cpp

@@ -5,6 +5,7 @@
 #include <structmember.h>
 
 #include <ATen/core/GeneratorForPrivateuseone.h>
+#include <ATen/detail/XPUHooksInterface.h>
 #include <torch/csrc/Device.h>
 #include <torch/csrc/Exceptions.h>
 #include <torch/csrc/THP.h>
@@ -71,7 +72,9 @@ static PyObject* THPGenerator_pynew(
     self->cdata = make_generator<MPSGeneratorImpl>();
   }
 #endif
-  else if (device.type() == at::kPrivateUse1) {
+  else if (device.type() == at::kXPU) {
+    self->cdata = at::detail::getXPUHooks().getXPUGenerator(device.index());
+  } else if (device.type() == at::kPrivateUse1) {
     self->cdata = at::GetGeneratorForPrivateuse1(device.index());
   } else {
     AT_ERROR(

torch/random.py

@@ -43,6 +43,9 @@ def manual_seed(seed) -> torch._C.Generator:
     if not torch.mps._is_in_bad_fork():
         torch.mps.manual_seed(seed)
 
+    if hasattr(torch, 'xpu') and not torch.xpu._is_in_bad_fork():
+        torch.xpu.manual_seed(seed)
+
     _seed_custom_device(seed)
 
     return default_generator.manual_seed(seed)
@@ -62,6 +65,9 @@ def seed() -> int:
     if not torch.mps._is_in_bad_fork():
         torch.mps.manual_seed(seed)
 
+    if hasattr(torch, 'xpu') and not torch.xpu._is_in_bad_fork():
+        torch.xpu.manual_seed(seed)
+
     _seed_custom_device(seed)
 
     return seed