xref: /aosp_15_r20/external/pytorch/torch/csrc/distributed/c10d/GlooDeviceFactory.hpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #ifdef USE_C10D_GLOO
4 
5 #include <string>
6 
7 #include <c10/util/Registry.h>
8 #include <gloo/config.h>
9 #include <gloo/transport/device.h>
10 
11 namespace c10d {
12 
13 class TORCH_API GlooDeviceFactory {
14  public:
15   // Create new device instance for specific interface.
16   static std::shared_ptr<::gloo::transport::Device> makeDeviceForInterface(
17       const std::string& interface);
18 
19   // Create new device instance for specific hostname or address.
20   static std::shared_ptr<::gloo::transport::Device> makeDeviceForHostname(
21       const std::string& hostname);
22 };
23 
24 TORCH_DECLARE_SHARED_REGISTRY(
25     GlooDeviceRegistry,
26     ::gloo::transport::Device,
27     const std::string&, /* interface */
28     const std::string& /* hostname */);
29 
30 } // namespace c10d
31 
32 #endif // USE_C10D_GLOO
33