Управление устройствами

Общие сведения

Управление устройствами включает основные операции, такие как получение сведений о количестве доступных устройств и переключение между ними. Бэкенды ускорителей оборачивают API среды выполнения устройств и предоставляют их PyTorch.

Проектирование

Поставщикам ускорителей следует реализовать следующие основные функции:

Эти функции являются строительными блоками для потоков, событий и управления памятью. Проверяйте входные данные и корректно обрабатывайте ошибки.

Реализация

В этом разделе на примере set_device показано управление устройствами. Для реализации необходимы:

1. Обёртки C++ для среды выполнения устройства
2. Привязки Python для предоставления функций C++
3. Удобные для пользователя API Python

В качестве примера рассмотрим OpenReg (Open Registration) — пример интеграции с PyTorch, который устраняет пробел в интеграции внешних бэкендов ускорителей. Его реализация (OpenRegFunctions.h/cpp) показывает, как аккуратно обернуть стороннюю среду выполнения. Эти функции повторно используются во всём бэкенде — для потоков, событий, генераторов и привязок Python.

1orError_t SetDevice(DeviceIndex device) {
2  int cur_device = -1;
3  OPENREG_CHECK(orGetDevice(&cur_device));
4  if (device == cur_device) {
5    return orSuccess;
6  }
7  return orSetDevice(device);
8}

1OPENREG_EXPORT void set_device(DeviceIndex device) {
2  check_device_index(device);
3  OPENREG_CHECK(SetDevice(device));
4}

 1PyObject* _setDevice(PyObject* self, PyObject* arg) {
 2  HANDLE_TH_ERRORS
 3  TORCH_CHECK(THPUtils_checkLong(arg), "invalid argument to setDevice");
 4  auto device = THPUtils_unpackDeviceIndex(arg);
 5  torch::utils::device_lazy_init(at::kPrivateUse1);
 6  c10::openreg::set_device(device);
 7
 8  Py_RETURN_NONE;
 9  END_HANDLE_TH_ERRORS
10}

1static PyMethodDef methods[] = {
2    {"_init", _initExtension, METH_NOARGS, nullptr},
3    {"_isInBadFork", _isInBadFork, METH_NOARGS, nullptr},
4    {"_get_default_generator", _getDefaultGenerator, METH_O, nullptr},
5    {"_get_device", _getDevice, METH_NOARGS, nullptr},
6    {"_set_device", _setDevice, METH_O, nullptr},
7    {"_exchangeDevice", _exchangeDevice, METH_O, nullptr},
8    {"_get_device_count", _getDeviceCount, METH_NOARGS, nullptr},
9    {nullptr, nullptr, 0, nullptr}};

1def set_device(device) -> None:
2    if device >= 0:
3        torch_openreg._C._set_device(device)
4
5

Защитный объект устройства

Защитные объекты устройств автоматически переключают устройства с безопасной обработкой исключений. Они похожи на блокировки C++: переключают устройства при создании и восстанавливают исходное состояние при уничтожении.

Реализуйте DeviceGuardImplInterface для интеграции с системой защитных объектов PyTorch:

 1  /**
 2   * Return the type of device managed by this guard implementation.
 3   */
 4  DeviceType type() const override {
 5    return static_type;
 6  }
 7  /**
 8   * Set the current device to device d, and return the previous Device.
 9   */
10  // LITERALINCLUDE START: OPENREG GUARD DEVICE MANAGEMENT
11  Device exchangeDevice(Device d) const override {
12    TORCH_CHECK(d.is_privateuseone(), "Expected a PrivateUse1 device, but got ", d);
13
14    auto old_device_index = ExchangeDevice(d.index());
15    return Device(static_type, old_device_index);
16  }
17  // LITERALINCLUDE END: OPENREG GUARD DEVICE MANAGEMENT
18
19  /**
20   * Get the current device.
21   */
22  Device getDevice() const override {
23    int device_index = current_device();
24    return c10::Device(static_type, device_index);
25  }
26
27  /**
28   * Get the device capability for a given device.
29   * By default, OpenReg has 2 same devices with the same capability.
30   */
31  DeviceCapability getDeviceCapability(Device /*unused*/) const override {
32    return DeviceCapability();
33  }
34
35  /**
36   * Set the current device to c10::Device.
37   */
38  void setDevice(Device d) const override {
39    TORCH_CHECK(d.is_privateuseone(), "Expected a PrivateUse1 device, but got ", d);
40
41    set_device(d.index());
42  }
43
44  /**
45   * Set the current device to device d, without checking for errors
46   * (so, e.g., this can be called from a destructor).
47   */
48  void uncheckedSetDevice(Device d) const noexcept override {
49    set_device(d.index());
50  }
51
52  /**
53   * Get the number of devices.
54   *
55   * WARNING: This is REQUIRED to not raise an exception.
56   * If there is some sort of problem, e.g., driver error,
57   * you should report that there are zero available devices.
58   */
59  DeviceIndex deviceCount() const noexcept override {
60    return device_count();
61  }
62
63  /**
64   * Wait (by blocking the calling thread) until all the work has
65   * completed running on the device.
66   */
67  void synchronizeDevice(const DeviceIndex device_index) const override {
68    OPENREG_CHECK(orDeviceSynchronize());
69  }

Это позволит использовать защитный объект в PyTorch для типа устройства PrivateUse1; после этого пользователи смогут применять стандартные защитные объекты устройств PyTorch с пользовательским бэкендом.